SlowGuess/NeurIPS_2023_OCR huggingface dataset:
It can be seen from eq. (2) that the boundedness of  $v$  is necessary to guarantee the smoothness (w.r.t.  $x, y$ ) and bounded variance in solving the local and global LS problems. Projecting the global  $v^{(t)}$  vector and the local  $v_{i}^{(t,k)}, k \geq 1$  vectors onto a bounded set can be a feasible solution, but in this case, the local aggregation  $q_{v,i}^{(t)}$  is no longer a linear combination of local gradients. This can complicate the implementation and analysis, and degrade the flexibility of the framework. Fortunately, we show via induction that the projection of the server-side vector  $v^{(t)}$  on a bounded set suffices to guarantee the boundedness of local vectors  $v_{i}^{t,k}$ .
Proposition 1 (Boundedness of Local  $v$ ). Under Assumptions 1 and 2, for each iteration  $t$ , client  $i$ , and local iteration  $k = 1, 2, \ldots, \tau_i^{(t)}$ , we have  $r_i := \| v_i^{(t,k)}\| \leq \left(1 + \frac{\alpha_{\max}}{\alpha_{\min}}\right)r$ , where the radius  $r = \frac{L_f}{\mu_g}$  and  $\alpha_{\min}, \alpha_{\max}$  are chosen such that  $\alpha_{\min} \leq a_i^{(t,k)} \leq \alpha_{\max}$ .
Next, we show an important proposition in characterizing the per-iteration progress of the global  $v^{(t)}$  updates in approximating the solution of a reweighted global LS problem. Let  $\Delta_v^{(t)} = \mathbb{E}\| v^{(t)} - \widetilde{v}^* (x^{(t)})\|^2$  denote the approximation error, where  $\widetilde{v}^*$  be the minimizer of  $\sum_{i = 1}^{n}w_{i}R_{i}(x,\widetilde{y}^{*},\cdot)$ .
Proposition 2. Under the Assumption 1, 2 and 3, the iterates  $v^{(t)}$  in solving the global LS problem generated by Algorithm 1 satisfy
$$
\begin{array}{l} \mathbb {E} \| v ^ {(t + 1)} - \widetilde {v} ^ {*} (x ^ {(t + 1)}) \| ^ {2} - \mathbb {E} \| v ^ {(t)} - \widetilde {v} ^ {*} (x ^ {(t)}) \| ^ {2} \\ \leq (\delta_ {t} ^ {\prime} - \rho^ {(t)} \gamma_ {v} \mu_ {g} - \delta_ {t} ^ {\prime} \rho^ {(t)} \gamma_ {v} \mu_ {g}) \mathbb {E} \| v ^ {(t)} - \widetilde {v} ^ {*} (x ^ {(t)}) \| ^ {2} + (1 + \delta_ {t} ^ {\prime}) (\rho^ {(t)} \gamma_ {v}) ^ {2} \mathbb {E} \bigg \| \sum_ {i \in C ^ {(t)}} \widetilde {w} _ {i} h _ {v, i} ^ {(t)} \bigg \| ^ {2} \\ + \left(1 + \delta_ {t} ^ {\prime}\right) \rho^ {(t)} \gamma_ {v} \frac {4 L _ {R} ^ {2}}{\mu_ {g}} \mathbb {E} \left\| y ^ {(t)} - \widetilde {y} ^ {*} (x ^ {(t)}) \right\| ^ {2} + \left(\rho^ {(t)} \gamma_ {x}\right) ^ {2} \left(L _ {v} ^ {2} + \frac {L _ {v x}}{4}\right) \mathbb {E} \left\| \sum_ {i \in C ^ {(t)}} \widetilde {w} _ {i} h _ {x, i} ^ {(t)} \right\| ^ {2} \\ + (1 + \delta_ {t} ^ {\prime}) \rho^ {(t)} \gamma_ {v} \frac {4 L _ {R} ^ {2}}{\mu_ {g}} \sum_ {i = 1} ^ {n} w _ {i} \sum_ {k = 0} ^ {\tau_ {i} - 1} \frac {a _ {i} ^ {(t , k)}}{\| a _ {i} ^ {(t)} \| _ {1}} \mathbb {E} \Big [ \| x ^ {(t)} - x _ {i} ^ {(t, k)} \| ^ {2} + \| y ^ {(t)} - y _ {i} ^ {(t, k)} \| ^ {2} \\ \left. + \left\| v ^ {(t)} - v _ {i} ^ {(t, k)} \right\| ^ {2} \right] + (\rho^ {(t)}) \gamma_ {x}) ^ {2} \frac {2 L _ {v}}{\delta_ {t , 1} ^ {\prime}} \mathbb {E} \left\| \sum_ {i = 1} ^ {n} w _ {i} \widetilde {h} _ {x, i} ^ {(t)} \right\| ^ {2}. \\ \end{array}
$$
for all  $t \in \{0, 1, \dots, T - 1\}$ ,  $k \in \{0, 1, \dots, \tau_i^{(t)} - 1\}$  and  $i \in \{1, 2, \dots, n\}$ , where  $\widetilde{w}_i \coloneqq \frac{n}{|C^{(t)}|} w_i$ .
Similarly, we can provide a per-iteration process of  $y^{(t)}$  in approximating the solution  $\widetilde{y}^*$  of the reweighted lower-level global function  $\sum_{i=1}^{n} w_i g_i(x, \cdot)$ . Note that such characterizations do not exist in previous studies in single-level or minimax federated optimization with a single objective (e.g., [57]) because our analysis needs to handle three different lower-level, LS and upper-level objectives. As shown in Proposition 2, the bound involves the client drift term  $\mathbb{E}\|v^{(t)} - v_i^{(t,k)}\|^2$  (similarly for  $y, x$ ), so the next step is to characterize this important quantity.
Proposition 3. Under Assumption 1 and 2, the local iterates client drift of  $v_{i}^{(t,k)}$  is bounded as
$$
\sum_ {i = 1} ^ {n} w _ {i} \frac {1}{\| a _ {i} ^ {(t)} \| _ {1}} \sum_ {k = 1} ^ {\tau_ {i} - 1} a _ {i} ^ {(t, k)} \mathbb {E} \| v _ {i} ^ {(t, k)} - v ^ {(t)} \| ^ {2} \leq \eta_ {v} ^ {2} \bar {\tau} \sigma_ {M 1} ^ {2},
$$
for all  $t\in \{0,1,\dots,T - 1\}$ $k\in \{0,1,\dots,\tau_i - 1\}$  and  $i\in \{1,2,\dots,n\}$ . We define  $\bar{\tau} := \sum_{i = 1}^{n}\tau_{i} / n$  and  $\sigma_{M1}^{2}:= \alpha_{\max}^{2}(\sigma_{f}^{2} + r_{\max}^{2}\sigma_{gg}^{2}) + \alpha_{\max}(L_{f}^{2} + r_{\max}^{2}L_{1}^{2})$ .
It can be seen from Proposition 3 that the bound on the client drift of the local updates on  $v$  is proportional to  $\eta_v$  and  $\| a_i^{(t)}\| _1$ . Since  $\alpha_{\mathrm{min}}\leq a_i^{(t,k)}\leq \alpha_{\mathrm{max}}$ ,  $\| a_i^{(t)}\| _1$  is proportional to the number  $\tau_i^{(t)}$  of local steps. Thus, this client drift is controllable by choosing  $\tau_i^{(t)}$  and the local step sizes  $\eta_v$  properly. Then, combining the results in the above Proposition 1, 2, 3, and under a proper Lyapunov function, we obtain the following theorem. Let  $P = |C^{(t)}|$  be the number of sampled clients.
Theorem 1. Define  $\widetilde{\Phi}(x) = \widetilde{F}(x, \widetilde{y}^*)$  as the objective function by replacing  $p_i$  in eq. (1) with  $w_i$ . Suppose Assumptions 1, 2 and 3 are satisfied. The iterates by SimFBO in Algorithm 1 satisfy
$$
\min  _ {t} \mathbb {E} \left\| \nabla \widetilde {\Phi} (x ^ {(t)}) \right\| ^ {2} = \underbrace {\mathcal {O} \left(\frac {M _ {1} (n - P)}{n} \sqrt {\frac {\bar {\tau}}{P T}}\right)} _ {\text {p a r t i a l p a r t i c i p a t i o n e r r o r}} + \underbrace {\mathcal {O} \left(M _ {2} \sqrt {\frac {1}{P \bar {\tau} T}}\right)} _ {\text {f u l l s y n c h r o n i z a t i o n e r r o r}} + \underbrace {\mathcal {O} \left(\frac {M _ {3}}{\bar {\tau} T}\right)} _ {\text {l o c a l u p d a t e s e r r o r}}, \tag {9}
$$
where  $\gamma_{x},\gamma_{y},\gamma_{v},\eta_{x},\eta_{y},\eta_{v}$  are set in eq. (38) and  $M_1,M_2,M_3$  are defined by eq. (40) in appendix. For the full client participation (i.e.,  $P = n$ ), the sample complexity is  $\bar{\tau} T = \mathcal{O}(n^{-1}\epsilon^{-2})$ , and the number of communication rounds is  $T = \mathcal{O}(\epsilon^{-1})$ . For partial client participation, the sample complexity is  $\bar{\tau} T = \mathcal{O}(P^{-1}\epsilon^{-2})$ , and the number of communication rounds is  $T = \mathcal{O}(P^{-1}\epsilon^{-2})$ .
First, when set  $\bar{\tau} = \mathcal{O}(1)$ , Theorem 1 shows that SimFBO converges to a stationary point of an objective function  $\widetilde{\Phi}(x)$  with a rate of  $\mathcal{O}\left(\frac{1}{\sqrt{PT}} + \frac{1}{T}\right)$ , which, to the best of our knowledge, is the first linear speedup result under partial client participation without replacement. Note that without system-level heterogeneity, i.e.,  $\|a_1^{(t)}\| = \ldots = \|a_n^{(t)}\|$ ,  $w_i = \frac{p_i \|a_i^{(t)}\|_1}{\sum_{j=1}^n p_j \|a_j^{(t)}\|_1} = p_i$ , and hence SimFBO converges to the stationary point of the original objective in eq. (1). However, in the presence of system-level heterogeneity, SimFBO may converge to the stationary point of a different objective. Second, when nearly full clients participate, the partial participation error is approximately zero. Then we can see that setting local update round  $\bar{\tau}$  to its upper-bound results in the best performance.
Theorem 2. Define  $\Phi(x) = F(x, y^*)$  as eq. (1). Suppose Assumptions 1, 2 and 3 are satisfied. The iterates generated by ShroFBO in Algorithm 1 satisfy
$$
\min  _ {t} \mathbb {E} \left\| \nabla \Phi (x ^ {(t)}) \right\| ^ {2} = \mathcal {O} \left(\frac {M _ {1} (n - P)}{n} \sqrt {\frac {\bar {\tau}}{P T}}\right) + \mathcal {O} \left(M _ {2} \sqrt {\frac {1}{P \bar {\tau} T}}\right) + \mathcal {O} \left(\frac {M _ {3}}{\bar {\tau} T}\right), \tag {10}
$$
by setting the same server-side and local step sizes and  $M_1, M_2$  and  $M_3$  as in Theorem 1. For full client participation, the sample complexity is  $\bar{\tau} T = \mathcal{O}(n^{-1}\epsilon^{-2})$ , and the number of communication rounds is  $T = \mathcal{O}(\epsilon^{-1})$ . For partial client participation, the sample complexity is  $\bar{\tau} T = \mathcal{O}(P^{-1}\epsilon^{-2})$ , and the number of communication rounds is  $T = \mathcal{O}(P^{-1}\epsilon^{-2})$ .
In Theorem 2, we show that even under the system-level heterogeneity, ShroFBO can converge to the original objective function with the same convergence rate as SimFBO. This justifies the design principle of robust server-side updates.
Bilevel optimization. Bilevel optimization, first introduced by [3], has been studied for decades. A class of constraint-based bilevel methods was then proposed [20, 16, 59, 61], whose idea is to replace the lower-level problem by the optimality conditions. Gradient-based bilevel algorithms have attracted considerable attention due to the effectiveness in machine learning. Among them, AID-based approaches [8, 51, 38, 1] leveraged the implicit derivation of the hypergradient, which was then approximated via solving a linear system. ITD-based approaches [45, 12, 11, 56, 17] approximated the hypergradient based on automatic differentiation via the forward or backward mode. A group of stochastic bilevel approaches has been developed and analyzed recently based on Neumann series [5, 26, 1], recursive momentum [72, 23, 19] and variance reduction [72, 7], etc. For the lower-level problem with multiple solutions, several approaches were proposed based on the upper- and lower-level gradient aggregation [55, 43, 34], barrier types of regularization [41, 39], penalty-based formulations [58], primal-dual technique [62], and dynamic system-based methods [42].
Federated (bilevel) learning. Federated Learning was proposed to enable collaborative model training across multiple clients without compromising the confidentiality of individual data [32, 60, 49]. As one of the earliest methods of federated learning [47], FedAvg has inspired an increasing number of approaches to deal with different limitations such as slower convergence, high communication cost and undesired client drift by leveraging the techniques including proximal regularization [37], periodic variance reduction [48, 28], proximal splitting [50], adaptive gradients [53]. Theoretically, the convergence of FedAvg and its variants has been analyzed in various settings with the homogeneous [63, 68, 64, 2] or heterogeneous datasets [37, 66, 48, 29]. [67] analyzed the impact of the
system-level heterogeneity such as heterogeneous local computing on the convergence. [57] further extended the analysis and the methods to the minimax problem setting.
Federated bilevel optimization has not been explored well except for a few attempts recently. For example, [14, 35] proposed momentum-based bilevel algorithms, and analyzed their convergence in the setting with homogeneous datasets. In the setting with non-i.i.d. datasets, [65] and [24] proposed FedNest and FedMBO via AID-based federated hypergradient estimation, and [69] proposed an ITD-based aggregated approach named Agg-ITD. Momentum-based techniques have been also used by [22, 36] to improve the sample complexity. Moreover, there are some studies that focus on other distributed scenarios, including decentralized bilevel optimization [6, 73, 44], asynchronous bilevel optimization over directed network [74], and distributed bilevel network utility maximization [27].
Figure 2: Comparison among our SimFBO, FBO-AggITD [69], FedNest [65] and LFedNest [65]. The left and middle ones plot the training accuracy v.s. # of communication rounds on i.i.d. MNIST datasets with MLP networks, and the right one plots the training accuracy v.s. # of rounds on i.i.d. CIFAR-10 datasets with a 7-layer CNN.
Figure 3: Comparison of different methods: the test accuracy v.s. # of communication rounds.
In this section, we perform two hyper-representation experiments to compare the performance of our proposed SimFBO algorithm with FBO-AggITD [69], FedNest [69], and LFedNest [65], and validate the better performance of ShroFBO in the presence of heterogeneous local computation. We test the performance on MNIST and CIFAR datasets with MLP and CNN backbones. We follow the same experimental setup and problem formulation as in [65, 69]. The details of all experimental specifications can be found in Appendix A.1.
Comparison to existing methods. The comparison results are presented in Figure 2 and Figure 3. It can be seen that across different datasets and backbones, our proposed SimFBO consistently converges much faster than other comparison methods, while achieving a much higher training and test accuracy. We do not plot the curves of FedNest and LFedNest on CIFAR and CNN, because they are hard to converge under various hyperparameter configurations using their source codes.
Performance under heterogeneous local computation. We now test the performance in the setting where a total of 10 clients perform a variable number of local steps. This is to simulate the scenario where clients have heterogeneous computing capabilities and hence can perform an uneven number of local updates. In this experiment, we choose the number  $\tau_{i}$  of the client  $i$ 's local steps from the set  $\{1,\dots,10\}$  uniform at random. As shown in Figure 4, the proposed ShroFBO method performs the best due to the better resilience to such client heterogeneity. We also compare the convergence rate of our proposed SimFBO, FedNest and FBO-AggITD w.r.t. running time. The results are provided in Figure 5. All the settings for different algorithms are the same as in Appendix A.2. It can be seen that the proposed SimFBO still converges fastest with a higher test accuracy in terms of running time.
Figure 4: Comparison among different algorithms in the presence of heterogeneous local computation.
Figure 5: Time Comparison among different algorithms.
In this paper, we propose a simple and communication-efficient federated bilevel algorithm named SimFBO and its variant ShroFBO with better resilience to the system-level heterogeneity. We show that both SimFBO and ShroFBO allow for more practical client sampling without replacement, and achieve better sample and communication complexities. Experiments demonstrate the great promise of the proposed methods. We anticipate that the proposed algorithms and the developed analysis can be applied to other distributed settings such as decentralized or asynchronous bilevel optimization, and the algorithms may be useful in applications such as hyperparameter tuning or fine-tuning in federated learning or AI-aware edge computing.
[1] M. Arbel and J. Mairal. Amortized implicit differentiation for stochastic bilevel optimization. arXiv preprint arXiv:2111.14580, 2021.
[2] D. Basu, D. Data, C. Karakus, and S. Diggavi. Qsparse-local-sgd: Distributed sgd with quantization, sparsification and local computations. Advances in Neural Information Processing Systems, 32, 2019.
[3] J. Bracken and J. T. McGill. Mathematical programs with optimization problems in the constraints. Operations research, 21(1):37-44, 1973.
[4] T. Chen, Y. Sun, and W. Yin. Closing the gap: Tighter analysis of alternating stochastic gradient methods for bilevel problems. Advances in Neural Information Processing Systems, 34:25294-25307, 2021.
[5] T. Chen, Y. Sun, and W. Yin. A single-timescale stochastic bilevel optimization method. arXiv preprint arXiv:2102.04671, 2021.
[6] X. Chen, M. Huang, and S. Ma. Decentralized bilevel optimization. arXiv preprint arXiv:2206.05670, 2022.
[64] S. U. Stich and S. P. Karimireddy. The error-feedback framework: Better rates for sgd with delayed gradients and compressed updates. The Journal of Machine Learning Research, 21(1):9613-9648, 2020.
[65] D. A. Tarzanagh, M. Li, C. Thrampoulidis, and S. Oymak. Fednest: Federated bilevel, minimax, and compositional optimization. In International Conference on Machine Learning, pages 21146-21179. PMLR, 2022.
[66] J. Wang and G. Joshi. Cooperative sgd: A unified framework for the design and analysis of local-update sgd algorithms. The Journal of Machine Learning Research, 22(1):9709-9758, 2021.
[67] J. Wang, Q. Liu, H. Liang, G. Joshi, and H. V. Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems, 33:7611-7623, 2020.
[68] Z. Wang, K. Balasubramanian, S. Ma, and M. Razaviyayn. Zeroth-order algorithms for nonconvex minimax problems with improved complexities. arXiv preprint arXiv:2001.07819, 2020.
[69] P. Xiao and K. Ji. Communication-efficient federated hypergradient computation via aggregated iterative differentiation. arXiv preprint arXiv:2302.04969, 2023.
[70] E. P. Xing, Q. Ho, P. Xie, and D. Wei. Strategies and principles of distributed machine learning on big data. Engineering, 2(2):179-195, 2016.
[71] P. Xing, S. Lu, L. Wu, and H. Yu. Big-fed: Bilevel optimization enhanced graph-aided federated learning. IEEE Transactions on Big Data, 2022.
[72] J. Yang, K. Ji, and Y. Liang. Provably faster algorithms for bilevel optimization. Advances in Neural Information Processing Systems, 34:13670-13682, 2021.
[73] S. Yang, X. Zhang, and M. Wang. Decentralized gossip-based stochastic bilevel optimization over communication networks. arXiv preprint arXiv:2206.10870, 2022.
[74] F. Yousefian. Bilevel distributed optimization in directed networks. In 2021 American Control Conference (ACC), pages 2230-2235. IEEE, 2021.
[75] Y. Zeng, H. Chen, and K. Lee. Improving fairness via federated learning. arXiv preprint arXiv:2110.15545, 2021.
[76] Y. Zhang, G. Zhang, P. Khanduri, M. Hong, S. Chang, and S. Liu. Revisiting and advancing fast adversarial training through the lens of bi-level optimization. In International Conference on Machine Learning, pages 26693–26712. PMLR, 2022.