SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Remark 4.6. We choose  $b = O(1)$ ,  $B = O(1)$ , and  $m = O(\varepsilon^{-2})$ . According to Corollary 4.5, to make  $\frac{1}{T}\sum_{t=0}^{T-1}\|\nabla F(\bar{x}_t)\|^2 \leq \varepsilon^2$ , we get  $T = O(N^{-1}\varepsilon^{-4})$ . Given  $q = (T/N^3)^{1/4}$ , the communication complexity is  $\frac{T}{q} = (NT)^{3/4} = O(\varepsilon^{-3})$ . The sample complexity is  $mT = O(N^{-1}\varepsilon^{-6})$  which indicates FCSG-M also has the linear speedup with respect to the number of clients.
4.3 Acc-FCSG-M
In the single-machine setting, [17] presents that under the general nonconvex conditional stochastic optimization objective, the lower bound of sample complexity is  $O(\varepsilon^{-5})$ . It means that the sample complexity achieved by FCSG and FCSG-M could be improved for nonconvex smooth conditional stochastic optimization objectives. To match the above lower bound of sample complexity, we propose an accelerated version of the FCSG-M (Acc-FCSG-M) based on the momentum-based variance reduction technique. The details of the method are shown in Algorithm 2.
Similar to the FCSG-M, in the beginning, each client initializes the model parameters and utilizes the stochastic gradient descent method to calculate the gradient estimator. Subsequently, in every  $q$  iterations, all clients perform communication with the central server, and the model parameters and gradient estimators are averaged. The key difference is that at Line 14 in Acc-FCSG-M, we use the momentum-based variance reduction gradient estimator  $u_{t + 1}^n$  to track the gradient and update the model. where  $\beta \in (0,1)$ . We establish the theoretical convergence guarantee of our new algorithm Acc-FCSG-M. All proofs are provided in the supplementary materials.
Algorithm 2 Acc-FCSG-M Algorithm
1: Input:  $T$ , momentum weight  $\beta$ , learning rate  $\alpha$ , the number of local updates  $q$ , inner batch size  $m$  and outer batch size  $b$ , as well as the initial outer batch size B;  
2: Initialize:  $x_0^n = \frac{1}{N}\sum_{k=1}^N x_0^n$ . Draw  $B$  samples of  $\{\xi_1^n, \dots, \xi_B^n\}$  and draw  $m$  samples  $\mathcal{B}_{0,i}^n = \{\eta_{ij}^n\}_{j=1}^m$  from  $P(\eta^n \mid \xi_i^n)$  for each  $\xi_i^n \in \{\xi_1^n, \dots, \xi_B^n\}$ , then  $u_1^n = \frac{1}{B}\sum_{i=1}^B \nabla \hat{F}^n(x_0^n; \xi_0^n, \mathcal{B}_{0,i}^n)$  for  $n \in [N]$ .  
3: for  $t = 1, 2, \dots, T$  do  
4: for  $n = 1, 2, \dots, N$  do  
5: if mod  $(t, q) = 0$  then  
6: Server Update:  
7:  $u_t^n = \bar{u}_t = \frac{1}{N}\sum_{i=1}^N u_t^n$   
8:  $x_t^n = \bar{x}_t = \frac{1}{N}\sum_{n=1}^N (x_{t-1}^n - \alpha u_t^n)$   
9: else  
10:  $x_{t,n} = x_{t-1}^n - \alpha u_t^n$   
11: end if  
12: Draw  $b$  samples of  $\{\xi_{t,1}^n, \dots, \xi_{t,b}^n\}$   
13: Draw  $m$  samples  $\mathcal{B}_{t,n}^n = \{\eta_{ij}^n\}_{j=1}^m$  from  $P(\eta^n \mid \xi_{t,i}^n)$  for each  $\xi_{t,i}^n \in \{\xi_{t,1}^n, \dots, \xi_{t,B}^n\}$ ,  
14:  $u_{t+1}^n = \frac{1}{b}\sum_{i=1}^b \nabla \hat{F}^n(x_t^n; \xi_{t,i}^n, \mathcal{B}_{t,i}^n) + (1 - \beta)(u_t^n - \frac{1}{b}\sum_{i=1}^b \nabla \hat{F}^n(x_{t-1}^n; \xi_{t,i}^n, \mathcal{B}_{t,i}^n))$   
15: end for  
16: end for  
17: Output:  $x$  chosen uniformly random from  $\{\bar{x}_t\}_{t=1}^T$ .
Theorem 4.7. Suppose Assumptions 3.1, 3.2 and 3.3 hold,  $\alpha \leq \frac{1}{6qS_F}$  and  $\beta = 5S_F\alpha$ . Acc-FCSG-M has the following
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} \| \nabla F (\bar {x} _ {t}) \| ^ {2} \leq \frac {2 [ F (\bar {x} _ {0}) - F (\bar {x} _ {T}) ]}{T \alpha} + \frac {3 L _ {f} ^ {2} L _ {g} ^ {2}}{\beta B N T} + \frac {1 3 L _ {f} ^ {2} L _ {g} ^ {2} c ^ {2}}{6 S _ {F} ^ {2}} \alpha^ {2} + \frac {3 L _ {g} ^ {2} S _ {f} ^ {2} \sigma_ {g} ^ {2}}{m} + \frac {6 \beta L _ {f} ^ {2} L _ {g} ^ {2} c ^ {2}}{N b}.
$$
Corollary 4.8. We choose  $q = \left(T / N^2\right)^{1 / 3}$ . Therefore,  $\alpha = \frac{1}{12qS_F} = \frac{N^{2 / 3}}{12S_F T^{1 / 3}}$ , since  $c = \frac{30S_F^2}{bN}$ , we have  $\beta = c\alpha^2 = \frac{5N^{1 / 3}}{24T^{2 / 3}b}$ . And let  $B = \frac{T^{1 / 3}}{N^{2 / 3}}$ . Therefore, we have
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \| \nabla F (\bar {x} _ {t}) \| ^ {2} \leq \frac {2 4 S _ {F} [ F (\bar {x} _ {0}) - F (\bar {x} _ {*}) ]}{(N T) ^ {2 / 3}} \\ + \frac {7 2 L _ {f} ^ {2} L _ {g} ^ {2} b}{5 (N T) ^ {2 / 3}} + \frac {3 2 5 L _ {f} ^ {2} L _ {g} ^ {2}}{2 4 b ^ {2} (T N) ^ {2 / 3}} + \frac {3 L _ {g} ^ {2} S _ {f} ^ {2} \sigma_ {g} ^ {2}}{m} + \frac {5 L _ {f} ^ {2}}{4 b ^ {2} (N T) ^ {2 / 3}} \\ \end{array}
$$
Remark 4.9. We choose  $\mathbf{b}$  as  $O(1)(b \geq 1)$  and  $m = O(\varepsilon^{-2})$ . According to Corollary 4.8 to make  $\frac{1}{T}\sum_{t=0}^{T-1}\|\nabla F(\bar{x}_t)\|^2 \leq \varepsilon^2$ , we get  $T = O(N^{-1}\varepsilon^{-3})$  and  $\frac{T}{q} = (NT)^{2/3} = O(\varepsilon^{-2})$ . The sample complexity is  $O(N^{-1}\varepsilon^{-5})$ . The communication complexity is  $O(\varepsilon^{-2})$ .  $T = O(N^{-1}\varepsilon^{-3})$  indicates the linear speedup of our algorithm.
The experiments are run on CPU machines with AMD EPYC 7513 32-Core Processors as well as NVIDIA RTX A6000. The code is available * and Federated Online AUPRC maximization task follow [38] *.
5.1 Invariant Logistic Regression
Invariant learning has an important role in robust classifier training [27]. In this section, We compare the performance of our algorithms, FCSG and FCSG-M, on the distributed invariant logistic regression
*https://github.com/xidongwu/Federated-Minimax-and-Conditional-Stochastic-Optimization/tree/main
*https://github.com/xidongwu/D-AUPRC
(a)  $m = 1$
(b)  $m = 10$
Figure 1: Test accuracy vs the number of communication rounds for different inner mini-batch ( $m = 1, 10, 100$ ) under different noise ratios ( $\sigma_2 / \sigma_1 = 1, 1.5, 2$ ).
(c)  $m = 100$
to evaluate the benefit from momentum and the effect of inner batch size, and the problem was formulated by [16]:
$$
\min  _ {x} \frac {1}{N} \sum_ {n = 1} ^ {N} \mathbb {E} _ {\xi^ {n} = (a, b)} [ l _ {n} (x) + g (x) ] \tag {4}
$$
$$
\text {w h e r e} l _ {n} (x) = \log \left(1 + \exp \left(- b \mathbb {E} _ {\eta^ {n} | \xi^ {n}} [ (\eta^ {n}) ^ {\top} x ]\right)\right) \quad g (x) = \lambda \sum_ {i = 1} ^ {d} \frac {\gamma x _ {i} ^ {2}}{1 + \gamma x _ {i} ^ {2}}
$$
where  $l_{n}(x)$  is the logistic loss function and  $g(x)$  is a non-convex regularization. We follow the experimental protocols in [17] and set the dimension of the model as 10 over 16 clients. We construct the dataset  $\xi^n = (a,b)$  and  $\eta$  as follows: We sample  $a\sim N(0,\sigma_1^2 I_d)$ , set  $b = \{\pm 1\}$  according to the sign of  $a^\top x^*$ , then sample  $\eta_{ij}^{n}\sim N(a,\sigma_{2}^{2}I_{d})$ . We choose  $\sigma_{1} = 1$ , and consider the  $\sigma_{2} / \sigma_{1}$  from  $\{1,1.5,2\}$ . At each local iteration, we use a fixed mini-batch size  $m$  from  $\{1,10,100\}$ . The outer batch size is set as 1. We test the model with 50000 outer samples to report the test accuracy. We carefully tune hyperparameters for both methods.  $\lambda = 0.001$  and  $\alpha = 10$ . We run a grid search for the learning rate and choose the learning rate in the set  $\{0.01,0.005,0.001\}$ .  $\beta$  in FCSG-M are chosen from the set  $\{0.001,0.01,0.1,0.5,0.9\}$ . The local update step is set as 50. The experiments are presented in Figure 1.
Figure 1 shows that when the noise ratio  $\sigma_2 / \sigma_1$  increases, larger inner samples  $m$  are needed, as suggested by the theory because a large batch size could reduce sample noise. In addition, when  $m = 100$ , FCSG and FCSG-M have similar performance. However, when batch size is small, compared with FCSG, FCSG-M has a more stable performance, because FCSG-M can use historic information to reduce the effect of stochastic gradient noise.
5.2 Federated Model-Agnostic Meta-Learning
Next, we evaluate our proposed algorithms with the few shot image classification task over the dataset with baselines: Local-SCGD and Local-SCGDM [11]. MOML [33] is not suitable for this task since MOML requires maintaining an inner state for each task which is not permitted due to the large number of tasks in the Omniglot dataset. Local-SCGD is the federated version of SCGD [36]. This task can be effectively solved via Model-Agnostic Meta-Learning [10].
Meta-learning aims to train a model on various learning tasks, such that the model can easily adapt to a new task using few training samples. Model-agnostic meta-learning (MAML) [10] is a popular meta-learning method to learn a good initialization with a gradient-based update, which can be formulated as the following conditional stochastic optimization problem:
$$
\min _ {x} \mathbb {E} _ {i \sim \mathcal {P} _ {\mathrm {t a s k}}, a \sim D _ {\mathrm {q u e r y}} ^ {i}} \mathcal {L} _ {i} \left(\mathbb {E} _ {b \sim D _ {\mathrm {s u p p o r t}} ^ {i}} \left(x - \lambda \nabla \mathcal {L} _ {i} (x, b)\right), a\right)
$$
where  $\mathcal{P}_{\mathrm{task}}$  denotes the learning tasks distribution,  $D_{\mathrm{support}}^i$  and  $D_{\mathrm{query}}^i$  are support (training) dataset and query (testing) dataset of the learning task  $i$ , respectively.  $\mathcal{L}_i(\cdot ,D_i)$  is the loss function on dataset  $D_{i}$  of task  $i$ . And the  $\lambda$  is a fixed meta step size. Assume  $\xi = (i,a)$  and  $\eta = b$ , the MAML problem is an example of conditional stochastic optimization where the support (training) samples in the inner
Figure 2: Convergence results of the 5-way-1-shot case over Omniglot Dataset.
Figure 3: Convergence results of the 5-way-5-shot case over Omniglot Dataset.
layer for the meta-gradient update are drawn from the conditional distribution of  $P(\eta \mid \xi)$  based on the sampled task in the outer layer.
Given there are a large number of pre-train tasks in MAML, federated learning is a good training strategy to improve efficiency because we can evenly distribute tasks over various clients and the global server coordinates clients to learn a good initial model collaboratively like MAML. Therefore, in Federated MAML, it is assumed that each device has part of the tasks. The optimization problem is defined as follows:
$$
\min  _ {x \in \mathbb {R} ^ {d}} \frac {1}{N} \sum_ {n = 1} ^ {N} F ^ {n} (x) \triangleq \frac {1}{N} \sum_ {n = 1} ^ {N} f ^ {n} \left(g ^ {n} (x)\right) \tag {5}
$$
$$
\mathrm {w h e r e} g ^ {n} (x) = \mathbb {E} _ {\eta^ {n} \sim \mathcal {D} _ {i, \mathrm {s u p p o r t}} ^ {n}} \left[ x - \lambda \nabla \mathcal {L} _ {i} ^ {n} \left(x; \eta^ {n}\right) \right], f ^ {n} (y) = \mathbb {E} _ {i \sim \mathcal {P} _ {\mathrm {t a k}} ^ {n}, a ^ {n} \sim \mathcal {D} _ {i, \mathrm {q u e r y}} ^ {n}} \mathcal {L} _ {i} ^ {n} \left(y; a ^ {n}\right).
$$
In this part, we apply our methods to few-shot image classification on the Omniglot [24, 10]. The Omniglot dataset contains 1623 different handwritten characters from 50 different alphabets and each of the 1623 characters consists of 20 instances drawn by different persons. We divide the characters to train/validation/test with 1028/172/423 by Torchmeta [7] and tasks are evenly partitioned into disjoint sets and we distribute tasks randomly among 16 clients. We conduct the fast learning of N-way-K-shot classification following the experimental protocol in [31]. The N-way-K-shot classification denotes we sample N unseen classes, randomly provide the model with K different instances of each class for training, and evaluate the model's ability to classify with new instances from the same N classes. We sample 15 data points for validation. We use a 4-layer convolutional neural model where each layer has  $3 \times 3$  convolutions and 64 filters, followed by a ReLU nonlinearity and batch normalization [10]. The images from Omniglot are downsampled to  $28 \times 28$ . For all methods, the model is trained using a single gradient step with a learning rate of 0.4. The model was evaluated using 3 gradient steps [10]. Then we use grid search and carefully tune other hyper-parameters for each method. We choose the learning rate from the set  $\{0.1, 0.05, 0.01\}$  and  $\eta$  as 1 [11]. We select the inner state momentum coefficient for Local-SCGD and Local-SCGDM from  $\{0.1, 0.5, 0.9\}$  and outside momentum coefficient for Local-SCGDM, FCSG-M and Acc-FCSG-M from  $\{0.1, 0.5, 0.9\}$ .
Figures 2 and 3 show experimental results in the 5-way-1-shot and 5-way-5-shot cases, respectively. Results show that our algorithms outperform baselines by a large margin. The main reason for Local-SCGD and Local-SCGDM to have bad performance is that converting the MAML optimization to the stochastic compositional optimization is unreasonable. It ignores the effect of task sampling on the training and inner training data distribution changes based on the sampled tasks in the outer layer. The use of the momentum-like inner state to deal with the MAML will slow down the convergence and we have to tune the extra momentum coefficient for the inner state. In addition, the momentum-like inner state also introduces extra communication costs because the server needs to average the inner state as in Local-SCGDM. In addition, comparing the results in Figures 2 and 3, we can see when
Table 2: Final averaged AP scores on the testing data.
Datasets	FedAvg	CODA+	FCSG	FCSG-M	Acc-FCSG-M
MNIST	0.9357	0.9733	0.9868	0.9878	0.9879
CIFAR-10	0.5059	0.6039	0.7130	0.7157	0.7184
the K increases in the few-shot learning, the training performance is improved, which matches the theoretical analysis that a large inner batch-size  $m$  benefits the model training.
5.3 Federated Online AUPRC maximization
AUROC maximization in FL has been studied in [13, 42, 40] and AUPRC maximization is also used to solve the imbalanced classification. Existing AUPRC algorithms maintain an inner state for each data point. [38] consider the online AUPRC in the decentralized learning. For the large-scale distributed data over multiple clients, algorithms for online AUPRC maximization in FL is necessary.
Following [28] and [38], the surrogate function of average precision (AP) for online AUPRC maximization is:
$$
\hat {\mathbf {A}} \mathbf {P} = \mathbb {E} _ {\xi^ {+} \sim \mathcal {D} ^ {+}} \frac {\mathbb {E} _ {\xi \sim \mathcal {D}} \mathbf {I} (y = 1) \ell (x ; z ^ {+} , z)}{\mathbb {E} _ {\xi \sim \mathcal {D}} \ell (x ; z ^ {+} , z)}
$$
where  $\ell(x;z^{+},z) = (\max\{s - h(x;z^{+}) + h(x;z), 0\})^{2}$  and  $h(x;z)$  is the prediction score function of input  $z$  with model  $x$ . Federated Online AUPRC maximization could be reformulated as:
$$
\min  _ {x} F (x) = \min  _ {x} \frac {1}{N} \sum_ {n = 1} ^ {n} \mathbb {E} _ {\xi_ {n} \sim \mathcal {D} _ {n} ^ {+}} f \left(\mathbb {E} _ {\xi_ {n} ^ {\prime} \sim \mathcal {D} _ {n}} g ^ {n} (\mathbf {x}; \xi_ {n}, \xi_ {n} ^ {\prime})\right) \tag {6}
$$
where  $\xi_{n} = (z_{n},y_{n})\sim \mathcal{D}_{n}$  and  $\xi_n^+ = (z_n^+,y_n^+)\sim \mathcal{D}_n^+$  are samples drawn from the whole datasets and positive datasets, respectively. It is a two-level problem and the inner objective depends on both  $\xi$  and  $\xi^{+}$ . Since federated online AUPRC is a special example of a federated conditional stochastic optimization, our algorithms could be directly applied to it.
We choose MNIST dataset and CIFAR-10 datasets. As AUROC maximization in federated settings has been demonstrated in existing works [13, 42], we use  $\mathrm{CODA + }$  in [42] as a baseline. Another baseline is the FedAvg with cross-entropy loss. Since AUPRC is used for binary classification, the first half of the classes in the MNIST and CIFAR10 datasets are designated to be the negative class, and the rest half of the classes are considered to be the positive class. Then, we remove  $80\%$  of the positive examples in the training set to make it imbalanced, while keeping the test set unchanged. The results in Table 2 show that our algorithms could be used to solve the online AUPRC maximization in FL and it largely improves the model's performance.
In this paper, we studied federated conditional stochastic optimization under the general nonconvex setting. To the best of our knowledge, this is the first paper proposing algorithms for the federated conditional stochastic optimization problem. We first used the biased stochastic first-order gradient to design an algorithm called FCSG, which we proved to have a sample complexity of  $O(\epsilon^{-6})$ , and communication complexity of  $O(\epsilon^{-3})$  to reach an  $\epsilon$ -stationary point. FCSG enjoys an appealing linear speedup with respect to the number of clients. To improve the empirical performances of FCSG, we also proposed a novel algorithm (i.e., FCSG-M), which achieves the same theoretical guarantees as FCSG. To fill the gap from lower-bound complexity, we introduced an accelerated version of FCSG-M, called Acc-FCSG-M, using variance reduction technique, which is optimal for the nonconvex smooth federated conditional stochastic optimization problems as it matches the best possible complexity result achieved by the single-machine method. It has a sample complexity of  $O(\epsilon^{-5})$ , and communication complexity of  $O(\epsilon^{-2})$ . And the communication complexity achieves the best communication complexity of the nonconvex optimization in federated learning. Experimental results on the machine learning tasks validate the effectiveness of our algorithms.
Limitation Conditional stochastic optimization has much broader applications and a more comprehensive evaluation of our proposed algorithms on other use cases would be a promising future direction. In addition, FCSG-M and accelerated FCSG-M require higher communication overhead.
[39] Xidong Wu, Feihu Huang, Zhengmian Hu, and Heng Huang. Faster adaptive federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 10379-10387, 2023.
[40] Xidong Wu, Jianhui Sun, Zhengmian Hu, Aidong Zhang, and Heng Huang. Solving a class of nonconvex minimax optimization in federated learning. Advances in Neural Information Processing Systems (NeurIPS), 2023.
[41] Huizhuo Yuan and Wenqing Hu. Stochastic recursive momentum method for non-convex compositional optimization. arXiv preprint arXiv:2006.01688, 2020.
[42] Zhuoning Yuan, Zhishuai Guo, Yi Xu, Yiming Ying, and Tianbao Yang. Federated deep auc maximization for heterogeneous data with a constant communication complexity. In International Conference on Machine Learning, pages 12219-12229. PMLR, 2021.
[43] Junyu Zhang and Lin Xiao. A composite randomized incremental gradient method. In International Conference on Machine Learning, pages 7454-7462. PMLR, 2019.
[44] Hanhan Zhou, Tian Lan, Guru Prasadh Venkataramani, and Wenbo Ding. Federated learning with online adaptive heterogeneous local models. In Workshop on Federated Learning: Recent Advances and New Challenges (in Conjunction with NeurIPS 2022), 2022.
[45] Hanhan Zhou, Tian Lan, Guru Prasadh Venkataramani, and Wenbo Ding. Every parameter matters: Ensuring the convergence of federated learning with dynamic heterogeneous models reduction. In Advances in Neural Information Processing Systems, volume 36, 2023.