SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In detail, since the objective function in Eq. (4) is a compositional function whose stochastic gradient regarding  $\mathbf{x}$ , i.e.,  $\nabla g^{(k)}(\mathbf{x}_t^{(k)};\xi_t^{(k)})^T\nabla_gf^{(k)}(g^{(k)}(\mathbf{x}_t^{(k)};\xi_t^{(k)}),\mathbf{y}_t^{(k)};\zeta_t^{(k)})$ , is a biased estimation for the full gradient, we employ the stochastic compositional gradient  $\nabla g^{(k)}(\mathbf{x}_t^{(k)};\xi_t^{(k)})^T\nabla_gf^{(k)}(\mathbf{h}_t^{(k)},\mathbf{y}_t^{(k)};\zeta_t^{(k)})$  to update the model parameter  $\mathbf{x}$ , where  $\mathbf{h}_t^{(k)}$  is the moving-average estimation of the inner-level function  $g^{(k)}(\mathbf{x}_t^{(k)};\xi_t^{(k)})$  on the  $k$ -th device, which is defined as follows:
$$
\mathbf {h} _ {t} ^ {(k)} = (1 - \alpha \eta) \mathbf {h} _ {t - 1} ^ {(k)} + \alpha \eta g ^ {(k)} \left(\mathbf {x} _ {t} ^ {(k)}; \xi_ {t} ^ {(k)}\right), \tag {11}
$$
where  $\alpha > 0$  and  $\eta > 0$  are two hyperparameters, and  $\alpha \eta \in (0,1)$ . The objective function in Eq. (4) is not compositional regarding  $\mathbf{y}$ , thus we can directly leverage its stochastic gradient to perform an update. Then, based on the obtained stochastic (compositional) gradient, we compute the momentum as follows:
$$
\mathbf {u} _ {t} ^ {(k)} = \left(1 - \beta_ {x} \eta\right) \mathbf {u} _ {t - 1} ^ {(k)} + \beta_ {x} \eta \nabla g ^ {(k)} \left(\mathbf {x} _ {t} ^ {(k)}; \xi_ {t} ^ {(k)}\right) ^ {T} \nabla_ {g} f ^ {(k)} \left(\mathbf {h} _ {t} ^ {(k)}, \mathbf {y} _ {t} ^ {(k)}; \zeta_ {t} ^ {(k)}\right), \tag {12}
$$
$$
\mathbf {v} _ {t} ^ {(k)} = (1 - \beta_ {y} \eta) \mathbf {v} _ {t - 1} ^ {(k)} + \beta_ {y} \eta \nabla_ {y} f ^ {(k)} (\mathbf {h} _ {t} ^ {(k)}, \mathbf {y} _ {t} ^ {(k)}; \zeta_ {t} ^ {(k)}),
$$
where  $\beta_{x} > 0$  and  $\beta_{y} > 0$  are two hyperparameters,  $\beta_{x}\eta \in (0,1)$ , and  $\beta_{y}\eta \in (0,1)$ . Based on the obtained momentum, each device updates its local model parameters as follows:
$$
\mathbf {x} _ {t + 1} ^ {(k)} = \mathbf {x} _ {t} ^ {(k)} - \gamma_ {x} \eta \mathbf {u} _ {t} ^ {(k)}, \mathbf {y} _ {t + 1} ^ {(k)} = \mathbf {y} _ {t} ^ {(k)} + \gamma_ {y} \eta \mathbf {v} _ {t} ^ {(k)}, \tag {13}
$$
where  $\gamma_{x} > 0$  and  $\gamma_y > 0$ .
As we mentioned before, to obtain the global inner-level function, our algorithm periodically communicates the moving-average estimation of the inner-level function, i.e.,  $\mathbf{h}_{t + 1}^{(k)}$ . In particular, at every  $p$  iterations, i.e.,  $\mathrm{mod}(t + 1,p) == 0$  where  $p > 1$  is the communication period, each device uploads  $\mathbf{h}_{t + 1}^{(k)}$  to the central server and the central server computes the average of all received variables, which will be further broadcast to all devices as follows:
$$
\mathbf {h} _ {t + 1} ^ {(k)} = \bar {\mathbf {h}} _ {t + 1} \triangleq \frac {1}{K} \sum_ {k ^ {\prime} = 1} ^ {K} \mathbf {h} _ {t + 1} ^ {\left(k ^ {\prime}\right)}. \tag {14}
$$
In this way, each device is able to obtain the estimate of the global inner-level function periodically. As for the model parameters and momentum, we employ the same strategy as traditional federated learning methods [23, 36] to communicate them periodically with the central server, which is shown in Step 6 in Algorithm 1.
In summary, we developed a novel local stochastic compositional gradient descent ascent with momentum algorithm for the compositional minimax problem, which shows how to deal with two distributed functions in federated learning. With our algorithm, we can enable federated learning for the compositional deep AUC maximization model, benefiting imbalanced data classification tasks.
In this section, we provide the convergence rate of our algorithm to show how it is affected by the number of devices and communication period.
To investigate the convergence rate of our algorithm, we introduce the following auxiliary functions:
$$
\mathbf {y} _ {*} (\mathbf {x}) = \arg \max  _ {\mathbf {y} \in \mathbb {R} ^ {d _ {2}}} \frac {1}{K} \sum_ {k = 1} ^ {K} f ^ {(k)} \left(\frac {1}{K} \sum_ {k ^ {\prime} = 1} ^ {K} g ^ {\left(k ^ {\prime}\right)} (\mathbf {x}), \mathbf {y}\right), \quad \Phi (\mathbf {x}) = \frac {1}{K} \sum_ {k = 1} ^ {K} f ^ {(k)} \left(\frac {1}{K} \sum_ {k ^ {\prime} = 1} ^ {K} g ^ {\left(k ^ {\prime}\right)} (\mathbf {x}), \mathbf {y} _ {*} (\mathbf {x})\right). \tag {15}
$$
Then, based on Assumptions 1-4, we can obtain that  $\Phi^{(k)}$  is  $L_{\Phi}$ -smooth, where  $L_{\Phi} = \frac{2C_g^2L_f^2}{\mu} + C_f L_g$ . The proof can be found in Lemma 2 of Appendix A. In terms of these auxiliary functions, we establish the convergence rate of our algorithm.
Theorem 1. Given Assumption 1-4, by setting  $\alpha > 0$ ,  $\beta_{x} > 0$ ,  $\beta_{y} > 0$
$$
\eta \leq \min \{\frac {1}{2 \gamma_ {x} L _ {\phi}}, \frac {1}{1 0 p ^ {2} \gamma_ {y} L _ {f}}, \frac {1}{\alpha}, \frac {1}{\beta_ {x}}, \frac {1}{\beta_ {y}}, 1 \}, \gamma_ {y} \leq \min \left\{\frac {1}{6 L _ {f}}, \frac {3 \mu \beta_ {y} ^ {2}}{4 0 0 L _ {f} ^ {2}}, \frac {3 \beta_ {x} ^ {2}}{1 6 \mu} \right\}, a n d \gamma_ {x} \leq
$$
$$
\min \left\{\frac {\alpha \mu}{1 0 0 C _ {g} ^ {2} L _ {f} \sqrt {1 + 6 L _ {f} ^ {2}}}, \frac {\beta_ {x}}{3 2 \sqrt {C _ {g} ^ {4} L _ {f} ^ {2} + C _ {f} ^ {2} L _ {g} ^ {2}}}, \frac {\beta_ {y} \mu}{1 4 4 C _ {g} ^ {2} L _ {f} ^ {2}}, \frac {\sqrt {\alpha} \mu}{2 4 C _ {g} \sqrt {1 0 0 C _ {g} ^ {2} L _ {f} ^ {4} + 2 C _ {g} ^ {2} \mu^ {2}}}, \frac {\gamma_ {y} \mu^ {2}}{2 0 C _ {g} ^ {2} L _ {f} ^ {2}} \right\},
$$
Algorithm 1 has the following convergence rate
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} [ \| \nabla \Phi (\bar {\mathbf {x}} _ {t}) \| ^ {2} ] \leq \frac {2 (\Phi (\mathbf {x} _ {0}) - \Phi (\mathbf {x} _ {*})}{\gamma_ {x} \eta T} + \frac {2 4 C _ {g} ^ {2} L _ {f} ^ {2}}{\gamma_ {y} \eta \mu T} \| \mathbf {y} _ {0} - \mathbf {y} ^ {*} (\mathbf {x} _ {0}) \| ^ {2} + O (\frac {\eta}{K}) + O (\frac {1}{\eta T}) \\ + O \left(p ^ {2} \eta^ {2}\right) + O \left(p ^ {4} \eta^ {4}\right) + O \left(p ^ {6} \eta^ {6}\right) + O \left(p ^ {8} \eta^ {8}\right) + O \left(p ^ {1 0} \eta^ {1 0}\right). \tag {16} \\ \end{array}
$$
Remark 1. In terms of Theorem 1, for sufficiently large  $T$ , by setting the learning rate  $\eta = O(K^{1/2}/T^{1/2})$ ,  $p = O(T^{1/4}/K^{3/4})$ , Algorithm 1 can achieve  $O(1/\sqrt{KT})$  convergence rate, which indicates a linear speedup with respect to the number of devices  $K$ . In addition, it is straightforward to show that the communication complexity of our algorithm is  $T/p = O(K^{3/4}T^{3/4})$ . Moreover, to achieve the  $\epsilon$ -accuracy solution, i.e.,  $\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}[\|\nabla\Phi(\bar{\mathbf{x}}_t)\|^2] \leq \epsilon^2$ , by setting  $\eta = O(K\epsilon^2)$  and  $p = O(1/(K\epsilon))$ , then the sample complexity on each device is  $O(1/(K\epsilon^4))$  and the communication complexity is  $O(1/\epsilon^3)$ .
Challenges. The compositional structure in the loss function, especially when the inner-level functions are distributed on different devices, makes the convergence analysis challenging. In fact, the existing federated compositional minimization algorithms [4, 27, 12, 28] only consider a much simpler case where the inner-level functions are not distributed across devices. Therefore, our setting is much more challenging than existing works. On the other hand, all existing federated compositional minimization algorithms [4, 27, 12, 28] fail to achieve linear speedup with respect to the number of devices. Thus, it is still unclear whether the linear speedup is achievable for federated compositional optimization algorithm. In this paper, we successfully addressed these challenges with novel theoretical analysis strategies and achieved the linear speedup for the first time for federated compositional minimax optimization algorithms. We believe our approaches, e.g., that for bounding consensus errors, can be applied to the minimization algorithms to achieve linear speedup.
In this section, we present the experimental results to demonstrate the performance of our algorithm.
(a) CATvsDOG
(d) STL10
(e) FashionMNIST
(f) Melanoma
Figure 1: Testing performance with AUC score versus the number of iterations when the communication period  $p = 4$ .
(f) Melanoma
Figure 2: Testing performance with AUC score versus the number of iterations when the communication period  $p = 8$ .
(a) CATvsDOG
(d) STL10
Figure 3: Testing performance with AUC score versus the number of iterations when the communication period  $p = 16$ .
Datasets. In our experiments, we employ six image classification datasets, including CIFAR10 [15], CIFAR100 [15], STL10 [1], FashionMNIST [32], CATvsDOG $^2$ , and Melanoma [22]. For the first four datasets, following [39], we consider the first half of classes to be the positive class, and the second half as the negative class. Then, in order to construct highly imbalanced data, we randomly drop some samples of the positive class in the training set. Specifically, the ratio between positive samples and all samples is set to 0.1. For the two-class dataset, CATvsDOG, we employ the same
$^{2}$ https://www.kaggle.com/c/dogs-vs-cats
Table 1: The comparison between the test AUC score of different methods on all datasets. Here,  $p$  denotes the communication period.
Datasets	Methods	AUC
p=4	p=8	p=16
CATvsDOG	LocalSCGDAM	0.933±0.000	0.936±0.000	0.928±0.000
CoDA	0.895±0.000	0.892±0.000	0.883±0.001
LocalSGDAM	0.899±0.000	0.884±0.000	0.884±0.001
LocalSGDM	0.888±0.001	0.889±0.000	0.887±0.000
LocalSGDM_RL	0.909±0.000	0.901±0.001	0.917±0.001
CIFAR10	LocalSCGDAM	0.914±0.000	0.914±0.000	0.916±0.000
CoDA	0.890±0.000	0.886±0.000	0.883±0.000
LocalSGDAM	0.893±0.000	0.880±0.000	0.880±0.000
LocalSGDM	0.883±0.001	0.871±0.000	0.874±0.000
LocalSGDM_RL	0.890±0.000	0.904±0.001	0.883±0.001
CIFAR100	LocalSCGDAM	0.702±0.000	0.704±0.001	0.703±0.001
CoDA	0.694±0.001	0.681±0.001	0.685±0.000
LocalSGDAM	0.692±0.000	0.694±0.000	0.689±0.001
LocalSGDM	0.675±0.001	0.669±0.000	0.669±0.000
LocalSGDM_RL	0.676±0.000	0.690±0.001	0.682±0.001
STL10	LocalSCGDAM	0.820±0.001	0.817±0.000	0.801±0.000
CoDA	0.801±0.000	0.784±0.000	0.783±0.000
LocalSGDAM	0.792±0.000	0.790±0.000	0.780±0.000
LocalSGDM	0.760±0.001	0.808±0.000	0.757±0.001
LocalSGDM_RL	0.773±0.001	0.771±0.003	0.752±0.003
FashionMNIST	LocalSCGDAM	0.980±0.000	0.980±0.000	0.980±0.000
CoDA	0.976±0.000	0.976±0.000	0.976±0.000
LocalSGDAM	0.977±0.000	0.977±0.000	0.976±0.000
LocalSGDM	0.963±0.000	0.956±0.000	0.955±0.000
LocalSGDM_RL	0.980±0.000	0.957±0.000	0.962±0.000
Melanoma	LocalSCGDAM	0.876±0.000	0.880±0.000	0.870±0.000
CoDA	0.734±0.002	0.721±0.000	0.725±0.003
LocalSGDAM	0.730±0.000	0.729±0.000	0.721±0.003
LocalSGDM	0.774±0.001	0.766±0.001	0.750±0.000
LocalSGDM_RL	0.737±0.001	0.721±0.000	0.773±0.001
strategy to construct the imbalanced training data. For these synthetic imbalanced datasets, the testing set is balanced. Melanoma is an intrinsically imbalanced medical image classification dataset, which we do not modify. The details about these benchmark datasets are summarized in Table 4.
Experimental Settings. For Melanoma, we use DenseNet121 [13] where the dimensionality of the last layer is set to 1 for binary classification. The details for the classifier for FashionMNIST can be found in Appendix B. For the other datasets, we use ResNet20 [10], where the last layer is also set to 1. To demonstrate the performance of our algorithm, we compare it with three state-of-the-art methods: LocalSGDM [36], LocalSGDM_RL [30], CoDA [8], LocalSGDAM [23]. Specifically, LocalSGDM uses momentum SGD to optimize the standard cross-entropy loss function. LocalSGDAM employs momentum SGD to optimize a Ratio Loss function, which is to add a regularization term to the standard cross-entropy loss function to address the imbalance distribution issue. CoDA leverages SGDA to optimize AUC loss, while LocalSGDAM exploits momentum SGDA to optimize AUC loss. For a fair comparison, we use similar learning rates for all algorithms. The details can be found in Appendix B. We use 4 devices (i.e., GPUs) in our experiment. The batch size on each device is set to 8 for STL10, 16 for Melanoma, and 32 for the others.
In Table 1, we report the AUC score of the test set for all methods, where we show the average and variance computed across all devices. Here, the communication period is set to 4, 8, and 16, respectively. It can be observed that our LocalSCGDAM algorithm outperforms all competing methods
for all cases. For instance, our LocalSCGDAM can beat baseline methods on CATvsDOG dataset with a large margin for all communication periods. These observations confirm the effectiveness of our algorithm. In addition, we plot the average AUC score of the test set versus the number of iterations in Figures 1, 2, 3. It can also be observed that our algorithm outperforms baseline methods consistently, which further confirms the efficacy of our algorithm.
To further demonstrate the performance of our algorithm, we apply these algorithms to the dataset with different imbalance ratios. Using the CATvsDOG dataset, we set the imbalance ratio to 0.01, 0.05, and 0.2 to construct three imbalanced training sets. The averaged testing AUC score of these three datasets versus the number of iterations is shown in Figure 4. It can be observed that our algorithm outperforms competing methods consistently and is robust to large imbalances in the training data. Especially when the training set is highly imbalanced, e.g., the imbalance ratio is 0.01, all AUC based methods outperform the cross-entropy loss based method significantly, and our LocalSCGDAM beats other AUC based methods with a large margin.
(a) Imbalance Ratio:0.01
(b) Imbalance Ratio:0.05
(c) Imbalance Ratio:0.2
Finally, we compare our algorithm with two additional baseline methods: SCAFFOLD [14] and FedProx [16]. These two methods assume the local data distribution is imbalanced but the global one is balanced. Then, they use the global gradient to correct the local one. However, this kind of methods do not work when the global data distribution is imbalanced. Specifically, when the global gradient itself is computed on the imbalanced data, rather than the balanced one, it cannot alleviate the imbalance issue in the local gradient. In Figure 5, we show the test AUC score of the STL10 dataset, where we use the same experimental setting as that of Figure 1, i.e., both the local and global data distributions are imbalanced. It can be observed that
Figure 4: The test AUC score versus the number of iterations when using different imbalance ratios for CATvsDOG.
Figure 5: The test AUROC score for STL10.
our algorithm outperforms those two baselines with a large margin, which confirms the effectiveness of our algorithm in handling the global imbalanced data distribution.
In this paper, we developed a novel local stochastic compositional gradient descent ascent algorithm to solve the federated compositional deep AUC maximization problem. On the theoretical side, we established the convergence rate of our algorithm, which enjoys a linear speedup with respect to the number of devices. On the empirical side, extensive experimental results on multiple imbalanced image classification tasks confirm the effectiveness of our algorithm.
We thank anonymous reviewers for constructive comments. T. Yang was partially supported by NSF Career Award 2246753, NSF Grant 2246757.
[38] H. Yuan, X. Lian, and J. Liu. Stochastic recursive variance reduction for efficient smooth non-convex compositional optimization. arXiv preprint arXiv:1912.13515, 2019.
[39] Z. Yuan, Z. Guo, N. Chawla, and T. Yang. Compositional training for end-to-end deep auc maximization. In International Conference on Learning Representations, 2021.
[40] Z. Yuan, Z. Guo, Y. Xu, Y. Ying, and T. Yang. Federated deep auc maximization for heterogeneous data with a constant communication complexity. In International Conference on Machine Learning, pages 12219-12229. PMLR, 2021.
[41] J. Zhang and L. Xiao. A composite randomized incremental gradient method. In International Conference on Machine Learning, pages 7454-7462, 2019.
[42] J. Zhang, P. Xiao, R. Sun, and Z. Luo. A single-loop smoothed gradient descent-ascent algorithm for nonconvex-concave min-max problems. Advances in Neural Information Processing Systems, 33:7377-7389, 2020.