SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Federated learning (FL) is a promising privacy-preserving machine learning paradigm over distributed data. In this paradigm, each client trains the parameter of a model locally and the server aggregates the parameter from clients periodically. Therefore, we perform the learning and communication over the same set of parameters. However, we find that learning and communication have fundamentally divergent requirements for parameter selection, akin to two opposite teams in a tug-of-war game. To mitigate this discrepancy, we introduce FedSep, a novel two-layer federated learning framework. FedSep consists of separated communication and learning layers for each client and the two layers are connected through decode/encode operations. In particular, the decoding operation is formulated as a minimization problem. We view FedSep as a federated bilevel optimization problem and propose an efficient algorithm to solve it. Theoretically, we demonstrate that its convergence matches that of the standard FL algorithms. The separation of communication and learning in FedSep offers innovative solutions to various challenging problems in FL, such as Communication-Efficient FL and Heterogeneous-Model FL. Empirical validation shows the superior performance of FedSep over various baselines in these tasks.
In Federated Learning (FL) [37], a set of clients jointly solve a machine learning task under the coordination of a central server. The process of FL involves two categories of operations: Learning and Communication. For the learning operation, each client optimizes their local objectives, and for the communication operation, clients exchange local parameters to facilitate knowledge sharing. In the existing FL pipeline, both Learning and Communication operations hinge on the same set of parameters. However, these two operations have fundamentally divergent requirements, akin to two teams engaged in a tug-of-war, each pulling in opposite directions. In fact, on the communication side, it is imperative that the parameters of all clients reside in a uniform space. Moreover, to mitigate the high communication costs, a major bottleneck in FL, it is beneficial to maintain these parameters within a low-dimensional space. In contrast, on the learning side, given the heterogeneity of devices, including variations in hardware and data distribution, it is advantageous to allow the parameter space to differ across clients. Furthermore, to accommodate the implementation of state-of-the-art large-scale machine learning models (such as Transformers [47]), a high-dimensional parameter space is desirable. In summary, a huge discrepancy of requirements to parameter selection exists between the communication and the learning in FL.
To mitigate this discrepancy, we opt to 'break the rope' in this tug-of-war scenario. More precisely, we propose a two-layer structure on clients: a communication layer and a learning layer. As shown in Figure 1, we denote this framework as FedSep: The communication layer connects with the server and uploads/downloads parameter, while the learning layer performs local objective optimization. Then the communication layer and the learning layer are connected through encode/decode operations. The decode operation maps the parameter of the communication layer to the parameter of the learning layer and the encode operation performs the mapping in the opposite direction. More specifically, we formulate the decode operation as solving a minimization problem and does not assume an explicit relation between the parameter of communication layer and the parameter of the learning layer. In summary, we aim to use two distinct sets of parameters in FedSep to resolve the 'competition' between the communication and learning operations, meanwhile, the two set of parameters are close connected through decode/encode operations. Mathematically, FedSep can be formulated as solving a federated bilevel problem. The upper level problem corresponds to the learning problems on clients, while the lower level problem is the minimization problem in the decode operation. We propose an efficient algorithm to solve this bilevel problem, which is composed of three consecutive steps: Decode the communication parameter, Learning on local problems and Encode the learning parameter to the communication parameter. The convergence of the algorithm is guaranteed.
To demonstrate the flexibility of FedSep in incorporating the contradictory requirements of the communication and learning, we study two real-world FL tasks. In the first task, we study the communication-efficient FL task. In this task, we set the dimension of the communication parameter to be much smaller than the learning parameter, furthermore, we decode the communication parameter through a LASSO problem. In the second task, we study the model-heterogeneous FL task. In this task, the learning parameters on different clients have different dimension. We set the communication parameter to be the parameter of the server model, and the learning parameter be a subset of the communication parameter (server model). In particular, the decode operation is to select the subset adaptively based on the client's local data. We empirically verify the superior performance of our methods compared to other baselines. Finally, we summarize the main contributions of our work as follows:
Figure 1: The structure of FedSep. FedSep follows the one-server-multiple-clients structure, in particular, each client has two layers: the communication layer and the learning layer. The two layers are connected through the decode and encode operations.
1. We propose a novel two-layer Federated Learning framework, i.e. FedSep, where the communication layer and the learning layer are separated on clients;
2. We let the decode operation from the communication layer to the learning layer be solving a minimization problem, and therefore, the FedSep framework has a bilevel optimization formulation. We propose an efficient algorithm to solve the bilevel problem and show that its convergence rate matches that of the standard Federated Learning algorithms;
3. We apply the FedSep framework to solve two important FL tasks: Communication-Efficient FL and Model-Heterogeneous FL. Numerical experiments show the superior performance of the methods based on FedSep.
Notations.  $\nabla$  denotes the full gradient,  $\nabla_{x}$  is the partial derivative for variable  $x$ , and higher-order derivatives follow similar rules.  $||\cdot ||_2$  is  $\ell_2$ -norm for vectors and the spectral norm for matrices.  $AB$  denotes the multiplication of the matrix between the matrix  $A$  and  $B$ .  $[K]$  represents the sequence of integers from 1 to  $K$ .
Federated Learning. FL is a novel paradigm for performing machine learning tasks over distributed data, but it also poses challenges such as heterogeneity [24], privacy [38] and communication
bottleneck [48]. In particular, communication cost is one of the major bottlenecks in FL. A widely used approach to reduce communication burden is through compression. Various compressors [48, 35, 45, 25] have been studied in the literature. Next, model-heterogeneous FL is also studied to solve the device heterogeneity issue. There are two main categories of heterogeneous model methods: knowledge-distillation (KD) [34, 22, 9] based and partial-training (PT) [11, 6, 1] based. KD-based methods treat the server model as the teacher and the clients' model as the student. In contrast, PT-based methods select a subset of the server model for each client, and in the aggregation phase, the server simply aggregates the updated sub-models.
Bilevel Optimization. Bilevel problem is a type of two-layer optimization problem [49]. In the machine learning community, Hyper-parameter Optimization problems [28, 7, 4, 12] are one of the mostly studied bilevel problems. Early bilevel optimization algorithms solved the lower problem exactly to update the upper variable. Recently, researchers developed algorithms that solve the lower problem approximately with a fixed number of steps and use the back-propagation technique to compute the approximate hyper-gradient [13, 36, 14, 40, 43]. More recently, single-loop algorithms [16, 18, 23, 30, 26, 8, 51, 31, 10, 20, 21] based on alternative updates of lower and upper level problems are proposed. However, less work studies bilevel problems under the federated setting. [46, 33, 32] studies a type of federated bilevel problems where the lower level problems are federated. In particular, [33] solves the noisy label problem in FL. Note that our FedSep is formulated as a type of federated bilevel problems, however, our problem is different from that in [46, 33]. Firstly, our problem has a unique lower problem for each client, moreover, FedSep impose unique constraints not covered by the existing methods. Please refer to Section B in Appendix for more related works.
3 Separating Communication and Learning in Federated Learning
In this section, we propose FedSep, a novel two-layer Federated Learning Framework. As shown in Figure 1, FedSep has the one-server-multiple-clients structure as in the classical Federated learning framework, we assume  $M$  clients in the training. The difference between the FedSep and classical FL framework lies in the client. Different from the classical FL, FedSep includes two separated layers on clients: a communication layer and a learning layer. The two layers have different responsibilities. As the name shows, the learning layer is responsible of finishing the learning task, e.g. fitting a neural network over a dataset. We assume that the learning layer on client  $m \in [M]$  has parameter  $y^{(m)} \in \mathbb{R}^{d^{(m)}}$ . Note that the choice of  $y^{(m)}$  should be adapted to the specific setting of each client, such as the data distribution, hardware resources etc.. Next, the communication layer performs communication with the server (other clients), which includes the upload and download operations. We assume the communication layer on client  $m$  has parameter  $x \in \mathbb{R}^p$ . Note that the parameter of communication layer has identical formulation for all clients. Finally, the communication layer and the learning layer are connected through the encode and decode operations.
To model the various potential relationship between the communication and learning layer, we abstract the decode operation as solving a minimization problem, and the encode operation be its reverse operation. As a result, we have the following bilevel formulation for the FedSep framework.
$$
\min  _ {x \in \mathbb {R} ^ {p}} h (x) := \frac {1}{M} \sum_ {m = 1} ^ {M} h ^ {(m)} (x) := \frac {1}{M} \sum_ {m = 1} ^ {M} f ^ {(m)} \left(y _ {x} ^ {(m)}\right), y _ {x} ^ {(m)} = \underset {y ^ {(m)} \in \mathbb {R} ^ {d ^ {(m)}}} {\arg \min } g ^ {(m)} \left(x, y ^ {(m)}\right) \tag {1}
$$
In Eq. (1),  $M$  is the number of clients;  $f^{(m)}(y), m \in [M]$  denotes the local problem solved by the client  $m$ ;  $g^{(m)}(x, y)$  is the minimization problem solved in the decode operation by the client  $m$ ;  $x$  is the parameter of the communication layer;  $y^{(m)}$  denotes the parameter of the learning layer on the client  $m$ .
Next, we propose a three-stage algorithm to solve Eq. (1). As shown by Algorithm 1, we perform  $T$  global steps, and at each step  $t \in T$ , we randomly choose a subset of clients  $\mathcal{M}_t$  to perform the training and then the server performs aggregation. The Client training is divided into three stages: Decode stage (lines 7-12), Learning stage (lines 14-17), and Encode Stage (line 19).
Decode Stage. In Eq. (1), the decode operation solves the lower optimization problem  $g^{(m)}(x,y)$  to obtain  $y_{x}^{(m)}$ . We denote the decode operator on the  $m_{th}$  client by  $Dec^{(m)}\{\cdot\}$ , thus we have:  $y_{x}^{(m)} = Dec^{(m)}\{x\}$ . In Algorithm 1, we solve  $g^{(m)}(x,y)$  approximately with  $I_{dec}$  steps of stochastic gradient
Algorithm 1 Separating Communication and Learning in FL (FedSep)
1: Input: Initial states  $x_{1}$ ; learning rates  $\gamma$ ,  $\eta$ ; mini-batchsize  $b_{x}$ ,  $b_{y}$   
2: for  $t = 1$  to  $T$  do  
3: Randomly sample a subset  $\mathcal{M}_t$  of clients;  
4: for  $m \in \mathcal{M}_t$  in parallel do  
5: // Decode stage, estimate  $y_{x}^{(m)} = Dec^{(m)}\{x\}$ ;  
6: Receive the global state  $x_{t}$  from the server and set  $x^{(m)} = x_{t}$ , and initialize  $y_0^{(m)}$ ;  
7: for  $i = 1$  to  $I_{dec}$  do  
8: Randomly sample a minibatch of  $b_{y}$  samples  $\mathcal{B}_{y}$ ;  
9:  $y_{i}^{(m)} = y_{i - 1}^{(m)} - \gamma \nabla_{y}g^{(m)}(x^{(m)},y_{i - 1}^{(m)},\mathcal{B}_{y})$   
10: end for  
11: Set  $\hat{y}_0^{(m)} = y_{I_{dec}}^{(m)}$  as the estimation of  $Dec^{(m)}(x^{(m)})$ ;  
12: // Learning stage, optimize  $f^{(m)}(y)$ ;  
13: for  $i = 1$  to  $I$  do  
14: Randomly sample a minibatch of  $b_{y}$  samples  $\mathcal{B}_{\hat{y}}$ ;  
15:  $\hat{y}_i^{(m)} = \hat{y}_{i - 1}^{(m)} - \eta \nabla f^{(m)}(\hat{y}_{i - 1}^{(m)},\mathcal{B}_{\hat{y}})$ ;  
16: end for  
17: // Encode stage, encode the update of the learning layer back to the communication layer;  
18: Set  $\Delta \hat{y}^{(m)} = \hat{y}_I^{(m)} - \hat{y}_0^{(m)}$  and compute  $\Delta \hat{x}_t^{(m)} = \widetilde{Enc}^{(m)}\{\Delta \hat{y}^{(m)}\}$ , where  $\widetilde{Enc}^{(m)}\{\cdot\}$  is defined in Eq. (3) and we choose  $|\mathcal{B}_x| = b_x$ .  
19: end for  
20:  $x_{t + 1} = x_t - \frac{1}{|\mathcal{M}_t|}\sum_{m\in \mathcal{M}_t}\eta_g\Delta \hat{x}_t^{(m)}$   
21: end for
descent (Line 8-11) and use the output as an approximation of  $y_{x}^{(m)}$ , i.e.,  $y_{I_{dec}}^{(m)} = \widetilde{Dec}^{(m)}(x^{(m)}) \approx y_{x}^{(m)}$ , where  $\widetilde{Dec}^{(m)}(\cdot)$  represents the approximation of the exact decoder  $Dec^{(m)}\{\cdot\}$ .
Learning Stage. Next, we perform  $I$  steps of stochastic gradient descent to solve the local learning problem  $f^{(m)}(y)$  (Line 14 - 17). This stage is similar to the local gradients in classical FL framework. Note we get  $\hat{y}_I^{(m)}$  as the updated learning parameter  $y^{(m)}$ , more formally, we use  $\Delta \hat{y}^{(m)} = \hat{y}_I^{(m)} - \hat{y}_0^{(m)} = \hat{y}_I^{(m)} - y_{I_{dec}}^{(m)}$  to represent the update in the learning stage.
Encode Stage. After the learning stage, we get the update of the local learning problem  $\Delta \hat{y}^{(m)}$ . Suppose we denote the encode operator  $Enc^{(m)}\{\cdot\}$ , then we get the update of the communication parameter as  $\Delta \hat{x}^{(m)} = Enc^{(m)}(\Delta \hat{y}^{(m)})$  (Line 19). In particular, we choose the following encode operator:
$$
E n c ^ {(m)} \{\cdot \} := - \nabla_ {x} y _ {x} ^ {(m)} = \nabla_ {x y} g ^ {(m)} (x, y _ {x} ^ {(m)}) \left(\nabla_ {y y} g ^ {(m)} (x, y _ {x} ^ {(m)})\right) ^ {- 1} \tag {2}
$$
where  $y_{x}^{(m)} = \text{Dec}^{(m)}(x)$  is the output of the exact decode operation and the second equality can be derived following the implicit function theorem under mild assumptions [16]. To reduce the computation complexity, we use Neumann series to approximate the inverse operation in Eq. (2). More specifically, we have the following approximation:  $\left(\nabla_{yy}g^{(m)}(x,y)\right)^{-1}\approx \tau \sum_{q = 0}^{Q}(I - \tau \nabla_{yy}g^{(m)}(x,y))^{q}$ , where  $Q$  and  $\tau$  are some constants. Since in the decode stage, we use  $y_{x}^{(m)}\approx y_{I_{dec}}^{(m)}$ , we get a stochastic approximation of Eq. (2) as:
$$
\widetilde {E n c} ^ {(m)} \{\cdot \} := \sum_ {q = 0} ^ {Q} \tau \nabla_ {x y} g ^ {(m)} (x, y _ {I _ {d e c}} ^ {(m)}; \mathcal {B} _ {x}) (I - \tau \nabla_ {y y} g ^ {(m)} (x, y _ {I _ {d e c}} ^ {(m)}; \mathcal {B} _ {x})) ^ {q} \tag {3}
$$
Remark 1. We use Eq. (2) as the encode operator due to the following fact about the hyper-gradient:
$$
\nabla h ^ {(m)} (x) = \nabla_ {x} y _ {x} ^ {(m)} (\nabla f ^ {(m)} (y _ {x} ^ {(m)}))
$$
Note that suppose  $\Delta \hat{y}^{(m)} = -\nabla f^{(m)}(y_x^{(m)})$ , which means the learning layer performs one step of gradient descent with learning rate 1 in the learning stage of Algorithm 1, then we have:  $\nabla h^{(m)}(x) =$
$\mathit{Enc}\{\Delta \hat{y}^{(m)}\} = \Delta x^{(m)}$ . Therefore, we update the communication parameter  $x$  such that it optimizes the overall problem  $\nabla h^{(m)}(x)$ .
Server Aggregation. Finally, the server aggregates the local updates of communication parameter  $\Delta \hat{x}^{(m)}, m \in \mathcal{M}_t$  to get the new communication parameter as shown in Line 21 of Algorithm 1.
Difference with existing bilevel algorithms. Bilevel optimization problems are widely studied in the literature [16, 23]. However, FedSep cannot directly apply existing algorithms. In a standard bilevel algorithm, each step of update to the upper level variable requires (approximately) solving of the corresponding lower level problem. However, in each epoch of FedSep, we only solve the lower level problem once and perform multiple steps of update to the upper level variable.
3.1 Convergence Analysis
In this section, we provide the convergence analysis of Algorithm 1. Before we state the convergence result, we first make the following assumptions.
Assumption 3.1. Function  $f^{(m)}(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)}(y)$  is  $L$ -smooth and has  $C_f$ -bounded gradient.
Assumption 3.3. Function  $g^{(m)}(x,y)$  is  $L$ -smooth;  $\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.
Note that Assumption 3.1-Assumption 3.3 are standard assumptions used to analyze bilevel problems [23, 16]. Assumption 3.4 is standard in analyzing stochastic optimization problems. For a full version of Assumption 3.4, please refer to Assumption C.1 in the Appendix.
In the learning stage of Algorithm 1, we perform  $I$  steps of updates to optimize the learning problem, and the update has the form of:  $\Delta \hat{y}^{(m)} = \hat{y}_I^{(m)} - \hat{y}_0^{(m)} = \sum_{i=0}^{I-1} -\eta \nabla f^{(m)}(\hat{y}_i^{(m)};\mathcal{B}_{\hat{y}})$ , therefore we have:  $\Delta \hat{x}_t^{(m)} = \widetilde{Enc}^{(m)}\{\Delta \hat{y}^{(m)}\} = \sum_{i=0}^{I-1}\eta \widetilde{Enc}^{(m)}\{-\nabla f^{(m)}(\hat{y}_i^{(m)};\mathcal{B}_{\hat{y}})\}$ . We denote  $\Delta \hat{x}_{t,i}^{(m)} = \widetilde{Enc}^{(m)}\{\nabla f^{(m)}(\hat{y}_i^{(m)};\mathcal{B}_{\hat{y}})\}$ . Then the main effort of the proof is bounding the estimation error of  $\Delta \hat{x}_{t,i}^{(m)}$  to the hyper-gradient  $\nabla h^{(m)}(\boldsymbol{x})$ .
More specifically, we first show that  $\Delta \hat{x}_{t,i}^{(m)}$  estimates a term  $\mu_i^{(m)}$  with bounded variance and bias as stated in the following proposition:
Proposition 3.5. Suppose Assumptions 3.1-3.3 and 3.4 hold and  $\tau < \frac{1}{L}$ , we have  $\| \mathbb{E}_{\xi}[\Delta \hat{x}_{t,i}^{(m)}] - \mu_i^{(m)}\| \leq G_1$ , where  $\mu_i^{(m)} = -\nabla_{xy}g^{(m)}(x,y_{I_{dec}}^{(m)})\bigl (\nabla_{yy}g^{(m)}(x,y_{I_{dec}}^{(m)})\bigr)^{-1}\nabla f^{(m)}(\hat{y}_i^{(m)})$ , and  $\mathbb{E}\| \Delta \hat{x}_{t,i}^{(m)} - \mathbb{E}_{\xi}[\Delta \hat{x}_{t,i}^{(m)}]\| ^2\leq G_2^2$ .  $G_{1}$  and  $G_{2}$  are some constants related to  $Q$  in Eq. (3) and mini-batch size.
Please refer to Proposition C.3 for the specific form of the constants in Proposition 3.5. Next, we have that  $\mu_i^{(m)}\approx \nabla h^{(m)}(x)$  as shown in the following proposition:
Proposition 3.6. Suppose Assumptions 3.2 and 3.3 hold, the following statements hold:
a)  $\| \mu_i^{(m)} - \nabla h^{(m)}(x)\| ^2\leq 2\hat{L}^2\| y_{I_{dec}}^{(m)} - y_x^{(m)}\| ^2 +2\kappa^2 L^2\| \hat{y}_i^{(m)} - y_x^{(m)}\| ^2$  where  $\hat{L} = O(\kappa^2)$
b)  $h^{(m)}(x)$  is Lipschitz continuous in  $x$  with constant  $\bar{L}$  i.e., for any given  $x_1, x_2 \in \mathbb{R}^p$ , we have  $\| \nabla h^{(m)}(x_2) - \nabla h^{(m)}(x_1) \|^2 \leq \bar{L}^2 \| x_2 - x_1 \|$ , where  $\bar{L} = O(\kappa^3)$ .
where we denote the condition number as  $\kappa = L / \mu$ .
Please refer to Proposition C.2 for the proof. Proposition 3.6 shows that  $\mu_i^{(m)}$  estimates the true hyper-gradient  $\nabla h^{(m)}(x)$  with two types of errors: the estimation error of the decode operation  $(\| y_{I_{dec}}^{(m)} - y_x^{(m)}\|^2)$  and the drift of the learning process  $(\| \hat{y}_i^{(m)} - y_x^{(m)}\|^2)$ , where Lemma C.4 and
Lemma C.5 show bounds for these two errors. Furthermore, Proposition 3.6.b) shows that  $h^{(m)}(x)$  is smooth. Next, we are ready to show the convergence of Algorithm 1:
Theorem 3.7. Suppose Assumptions 3.1-3.4 hold, and we run  $T$  iterations of Algorithm 1, with the learning rates satisfy  $\gamma < \frac{1}{L}$ ,  $\eta \eta_{g} < \frac{1}{2IL}$ , then we have:
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 1} ^ {T} \left(\mathbb {E} \| \nabla h (x _ {t}) \| ^ {2} + \frac {1}{2 I} \sum_ {i = 1} ^ {I} \mathbb {E} \| \mathbb {E} _ {\xi} [ \bar {\Delta} \hat {x} _ {t, i} ] \| ^ {2}\right) \\ \leq \frac {2 h (x _ {1})}{T I \eta \eta_ {g}} + \frac {\eta \eta_ {g} \bar {L} G _ {2} ^ {2}}{b _ {x} M} + 1 2 I ^ {2} \kappa^ {2} L ^ {2} \eta^ {2} C _ {f} + \frac {4 I ^ {2} \kappa^ {2} L ^ {2} \eta^ {2} \sigma^ {2}}{b _ {y}} \\ + \frac {4 (3 \kappa^ {2} L ^ {2} + \hat {L} ^ {2}) \gamma \sigma^ {2}}{\mu b _ {y}} + 2 G _ {1} ^ {2} + 2 (3 \kappa^ {2} L ^ {2} + \hat {L} ^ {2}) (1 - \mu \gamma) ^ {I _ {d e c}} C _ {0} \\ \end{array}
$$
where the expectation  $\mathbb{E}_{\xi}$  is w.r.t the noise of stochastic gradient.  $b_{x}, b_{y}$  denotes the batch size, and  $G_{1}, G_{2}$  and  $C_0$  are some constants.
Please refer to Theorem C.8 in the Appendix for the proof. We can further control the noise terms in Theorem 3.7 by choosing the learning rates carefully:
Corollary 3.8. Suppose we choose the learning rates as  $\gamma = \min \left(\frac{1}{2L},\left(\frac{1}{C_{\gamma}T}\right)^{1 / 2}\right)$ ,  $\eta = \min \left(1,\left(\frac{8Ib_xM\bar{L}h(x_1)}{TG_2^2}\right)^{1 / 2},\left(\frac{4\bar{L}h(x_1)}{C_\eta I^2T}\right)^{1 / 3}\right)$  and  $\eta_{g} = \frac{1}{2IL}$ , then we have:
$$
\frac {1}{T} \sum_ {t = 1} ^ {T} \mathbb {E} \big (\| \nabla h (x _ {t}) \| ^ {2} + \frac {1}{2 I} \sum_ {i = 1} ^ {I} \| \mathbb {E} _ {\xi} [ \bar {\Delta} \hat {x} _ {t, i} ] \| ^ {2} \big) = O \left(\frac {\kappa^ {3}}{T} + \left(\frac {\kappa^ {5}}{T}\right) ^ {1 / 2} + \left(\frac {\kappa^ {6}}{T ^ {2}}\right) ^ {1 / 3} + \tilde {G}\right)
$$
where  $\tilde{G} = \kappa^2 (1 - \tau \mu)^{2(Q + 1)} + \kappa^4 (1 - \mu \gamma)^{I_{dec}}$ ,  $C_\eta$  and  $C_\gamma$  are some constants.
To reach an  $\epsilon$  stationary point, we need to run Algorithm 1 with  $T = O(\kappa^{5}\epsilon^{-2})$  number of iterations, furthermore, we need  $Q = O(\kappa \log (\frac{\kappa}{\epsilon}))$ ,  $I_{dec} = O(\kappa \log (\frac{\kappa}{\epsilon}))$ . Note that this matches the iteration complexity of standard FL algorithms [50] up to some logarithm factors ( $Q$  and  $I_{dec}$ ). Note that Corollary 3.8 shows that the following term converges to 0:  $\mathbb{E}\| \nabla h(x_t)\|^2 +\frac{1}{2I}\sum_{i = 1}^I\mathbb{E}\| \mathbb{E}_\xi [\bar{\Delta}\hat{x}_{t,i}]\|^2$ . In fact, for the first term, it shows that  $x_{t}$  reaches to a stationary point of the bilevel problem  $h(x)$ , meanwhile, since we have:
$$
\mathbb {E} _ {\xi} [ \bar {\Delta} \hat {x} _ {t, i} ] = \frac {1}{M} \sum_ {m = 1} ^ {M} \sum_ {q = 0} ^ {Q} \tau \nabla_ {x y} g ^ {(m)} (x, y _ {I _ {d e c}} ^ {(m)}) \times (I - \tau \nabla_ {y y} g ^ {(m)} (x, y _ {I _ {d e c}} ^ {(m)})) ^ {q} \nabla f ^ {(m)} (\hat {y} _ {i} ^ {(m)})
$$
Therefore, if  $\| \nabla_{xy}g^{(m)}(x,y)\|$  is lower-bounded and  $\| \nabla_{yy}g^{(m)}(x,y)\|$  is upper-bounded, we also get the learning parameter  $\hat{y}_i^{(m)}$  converges to the stationary point of the local learning problem  $f^{(m)}(y)$ .
4 Application of FedSep to Real-world FL Problems
In FedSep, the separation of the communication layer and the learning layer makes it able to solve various challenges in Federated Learning. In this section, we apply FedSep to solve two challenging problems in Federated Learning: communication-efficient FL and model-heterogeneous FL.
4.1 Communication-efficient Federated Learning