SlowGuess/NeurIPS_2023_OCR huggingface dataset:
we denote the average of the momentum similarly as  $\bar{\omega}_t, \bar{\nu}_t$  and  $\bar{q}_t$ . Then we consider the following Lyapunov function  $\mathcal{G}_t$ :
$$
\begin{array}{l} \mathcal {G} _ {t} = h (\bar {x} _ {t}) + \frac {1 8 \eta \tilde {L} ^ {2}}{\mu \gamma} (\| \bar {y} _ {t} - y _ {\bar {x} _ {t}} \| ^ {2} + \| \bar {u} _ {t} - u _ {\bar {x} _ {t}} \| ^ {2}) + \frac {9 b M \eta}{6 4 \alpha_ {t}} \| \bar {\omega} _ {t} - \frac {1}{M} \sum_ {m = 1} ^ {M} \nabla_ {y} g ^ {(m)} (x _ {t} ^ {(m)}, y _ {t} ^ {(m)}) \| ^ {2} \\ + \frac {9 b M \eta}{6 4 \alpha_ {t}} \| \bar {q} _ {t} - \frac {1}{M} \sum_ {m = 1} ^ {M} (\nabla_ {y ^ {2}} g ^ {(m)} (x _ {t} ^ {(m)}, y _ {t} ^ {(m)}) u _ {t} ^ {(m)} - \nabla_ {y} f ^ {(m)} (x _ {t} ^ {(m)}, y _ {t} ^ {(m)}))) \| ^ {2} \\ + \frac {9 b M \eta}{6 4 \alpha_ {t}} \| \bar {\nu} _ {t} - \frac {1}{M} \sum_ {m = 1} ^ {M} \left(\nabla_ {x} f ^ {(m)} \left(x _ {t} ^ {(m)}, y _ {t} ^ {(m)}\right) - \nabla_ {x y} g ^ {(m)} \left(x _ {t} ^ {(m)}, y _ {t} ^ {(m)}\right) u _ {t} ^ {(m)}\right) \| ^ {2} \tag {9} \\ \end{array}
$$
where  $y_{\bar{x}_t}$  denotes the solution of the lower level problem  $g(\bar{x}_t,\cdot),u_{\bar{x}_t} = [\nabla_{y^2}g(\bar{x},y_{\bar{x}})]^{-1}\nabla_yf(\bar{x},y_{\bar{x}})$  denotes the solution of Eq (4) at state  $\bar{x}_t$  . Besides,  $\gamma ,\eta ,\tau$  are learning rates and  $L,\tilde{L}$  are constants. Note that the first three terms of  $\mathcal{G}_t$  ..  $h(\bar{x}_t)$ $\| \bar{y}_t - y_{\bar{x}_t}\| ^2$ $\| \bar{u}_t - u_{\bar{x}_t}\| ^2$  measures the errors of three federated problems: the upper level problem, the lower level problem and the hyper-gradient estimation. Then the last three terms measure the estimation error of the momentum variables:  $\bar{\omega}_t$ $\bar{\nu}_{t}$  and  $\bar{q}_t$  . The convergence proof primarily concentrates on bounding these errors, please see Lemma C.2 - C.6 in the Appendix for more details. Meanwhile, as in the single level federated optimization problems, local updates lead to client-drift error. More specifically, we need to bound  $\| x_t^{(m)} - \bar{x}_t\| ^2$ $\| y_t^{(m)} - \bar{y}_t\| ^2$  and  $\| u_t^{(m)} - \bar{u}_t\| ^2$  , please see Lemma C.7 - C.11 for more details. Finally, we have the following convergence theorem:
Theorem 3.6. Suppose in Algorithm 1, we choose learning rate  $\alpha_{t} = \frac{\delta}{(u + t)^{1 / 3}}$ ,  $t\in [T]$ , for some constant  $\delta$  and  $u$ , and let  $c_{\nu}$ ,  $c_{\omega}$ ,  $c_{u}$  choose some value,  $\eta$ ,  $\gamma$  and  $\tau$ ,  $r$  be some small values decided by the Lipschitz constants of  $h(x)$ , we choose the minibatch size to be  $b_{x} = b_{y} = b$  and the first batch to be  $b_{1} = O(Ib)$ , then we have:
$$
\frac {1}{T} \sum_ {t = 1} ^ {T - 1} \mathbb {E} \left[ \| \nabla h (\bar {x} _ {t}) \| ^ {2} \right] = O \left(\frac {\kappa^ {1 9 / 3} I}{T} + \frac {\kappa^ {1 6 / 3}}{(b M T) ^ {2 / 3}}\right)
$$
To reach an  $\epsilon$ -stationary point, we need  $T = O(\kappa^{8}(bM)^{-1}\epsilon^{-1.5})$ ,  $I = O(\kappa^{5/3}(bM)^{-1}\epsilon^{-0.5})$ .
As stated in the Theorem, to reach an  $\epsilon$ -stationary point, we need  $T = O(\kappa^{8}(bM)^{-1}\epsilon^{-1.5})$ , then the sample complexity for each client is  $Gc(f,\epsilon) = O(M^{-1}\kappa^{8}\epsilon^{-1.5})$ ,  $Gc(g,\epsilon) = O(M^{-1}\kappa^{8}\epsilon^{-1.5})$ ,  $Jv(g,\epsilon) = O(M^{-1}\kappa^{8}\epsilon^{-1.5})$ ,  $Hv(g,\epsilon) = O(M^{-1}\kappa^{8}\epsilon^{-1.5})$ . So FedBiOAcc achieves the linear speed up w.r.t. to the number of clients  $M$ . Next, suppose we choose  $I = O(\kappa^{5/3}(bM)^{-1}\epsilon^{-0.5})$ , then the number of communication round  $E = O(\kappa^{19/3}\epsilon^{-1})$ . This matches the optimal communication complexity of the single level optimization problems as in the STEM [27]. Furthermore, compared to FedNest and its variants, FedBiOAcc has improved both the communication complexity and the iteration complexity. As for LocalBSCVR [12], FedBiOAcc obtains same rate, but incorporates the heterogeneous case. Note that it is much more challenging to analyze the heterogeneous case. In fact, if we assume homogeneous clients, we have local hyper-gradient (Eq. (3)) equals the global hyper-gradient (Eq. (2)), then we do not need to use the quadratic federated optimization problem view in Section 3.2, while the theoretical analysis is also simplified significantly.
4 Federated Bilevel Optimization with Local Lower Level Problems
In this section, we consider an alternative formulation of the Federated Bilevel Optimization problems as follows:
$$
\min  _ {x \in \mathbb {R} ^ {p}} h (x) := \frac {1}{M} \sum_ {m = 1} ^ {M} f ^ {(m)} \left(x, y _ {x} ^ {(m)}\right), \text {s . t .} y _ {x} ^ {(m)} = \underset {y \in \mathbb {R} ^ {d}} {\arg \min } g ^ {(m)} (x, y) \tag {10}
$$
Same as Eq. (1), Eq. (10) has a federated upper level problem, however, Eq. (10) has a unique lower level problem for each client, which is different from Eq. (1). In fact, federated bilevel optimization problem Eq (10) can be viewed as a special type of standard federated learning problems. If we denote  $h^{(m)}(x) = f^{(m)}(x,y_x^{(m)})$ , then Eq. (10) can be written as  $\min_{x\in \mathbb{R}^p}h(x)\coloneqq \frac{1}{M}\sum_{m = 1}^{M}h^{(m)}(x)$ .
But due to the bilevel structure of  $h^{(m)}(x)$ , Eq. (10) is more challenging than the standard Federated Learning problems.
Hyper-gradient Estimation. Assume Assumption 3.1~Assumption 3.3 hold, then the hyper-gradient is  $\Phi (x,y_x) = \frac{1}{M}\sum_{m = 1}^{M}\Phi^{(m)}(x,y_x)$ , where  $\Phi^{(m)}(x,y)$  is defined in Eq. (3), in other words, the local hyper-gradient  $\Phi^{(m)}(x,y)$  is an unbiased estimate of the full hyper-gradient. This fact makes it possible to solve Eq. (10) with local-sgd like methods. More specifically, we solve the local bilevel problem  $h^{(m)}(x)$  multiple steps on each client and then the server averages the local states from clients. Please refer to Algorithm 3 and the variance-reduction acceleration Algorithm 4 in the Appendix. For ease of reference, we name them FedBiO-Local and FedBiOAcc-Local, respectively.
Several challenges exist in analyzing FedBiO-Local and FedBiOAcc-Local. First, Eq. (3) involves Hessian inverse, so we only evaluate it approximately through the Neumann series [37] as:
$$
\begin{array}{l} \Phi^ {(m)} (x, y; \xi_ {x}) = \nabla_ {x} f ^ {(m)} (x, y; \xi_ {f}) - \tau \nabla_ {x y} g ^ {(m)} (x, y; \xi_ {g}) \\ \times \sum_ {q = - 1} ^ {Q - 1} \prod_ {j = Q - q} ^ {Q} (I - \tau \nabla_ {y ^ {2}} g ^ {(m)} (x, y; \xi_ {j})) \nabla_ {y} f ^ {(m)} (x, y; \xi_ {f}) \tag {11} \\ \end{array}
$$
where  $\xi_{x} = \{\xi_{j}(j = 1,\dots ,Q),\xi_{f},\xi_{g}\}$ , and we assume its elements are mutually independent.  $\Phi^{(m)}(x,y;\xi_x)$  is a biased estimate of  $\Phi^{(m)}(x,y)$ , but with bounded bias and variance (Please see Proposition D.2 for more details.) Furthermore, to reduce the computation cost, each client solves the local lower level problem approximately and we update the upper and lower level variable alternatively. The idea of alternative update is widely used in the non-distributed bilevel optimization [23, 56]. However, in the federated setting, client variables drift away when performing multiple local steps. As a result, the variable drift error and the bias caused by inexact solution of the lower level problem intertwined with each other. For example, in the local update, clients optimize the lower level variable  $y^{(m)}$  towards the minimizer  $y_{x^{(m)}}^{(m)}$ , but after the communication step,  $x^{(m)}$  is smoothed among clients, as a result, the target of  $y_{t}^{(m)}$  changes which causes a huge bias.
In the appendix, we show the FedBiOAcc-Local algorithm achieves the same optimal convergence rate as FedBiOAcc, which has iteration complexity  $\bar{O}(\epsilon^{-1.5})$  and communication complexity  $O(\epsilon^{-1})$ . However, since the lower level problem in Eq. (10) is unique for each client, FedBiOAcc-Local does not have the property of linear speed-up w.r.t the number of clients as FedBiOAcc does.
In this section, we assess the performance of the proposed FedBiOAcc algorithm through two federated bilevel tasks: Federated Data Cleaning and Federated Hyper-representation Learning. The Federated Data Cleaning task involves global lower level problems, while the Hyper-representation Learning task involves local lower level problems. The implementation is carried out using PyTorch, and the Federated Learning environment is simulated using the PyTorch.Distributed package. Our experiments were conducted on servers equipped with an AMD EPYC 7763 64-core CPU and 8 NVIDIA V100 GPUs.
Figure 1: Validation Error vs Communication Rounds. From Left to Right:  $\rho = 0.1, 0.4, 0.8, 0.95$ . The local step  $I$  is set as 5 for FedBiO, FedBiOAcc and FedAvg.
Figure 2: Validation Error vs Communication Rounds with different number of clients per epoch. From Left to Right:  $\rho = 0.1, 0.4, 0.8, 0.95$ . The local step  $I$  is set as 5.
5.1 Federated Data Cleaning
In this section, we consider the Federated Data Cleaning task. In this task, we are given a noisy training dataset whose labels are corrupted by noise and a clean validation set. Then we aim to find weights for training samples such that a model that is learned over the weighted training set performs well on the validation set. This is a federated bilevel problem when the noisy training set is distributed over multiple clients. The formulation of the task is included in Appendix B.1. This task is a specialization of Eq. (1).
Dataset and Baselines. We create 10 clients and construct datasets based on MNIST [31]. For the training set, each client randomly samples 4500 images (no overlap among clients) from 10 classes and then randomly uniformly perturb the labels of  $\rho$  ( $0 \leq \rho \leq 1$ ) percent samples. For the validation set, each client randomly selects 50 clean images from a different class. In other words, the  $m_{th}$  client only has validation samples from the  $m_{th}$  class. This single-class validation setting introduces a high level of heterogeneity, such that individual clients are unable to conduct local cleaning due to they only have clean samples from one class. In our experiments, we test our FedBiOAcc algorithm, including the FedBiO algorithm (Algorithm 2 in Appendix) which does not use variance reduction; additionally, we also consider some baseline methods: a baseline that directly performs FedAvg [40] on the noisy dataset, this helps to verify the usefulness of data cleaning; Local-BSGVR [12], FedNest [49], CommFedBiO [33], AggITD [53] and FedMBO [21]. Note that Local-BSGVR is designed for the homogeneous setting, and the last four baselines all need multiple rounds of client-server communication to evaluate the hyper-gradient at each global epoch. We perform grid search to find the best hyper-parameters for each method and report the best results. Specific choices are included in Appendix B.1.
In figure 1, we compare the performance of different methods at various noise levels  $\rho$ . Note that the larger the  $\rho$  value, the more noisy the training data are. The noise level can be illustrated by the performance of the FedAvg algorithm, which learns over the noisy data directly. As shown in the figure, FedAvg learns almost nothing when  $\rho = 0.95$ . Next, our algorithms are robust under various heterogeneity levels. When the noise level in the training set increases as the value of  $\rho$  increases, learning relies more on the signal from the heterogeneous validation set, and our algorithms consistently outperform other baselines. Finally, in figure 2, we vary the number of clients sampled per epoch, and the experimental results show that our FedBiOAcc converges faster with more clients in the training per epoch; in figure 3, we vary the number of local steps under different noisy levels. Interestingly, the algorithm benefit more from the local training under larger noise.
Figure 3: Validation Error vs Communication Rounds with different number of local steps  $I$ . From Left to Right:  $\rho = 0.1, 0.4, 0.8, 0.95$ .
Figure 4: Validation Error vs Communication Rounds. The top row shows the result for the Omniglot Dataset and the bottom row shows MiniImageNet. From Left to Right: 5-way-1-shot, 5-way-5-shot, 20-way-1-shot, 20-way-5-shot. The local step  $I$  is set to 5.
5.2 Federated Hyper-Representation Learning
In the Hyper-representation learning task, we learn a hyper-representation of the data such that a linear classifier can be learned quickly with a small number of data samples. A mathematical formulation of the task is included in Appendix B.2. Note that this task is an instantiation of Eq. (10), due to the fact that each client has its own tasks, and thus only the upper level problem is federated. We consider the Omniglot [29] and MiniImageNet [44] data sets. As in the non-distributed setting, we perform  $N$ -way- $K$ -shot classification.
In this experiment, we compare FedBiOAcc-Local (Algorithm 4 in the Appendix) with three baselines FedBiO-Local (Algorithm 3 in the Appendix), DistBiO and DistBiOAcc. Note that DistBiO and DistBiOAcc are the distributed version of FedBiO-Local and FedBiOAcc-Local, respectively. In the experiments, we implement DistBiO and DistBiOAcc by setting the local steps as 1 for FedBiO-Local and FedBiOAcc-Local. We perform grid search for the hyper-parameter selection for both methods and choose the best ones, the specific choices of hyper-parameters are deferred to Appendix B.2. The results are summarized in Figure 4 (full results are included in Figure 5 and Figure 6 of Appendix. As shown by the results, FedBiOAcc converges faster than the baselines on both datasets and on all four types of classification tasks, which demonstrates the effectiveness of variance reduction and multiple steps of local training.
In this paper, we study the Federated Bilevel Optimization problems and introduce FedBiOAcc. In particular, FedBiOAcc evaluates the hyper-gradient by solving a federated quadratic problem, and mitigates the noise through momentum-based variance reduction technique. We provide a rigorous convergence analysis for our proposed method and show that FedBiOAcc has the optimal iteration complexity  $O(\epsilon^{-1.5})$  and communication complexity  $O(\epsilon^{-1})$ , and it also achieves linear speed-up w.r.t the number of clients. Besides, we study a type of novel Federated Bilevel Optimization problems with local lower level problems. We modify FedBiO for this type of problems and propose
FedBiOAcc-Local. FedBiOAcc-Local achieves the same optimal convergence rate as FedBiOAcc. Finally, we validate our algorithms with real-world tasks.
[1] A. K. R. Bayoumi, K. Mishchenko, and P. Richtarik. Tighter theory for local sgd on identical and heterogeneous data. In International Conference on Artificial Intelligence and Statistics, pages 4519-4529, 2020.
[2] Y. Bengio. Gradient-based optimization of hyperparameters. Neural computation, 12(8):1889-1900, 2000.
[3] D. Chen and M. T. Hagan. Optimal use of regularization and cross-validation in neural network modeling. In IJCNN'99. International Joint Conference on Neural Networks. Proceedings (Cat. No. 99CH36339), volume 2, pages 1275-1280. IEEE, 1999.
[4] T. Chen, Y. Sun, and W. Yin. A single-timescale stochastic bilevel optimization method. arXiv preprint arXiv:2102.04671, 2021.
[5] X. Chen, M. Huang, and S. Ma. Decentralized bilevel optimization. arXiv preprint arXiv:2206.05670, 2022.
[6] A. Cutkosky and F. Orabona. Momentum-based variance reduction in non-convex sgd. In Advances in Neural Information Processing Systems, pages 15236–15245, 2019.
[7] M. Dagréou, P. Ablin, S. Vaiter, and T. Moreau. A framework for bilevel optimization that enables stochastic and global variance reduction algorithms. arXiv preprint arXiv:2201.13409, 2022.
[8] C. B. Do, C.-S. Foo, and A. Y. Ng. Efficient multiple hyperparameter learning for log-linear models. In NIPS, volume 2007, pages 377-384. CiteSeer, 2007.
[9] J. Domke. Generic methods for optimization-based modeling. In Artificial Intelligence and Statistics, pages 318-326. PMLR, 2012.
[10] M. C. Ferris and O. L. Mangasarian. Finite perturbation of convex programs. Applied Mathematics and Optimization, 23(1):263-273, 1991.
[11] L. Franceschi, M. Donini, P. Frasconi, and M. Pontil. Forward and reverse gradient-based hyperparameter optimization. In Proceedings of the 34th International Conference on Machine Learning-Volume 70, pages 1165–1173. JMLR.org, 2017.
[12] H. Gao. On the convergence of momentum-based algorithms for federated stochastic bilevel optimization problems. arXiv preprint arXiv:2204.13299, 2022.
[13] H. Gao, B. Gu, and M. T. Thai. On the convergence of distributed stochastic bilevel optimization algorithms over a network. In International Conference on Artificial Intelligence and Statistics, pages 9238-9281. PMLR, 2023.
[14] S. Ghadimi and M. Wang. Approximation methods for bilevel programming. arXiv preprint arXiv:1802.02246, 2018.
[15] Z. Guo, Q. Hu, L. Zhang, and T. Yang. Randomized stochastic variance-reduced methods for multi-task stochastic bilevel optimization. arXiv preprint arXiv:2105.02266, 2021.
[16] F. Haddadpour and M. Mahdavi. On the convergence of local descent methods in federated learning. arXiv preprint arXiv:1910.14425, 2019.
[17] M. Hong, H.-T. Wai, Z. Wang, and Z. Yang. A two-timescale framework for bilevel optimization: Complexity analysis and application to actor-critic. arXiv preprint arXiv:2007.05170, 2020.
[18] F. Huang. On momentum-based gradient methods for bilevel optimization with nonconvex lower-level. arXiv preprint arXiv:2303.03944, 2023.
[58] Y. Yang, P. Xiao, and K. Ji. Simfbo: Towards simple, flexible and communication-efficient federated bilevel learning. arXiv preprint arXiv:2305.19442, 2023.
[59] H. Yu, S. Yang, and S. Zhu. Parallel restarted sgd with faster convergence and less communication: Demystifying why model averaging works for deep learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 5693-5700, 2019.
[60] Y. Zhao, M. Li, L. Lai, N. Suda, D. Civin, and V. Chandra. Federated learning with non-iid data. arXiv preprint arXiv:1806.00582, 2018.
[61] L. Zintgraf, K. Shiarli, V. Kurin, K. Hofmann, and S. Whiteson. Fast context adaptation via meta-learning. In International Conference on Machine Learning, pages 7693-7702. PMLR, 2019.