SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Junyi Li
College Park, MD 20742
junyili.ai@gmail.com
Feihu Huang
Pittsburgh, PA 15261
huangfeihu2018@gmail.com
Heng Huang *
henghuanghh@gmail.com
Bilevel Optimization has witnessed notable progress recently with new emerging efficient algorithms. However, its application in the Federated Learning setting remains relatively underexplored, and the impact of Federated Learning's inherent challenges on the convergence of bilevel algorithms remain obscure. In this work, we investigate Federated Bilevel Optimization problems and propose a communication-efficient algorithm, named FedBiOAcc. The algorithm leverages an efficient estimation of the hyper-gradient in the distributed setting and utilizes the momentum-based variance-reduction acceleration. Remarkably, FedBiOAcc achieves a communication complexity  $O(\epsilon^{-1})$ , a sample complexity  $O(\epsilon^{-1.5})$  and the linear speed up with respect to the number of clients. We also analyze a special case of the Federated Bilevel Optimization problems, where lower level problems are locally managed by clients. We prove that FedBiOAcc-Local, a modified version of FedBiOAcc, converges at the same rate for this type of problems. Finally, we validate the proposed algorithms through two real-world tasks: Federated Data-cleaning and Federated Hyper-representation Learning. Empirical results show superior performance of our algorithms.
Bilevel optimization [51, 48] has increasingly drawn attention due to its wide-ranging applications in numerous machine learning tasks, including hyper-parameter optimization [42], meta-learning [61] and neural architecture search [36]. A bilevel optimization problem involves an upper problem and a lower problem, wherein the upper problem is a function of the minimizer of the lower problem. Recently, great progress has been made to solve this type of problems, particularly through the development of efficient single-loop algorithms that rely on diverse gradient approximation techniques [23]. However, the majority of existing bilevel optimization research concentrates on standard, non-distributed settings, and how to solve the bilevel optimization problems under distributed settings have received much less attention. Federated learning (FL) [40] is a recently promising distributed learning paradigm. In FL, a set of clients jointly solve a machine learning task under the coordination of a central server. To protect user privacy and mitigate communication overhead, clients perform multiple steps of local update before communicating with the server. A variety of algorithms [50, 59, 16, 26, 1] have been proposed to accelerate this training process. However, most of these algorithms primarily address standard single-level optimization problems. In this work, we study the bilevel optimization problems in the Federated Learning setting and investigate the
*This work was partially supported by NSF IIS 1838627, 1837956, 1956002, 2211492, CNS 2213701, CCF 2217003, DBI 2225775.
Table 1: Comparisons of the Federated/Non-federated bilevel optimization algorithms for finding an  $\epsilon$ -stationary point of (1).  $Gc(f, \epsilon)$  and  $Gc(g, \epsilon)$  denote the number of gradient evaluations w.r.t.  $f^{(m)}(x, y)$  and  $g^{(m)}(x, y)$ ;  $JV(g, \epsilon)$  denotes the number of Jacobian-vector products;  $HV(g, \epsilon)$  is the number of Hessian-vector products;  $\kappa = L / \mu$  is the condition number,  $p(\kappa)$  is used when no dependence is provided. Sample complexities are measured by client.
Setting	Algorithm	Communication	Gc(f,ε)	Gc(g,ε)	JV(g,ε)	HV(g,ε)	Heterogeneity
Non-Fed	StocBiO [24]	O(κ5ε-2)	O(κ9ε-2)	O(κ5ε-2)	O(κ6ε-2)
MRBO [56]	O(p(κ)e-1.5)	O(p(κ)e-1.5)	O(p(κ)e-1.5)	O(p(κ)e-1.5)
Federated	CommFedBiO [33]	O(p(κ)ε-2)	O(p(κ)ε-2)	O(p(κ)ε-2)	O(p(κ)ε-2)	O(p(κ)ε-2)	✓
FedNest [49]	O(κ9ε-2)	O(κ5ε-2)	O(κ9ε-2)	O(κ5ε-2)	O(κ9ε-2)	✓
AggITD [53]	O(p(κ)ε-2)	O(p(κ)ε-2)	O(p(κ)ε-2)	O(p(κ)ε-2)	O(p(κ)ε-2)	✓
FedMBO [21]	O(M-1p(κ)ε-2)	O(M-1p(κ)ε-2)	O(M-1p(κ)ε-2)	O(M-1p(κ)ε-2)	O(M-1p(κ)ε-2)	✓
SimFBO [58]	O(p(κ)ε-1)	O(M-1p(κ)ε-2)	O(M-1p(κ)ε-2)	O(M-1p(κ)ε-2)	O(M-1p(κ)ε-2)	✓
Local-BSGVR [12]	O(p(κ)ε-1)	O(M-1p(κ)ε-1.5)	O(M-1p(κ)ε-1.5)	O(M-1p(κ)ε-1.5)	O(M-1p(κ)ε-1.5)	✗
FedBiOAcc (Ours)	O(κ19/3ε-1)	O(M-1κ8ε-1.5)	O(M-1κ8ε-1.5)	O(M-1κ8ε-1.5)	O(M-1κ8ε-1.5)	✓
following research question: Is it possible to develop communication-efficient federated algorithms tailored for bilevel optimization problems that also ensure a rapid convergence rate?
More specifically, a general Federated Bilevel Optimization problem has the following form:
$$
\min  _ {x \in \mathbb {R} ^ {p}} h (x) := \frac {1}{M} \sum_ {m = 1} ^ {M} f ^ {(m)} (x, y _ {x}), \text {s . t .} y _ {x} = \underset {y \in \mathbb {R} ^ {d}} {\arg \min } \frac {1}{M} \sum_ {m = 1} ^ {M} g ^ {(m)} (x, y) \tag {1}
$$
A federated bilevel optimization problem consists of an upper and a lower level problem, the upper problem  $f(x,y) \coloneqq \frac{1}{M}\sum_{m = 1}^{M}f^{(m)}(x,y)$  relies on the solution  $y_{x}$  of the lower problem, and  $g(x,y) \coloneqq \frac{1}{M}\sum_{m = 1}^{M}g^{(m)}(x,y)$ . Meanwhile, both the upper and the lower level problems are federated: In Eq.(1), we have  $M$  clients, and each client has a local upper problem  $f^{(m)}(x,y)$  and a lower level problem  $g^{(m)}(x,y)$ . Compared to single-level federated optimization problems, the estimation of the hyper-gradient in federated bilevel optimization problems is much more challenging. In Eq.(1), the hyper-gradient is not linear w.r.t the local hyper-gradients of clients, whereas the gradient of a single-level Federated Optimization problem is the average of local gradients. Consequently, directly applying the vanilla local-sgd method [40] to federated bilevel problems results in a large bias. In the literature [49, 33, 21, 53], researchers evaluate the hyper-gradient through multiple rounds of client-server communication, however, this approach leads to high communication overhead. In contrast, we view the hyper-gradient estimation as solving a quadratic federated problem and solving it with the local-sgd method. More specifically, we formulate the solution of the federated bilevel optimization as three intertwined federated problems: the upper problem, the lower problem and the quadratic problem for the hyper-gradient estimation. Then we address the three problems using alternating gradient descent steps, furthermore, to manage the noise of the stochastic gradient and obtain the fast convergence rate, we employ a momentum-based variance reduction technique.
Beyond the standard federated bilevel optimization problem as defined in Eq. 1, another variant of Federated Bilevel Optimization problem, which entails locally managed lower-level problems, is also frequently utilized in practical applications. For this type of problem, we can get an unbiased estimate of the global hyper-gradient using local hyper-gradient, thus we can solve it with a local-SGD like algorithm, named FedBiOAcc-Local. However, it is challenging to analyze the convergence of the algorithm. In particular, we need to bound the intertwined client drift error, which is intrinsic to FL and the bilevel-related errors e.g. the lower level solution bias. In fact, we prove that the FedBiOAcc-Local algorithm attains the same fast rate as FedBiO algorithm.
Finally, we highlight the main contributions of our paper as follows:
1. We propose FedBiOAcc to solve Federated Bilevel Optimization problems, the algorithm evaluates the hypergradient of federated bilevel optimization problems efficiently and achieves optimal convergence rate through momentum-based variance reduction. FedBiOAcc has sample complexity of  $O(\epsilon^{-1.5})$ , communication complexity of  $O(\epsilon^{-1})$  and achieves linear speed-up w.r.t the number of clients.
2. We study Federated Bilevel Optimization problem with local lower level problem for the first time, where we show the convergence of a modified version of FedBiOAcc, named FedBiOAcc-Local for this type of problems.
3. We validate the efficacy of the proposed FedBiOAcc algorithm through two real-world tasks: Federated Data Cleaning and Federated Hyper-representation Learning.
Notations  $\nabla$  denotes full gradient,  $\nabla_{x}$  denotes partial derivative for variable  $x$ , higher order derivatives follow similar rules.  $[K]$  represents the sequence of integers from 1 to  $K$ ,  $\bar{x}$  represents average of the sequence of variables  $\{x^{(m)}\}_{m=1}^{M}$ .  $\bar{t}_s$  represents the global communication timestamp  $s$ .
Bilevel optimization dates back to at least the 1960s when [51] proposed a regularization method, and then followed by many research works [10, 48, 55, 45], while in machine learning community, similar ideas in the name of implicit differentiation were also used in Hyper-parameter Optimization [30, 3, 2, 8]. Early algorithms for Bilevel Optimization solved the accurate solution of the lower problem for each upper variable. Recently, researchers developed algorithms that solve the lower problem with a fixed number of steps, and use the 'back-propagation through time' technique to compute the hyper-gradient [9, 39, 11, 43, 47]. Very Recently, it witnessed a surge of interest in using implicit differentiation to derive single loop algorithms [14, 17, 23, 28, 4, 56, 19, 32, 7, 20, 18]. In particular, [32, 7] proposes a way to iteratively evaluate the hyper-gradients to save computation. In this work, we view the hyper-gradient estimation of Federated Bilevel Optimization as solving a quadratic federated optimization problem and use a similar iterative evaluation rule as [32, 7] in local update.
The bilevel optimization problem is also considered in the more general settings. For example, bilevel optimization with multiple lower tasks is considered in [15], furthermore, [5, 57, 38, 13] studies the bilevel optimization problem in the decentralized setting, [25] studies the bilevel optimization problem in the asynchronous setting. In contrast, we study bilevel optimization problems under Federated Learning [40] setting. Federated learning is a promising privacy-preserving learning paradigm for distributed data. Compared to traditional data-center distributed learning, Federated Learning poses new challenges including data heterogeneity, privacy concerns, high communication cost, and unfairness. To deal with these challenges, various methods [26, 35, 46, 60, 41, 34] are proposed. However, bilevel optimization problems are less investigated in the federated learning setting. [54] considered the distributed bilevel formulation, but it needs to communicate the Hessian matrix for every iteration, which is computationally infeasible. More recently, FedNest [49] has been proposed to tackle the general federated nest problems, including federated bilevel problems. However, this method evaluates the full hyper-gradient at every iteration; this leads to high communication overhead; furthermore, FedNest also uses SVRG to accelerate the training. Similar works that evaluate the hyper-gradient with multiple rounds of client-server communication are [33, 21, 53, 58]. Finally, there is a concurrent work [12] that investigates the possibility of local gradients on Federated Bilevel Optimization, however, it only considers the homogeneous case, this setting is quite constrained and much simpler than the more general heterogeneous case we considered. Furthermore, [12] only considers the case where both the upper and the lower problem are federated, and omit the equally important case where the lower level problem is not federated.
3 Federated Bilevel Optimization
3.1 Some Mild Assumptions
Note that the formulation of Eq.(1) is very general, and we consider the stochastic heterogeneous case in this work. More specifically, we assume:
$$
f ^ {(m)} (x, y) := \mathbb {E} _ {\xi \sim \mathcal {D} _ {f} ^ {(m)}} [ f ^ {(m)} (x, y, \xi) ], g ^ {(m)} (x, y) := \mathbb {E} _ {\xi \sim \mathcal {D} _ {g} ^ {(m)}} [ g ^ {(m)} (x, y; \xi) ]
$$
where  $\mathcal{D}_f^{(m)}$  and  $\mathcal{D}_g^{(m)}$  are some probability distributions. Furthermore, we assume the local objectives could be potentially different:  $f^{(m)}(x,y)\neq f^{(k)}(x,y)$  or  $g^{(m)}(x,y)\neq g^{(k)}(x,y)$  for  $m\neq k,m,k\in [M]$ . Furthermore, we assume the following assumptions in our subsequent discussion:
Assumption 3.1. Function  $f^{(m)}(x,y)$  is possibly non-convex and  $g^{(m)}(x,y)$  is  $\mu$ -strongly convex w.r.t  $y$  for any given  $x$ .
Assumption 3.2. Function  $f^{(m)}(x,y)$  is  $L$ -smooth and has  $C_f$ -bounded gradient;
Assumption 3.3. Function  $g^{(m)}(x,y)$  is  $L$ -smooth, and  $\nabla_{xy}g^{(m)}(x,y)$  and  $\nabla_{y^2}g^{(m)}(x,y)$  are Lipschitz continuous with constants  $L_{xy}$  and  $L_{y^2}$  respectively;
Algorithm 1 Accelerated Federated Bilevel Optimization (FedBiOAcc)
1: Input: Constants  $c_{\omega}, c_{\nu}, c_{u}, \gamma, \eta, \tau, r$ ; learning rate schedule  $\{\alpha_t\}$ ,  $t \in [T]$ , initial state  $(x_1, y_1, u_1)$ ;
2: Initialization: Set  $y_{1}^{(m)} = y_{1}$ ,  $x_{1}^{(m)} = x_{1}$ ,  $u_{1}^{(m)} = u_{1}$ ,  $\omega_{1}^{(m)} = \nabla_{y}g^{(m)}(x_{1},y_{1},\mathcal{B}_{y})$ ,  $\nu_{1}^{(m)} = \nabla_{x}f^{(m)}(x_{1},y_{1};\mathcal{B}_{f,1}) - \nabla_{xy}g^{(m)}(x_{1},y_{1};\mathcal{B}_{g,1})u_{1}$  and  $q_{1} = \nabla_{y^{2}}g^{(m)}(x_{1}^{(m)},y_{1}^{(m)};\mathcal{B}_{g,2})u_{1} - \nabla_{y}f^{(m)}(x_{1}^{(m)},y_{1}^{(m)};\mathcal{B}_{f,2})$  for  $m \in [M]$
3: for  $t = 1$  to  $T$  do  
4:  $\hat{y}_{t + 1}^{(m)} = y_t^{(m)} - \gamma \alpha_t\omega_t^{(m)},\hat{x}_{t + 1}^{(m)} = x_t^{(m)} - \eta \alpha_t\nu_t^{(m)},\hat{u}_{t + 1}^{(m)} = \mathcal{P}_r(u_t^{(m)} - \tau \alpha_tq_t^{(m)})$
5: if  $t$  mod  $\mathbf{I} = 0$  then  
6:  $y_{t + 1}^{(m)} = \frac{1}{M}\sum_{j = 1}^{M}\hat{y}_{t + 1}^{(j)}; x_{t + 1}^{(m)} = \frac{1}{M}\sum_{j = 1}^{M}\hat{x}_{t + 1}^{(j)}, u_{t + 1}^{(m)} = \frac{1}{M}\sum_{j = 1}^{M}\hat{u}_{t + 1}^{(j)}$
7: else
8:  $y_{t + 1}^{(m)} = \hat{y}_{t + 1}^{(m)},x_{t + 1}^{(m)} = \hat{x}_{t + 1}^{(m)},u_{t + 1}^{(m)} = \hat{u}_{t + 1}^{(m)}$
9: end if
10: Get  $\hat{\omega}_{t + 1}^{(m)},\hat{\nu}_{t + 1}^{(m)}$  and  $\hat{q}_{t + 1}^{(m)}$  following Eq. (7)
1: if  $t \mod I = 0$  then
2:  $\omega_{t + 1}^{(m)} = \frac{1}{M}\sum_{j = 1}^{M}\hat{\omega}_{t + 1}^{(j)},\nu_{t + 1}^{(m)} = \frac{1}{M}\sum_{j = 1}^{M}\hat{\nu}_{t + 1}^{(j)},q_{t + 1}^{(m)} = \frac{1}{M}\sum_{j = 1}^{M}\hat{q}_{t + 1}^{(j)},$
3: else
4:  $\omega_{t + 1}^{(m)} = \hat{\omega}_{t + 1}^{(m)},\nu_{t + 1}^{(m)} = \hat{\nu}_{t + 1}^{(m)},q_{t + 1}^{(m)} = \hat{q}_{t + 1}^{(m)}$
5: end if
6: end for
Assumption 3.4. We have unbiased stochastic first-order and second-order gradient oracle with bounded variance.
Assumption 3.5. For any  $m, j \in [M]$  and  $z = (x,y)$ , we have:  $\| \nabla f^{(m)}(z) - \nabla f^{(j)}(z) \| \leq \zeta_f$ ,  $\| \nabla g^{(m)}(z) - \nabla g^{(j)}(z) \| \leq \zeta_g$ ,  $\| \nabla_{xy} g^{(m)}(z) - \nabla_{xy} g^{(j)}(z) \| \leq \zeta_{g,xy}$ ,  $\| \nabla_{y^2} g^{(m)}(z) - \nabla_{y^2} g^{(j)}(z) \| \leq \zeta_{g,yy}$ , where  $\zeta_f, \zeta_g, \zeta_{g,xy}, \zeta_{g,yy}$ , are constants.
As stated in The assumption 3.1, we study the non-convex-strongly-convex bilevel optimization problems, this class of problems is widely studied in the non-distributed bilevel literature [22, 14]. Furthermore, Assumption 3.2 and Assumption 3.3 are also standard assumptions made in the non-distributed bilevel literature. Assumption 3.4 is widely used in the study of stochastic optimization problems. For Assumption 3.5, gradient difference is widely used in single level Federated Learning literature as a measure of client heterogeneity [28, 52]. Please refer to the full version of Assumptions in Appendix.
3.2 The FedBiOAcc Algorithm
A major difficulty in solving a Federated Bilevel Optimization problem Eq. (1) is evaluating the hyper-gradient  $\nabla h(x)$ . For the function class (non-convex-strongly-convex) we consider, the explicit form of hypergradient  $h(x)$  exists as  $\nabla h(x) = \Phi(x, y_x)$ , where  $\Phi(x, y)$  is denoted as:
$$
\Phi (x, y) = \nabla_ {x} f (x, y) - \nabla_ {x y} g (x, y) \times \left[ \nabla_ {y ^ {2}} g (x, y) \right] ^ {- 1} \nabla_ {y} f (x, y), \tag {2}
$$
Based on Assumption 3.1~3.3, we can verify  $\Phi (x,y_x)$  is the hyper-gradient [14]. But since the clients only have access to their local data, for  $\forall m\in [M]$ , the client evaluates:
$$
\Phi^ {(m)} (x, y) = \nabla_ {x} f ^ {(m)} (x, y) - \nabla_ {x y} g ^ {(m)} (x, y) \times [ \nabla_ {y ^ {2}} g ^ {(m)} (x, y) ] ^ {- 1} \nabla_ {y} f ^ {(m)} (x, y), \tag {3}
$$
It is straightforward to verify that  $\Phi^{(m)}(x,y)$  is not an unbiased estimate of the full hyper-gradient, i.e.  $\Phi (x,y_{x})\neq \frac{1}{M}\sum_{m = 1}^{M}\Phi^{(m)}(x,y_{x})$  . To address this difficulty, we can view the Hyper-gradient computation as the process of solving a federated optimization problem.
In fact, Evaluating Eq. (2) is equivalent to the following two steps: first, we solve the quadratic federated optimization problem  $l(u)$ :
$$
\min  _ {u \in \mathbb {R} ^ {d}} l (u) = \frac {1}{M} \sum_ {m = 1} ^ {M} u ^ {T} \left(\nabla_ {y ^ {2}} g ^ {(m)} (x, y)\right) u - \langle \nabla_ {y} f ^ {(m)} (x, y), u \rangle \tag {4}
$$
Suppose that we denote the solution of the above problem as  $u^{*}$ , then we have the following linear operation to get the hypergradient:
$$
\nabla h (x) = \frac {1}{M} \sum_ {m = 1} ^ {M} \left(\nabla_ {x} f ^ {(m)} \left(x, y _ {x}\right) - \nabla_ {x y} g ^ {(m)} \left(x, y _ {x}\right) u ^ {*}\right) \tag {5}
$$
Compared to the formulation Eq. (2), Eq. (4) and Eq. (5) are more suitable for the distributed setting. In fact, both Eq. (4) and Eq. (5) have a linear structure. Eq. (4) is a (single-level) quadratic federated optimization problem, and we could solve Eq. (4) through local-sgd [40], suppose that each client maintains a variable  $u_{t}^{(m)}$ , and performs the following update:
$$
u _ {t + 1} ^ {(m)} = \mathcal {P} _ {r} (u _ {t} ^ {(m)} - \tau_ {t} \nabla l ^ {(m)} (u _ {t} ^ {(m)}; \mathcal {B}))
$$
$$
\nabla l ^ {(m)} (u _ {t} ^ {(m)}; \mathcal {B}) = \nabla_ {y ^ {2}} g ^ {(m)} (x _ {t} ^ {(m)}, y _ {t} ^ {(m)}; \mathcal {B} _ {g, 2})) u _ {t} ^ {(m)} - \nabla_ {y} f ^ {(m)} (x _ {t} ^ {(m)}, y _ {t} ^ {(m)}; \mathcal {B} _ {f, 2})
$$
where  $\nabla l^{(m)}(u_t^{(m)};\mathcal{B})$  is client  $m$ 's the stochastic gradient of Eq. (4), and  $(x_{t}^{(m)},y_{t}^{(m)})$  denotes the upper and lower variable state at the timestamp  $t$ , the  $\mathcal{P}_r(\cdot)$  denotes the projection to a bounded ball of radius-  $r$ . Note that Clients perform multiple local updates of  $u_{t}^{(m)}$  before averaging. As for Eq. (5), each client evaluates  $\nabla h^{(m)}(x)$  locally:  $\nabla h^{(m)}(x) = \nabla_{x}f^{(m)}(x,y_{x}) - \nabla_{xy}g^{(m)}(x,y_{x})u^{*}$  and the server averages  $\nabla h^{(m)}(x)$  to get  $\nabla h(x)$ . In summary, the linear structure of Eq. (4) and Eq. (5) makes it suitable for local updates, therefore, reduce the communication cost.
More specifically, we perform alternative update of upper level variable  $x_{t}^{(m)}$ , the lower level variable  $y_{t}^{(m)}$  and hyper-gradient computation variable  $u_{t}^{(m)}$ . For example, for each client  $m \in [M]$ , we perform the following local updates:
$$
y _ {t + 1} ^ {(m)} = y _ {t} ^ {(m)} - \gamma_ {t} \nabla_ {y} g ^ {(m)} (x _ {t} ^ {(m)}, y _ {t} ^ {(m)}, \mathcal {B} _ {y}), u _ {t + 1} ^ {(m)} = \mathcal {P} _ {r} (u _ {t} ^ {(m)} - \tau_ {t} \nabla l ^ {(m)} (u _ {t} ^ {(m)}; \mathcal {B}))
$$
$$
x _ {t + 1} ^ {(m)} = x _ {t} ^ {(m)} - \eta_ {t} \big (\nabla_ {x} f ^ {(m)} (x _ {t} ^ {(m)}, y _ {t} ^ {(m)}; \mathcal {B} _ {f, 1}) - \nabla_ {x y} g ^ {(m)} (x _ {t} ^ {(m)}, y _ {t} ^ {(m)}; \mathcal {B} _ {g, 1}) u _ {t} ^ {(m)} \big) \qquad (6)
$$
Every  $I$  steps, the server averages clients' local states, this resembles the local-sgd method for single level federated optimization problems. Note that in the update of the upper variable  $x_{t}^{(m)}$ , we use  $u_{t}^{(m)}$  as an estimation of  $u^{*}$  in Eq. (5). An algorithm follows Eq. (6) is shown in Algorithm 2 of Appendix and we refer to it as FedBiO.
Comparison with FedNest. The update rule of Eq. 6 is very different from that of FedNest [49] and its follow-ups [21, 53]. In FedNest, a sub-routine named FedIHGP is used to evaluate Eq. (2) at every global epoch. This involves multiple rounds of client-server communication and leads to higher communication overhead. In contrast, Eq. (6) formulates the hyper-gradient estimation as an quadratic federated optimization problem, and then solves three intertwined federated problems through alternative updates of  $x$ ,  $y$  and  $u$ .
Note that Eq. 6 updates the related variables through vanilla gradient descent steps. In the nonfederated setting, gradient-based methods such as stocBiO [23] requires large-batch size  $(O(\epsilon^{-1}))$  to reach an  $\epsilon$ -stationary point, and we also analyze Algorithm 2 in Appendix to show the same dependence. To control the noise and remove the dependence over large batch size, we apply the momentum-based variance-reduction technique STORM [6]. In fact, Eq. (6) solves three intertwined optimization problems: the bilevel problem  $h(x)$ , the lower level problem  $g(x,y)$  and the hypergradient computation problem Eq (4). So we control the noise in the process of solving each of the three problems. More specifically, we have  $\omega_{t}^{(m)}, \nu_{t}^{(m)}$  and  $q_{t}^{(m)}$  to be the momentum estimator for  $x_{t}^{(m)}, y_{t}^{(m)}$  and  $u_{t}^{(m)}$  respectively, and we update them following the rule of STORM [6]:
$$
\begin{array}{l} \hat {\omega} _ {t + 1} ^ {(m)} = \nabla_ {y} g ^ {(m)} (x _ {t + 1} ^ {(m)}, y _ {t + 1} ^ {(m)}, \mathcal {B} _ {y}) + (1 - c _ {\omega} \alpha_ {t} ^ {2}) (\omega_ {t} ^ {(m)} - \nabla_ {y} g ^ {(m)} (x _ {t} ^ {(m)}, y _ {t} ^ {(m)}, \mathcal {B} _ {y})) \\ \hat {\nu} _ {t + 1} ^ {(m)} = \big (\nabla_ {x} f ^ {(m)} (x _ {t + 1} ^ {(m)}, y _ {t + 1} ^ {(m)}; \mathcal {B} _ {f, 1}) - \nabla_ {x y} g ^ {(m)} (x _ {t + 1} ^ {(m)}, y _ {t + 1} ^ {(m)}; \mathcal {B} _ {g, 1}) u _ {t + 1} ^ {(m)} \big) \\ + \left(1 - c _ {\nu} \alpha_ {t} ^ {2}\right) \left(\nu_ {t} ^ {(m)} - \left(\nabla_ {x} f ^ {(m)} \left(x _ {t} ^ {(m)}, y _ {t} ^ {(m)}; \mathcal {B} _ {f, 1}\right) - \nabla_ {x y} g ^ {(m)} \left(x _ {t} ^ {(m)}, y _ {t} ^ {(m)}; \mathcal {B} _ {g, 1}\right) u _ {t} ^ {(m)}\right)\right) \\ \hat {q} _ {t + 1} ^ {(m)} = \left(\nabla_ {y ^ {2}} g ^ {(m)} \left(x _ {t + 1} ^ {(m)}, y _ {t + 1} ^ {(m)}; \mathcal {B} _ {g, 2}\right) u _ {t + 1} ^ {(m)} - \nabla_ {y} f ^ {(m)} \left(x _ {t + 1} ^ {(m)}, y _ {t + 1} ^ {(m)}; \mathcal {B} _ {f, 2}\right)\right) \\ + \left(1 - c _ {u} \alpha_ {t} ^ {2}\right) \left(q _ {t} ^ {(m)} - \left(\nabla_ {y ^ {2}} g ^ {(m)} \left(x _ {t} ^ {(m)}, y _ {t} ^ {(m)} ; \mathcal {B} _ {g, 2}\right) u _ {t} ^ {(m)} - \nabla_ {y} f ^ {(m)} \left(x _ {t} ^ {(m)}, y _ {t} ^ {(m)} ; \mathcal {B} _ {f, 2}\right)\right)\right) \tag {7} \\ \end{array}
$$
where  $c_{\omega}, c_{\nu}$  and  $c_{u}$  are constants,  $\alpha_{t}$  is the learning rate. Then we update the  $x_{t}^{(m)}, y_{t}^{(m)}$  and  $u_{t}^{(m)}$  as follows:
$$
\hat {y} _ {t + 1} ^ {(m)} = y _ {t} ^ {(m)} - \gamma \alpha_ {t} \omega_ {t} ^ {(m)}, \hat {x} _ {t + 1} ^ {(m)} = x _ {t} ^ {(m)} - \eta \alpha_ {t} \nu_ {t} ^ {(m)}, \hat {u} _ {t + 1} ^ {(m)} = \mathcal {P} _ {r} (u _ {t} ^ {(m)} - \tau \alpha_ {t} q _ {t} ^ {(m)}) \tag {8}
$$
where  $\gamma$ ,  $\eta$ ,  $\tau$  are constants and  $\alpha_{t}$  is the learning rate. The FedBiOAcc algorithm following Eq. (8) is summarized in Algorithm 1. As shown in line 6 and 12 of Algorithm 1, Every  $I$  iterations, we average both variables and the momentum.
3.3 Convergence Analysis
In this section, we study the convergence property for the FedBiOAcc algorithm. For any  $t \in [T]$ , we define the following virtual sequence:
$$
\bar {x} _ {t} = \frac {1}{M} \sum_ {m = 1} ^ {M} x _ {t} ^ {(m)}, \bar {y} _ {t} = \frac {1}{M} \sum_ {m = 1} ^ {M} y _ {t} ^ {(m)}, \bar {u} _ {t} = \frac {1}{M} \sum_ {m = 1} ^ {M} u _ {t} ^ {(m)}
$$