Resolving the Tug-of-War: A Separation of Communication and Learning in Federated Learning
Authors: Junyi Li
Heng Huang
Venue: NeurIPS 2023 poster (2023)
Keywords: Federated Learning

Abstract: 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.

Resolving the Tug-of-War: A Separation of Communication and Learning in Federated Learning
Junyi Li Computer Science University of Maryland College Park junyili. ai @gmail com
Heng Huang
Computer Science
University of Maryland College Park
henghuanghhogmaal. com
Abstract
Federated learning (FL) is promising preeacyprese machine learning paradigm over distributed data. In this paradigm, each client trains the parameter model locally and the server aggregates the parameter from clients period- ically. Therefore, we perform the learning and communication over the same set of parameters. However, we find that learning and communication have fun- damentally divergent requirements for parameter selection, akin to two opposite teams in a tug-of-war game. To mitigate this discrepancy, we introduce FedSep, novel two-layer federated learning framework. FedSep consists of separated communication and learning layers for each client and the two layers are con- nected through decode/encode operations. In particular, the decoding operation is formulated as a minimization problem. We view FedSep as a federated bilevel opti- mization problem and propose efficient algorithm 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-EEfinn Heterogeneous- MModel FL. Empirical validation shows the superior performance of FedSep over various baselines in these tasks.
FL. and
Introduction
In Federated Learning (FL) (37], a set of clients jointly solve a machine learning task under the coordination of central server. The process of FL involves two category of operations: Learning and Communication. For the learning operation, each client optimizes their local objectives, and for the communication operation, clients exchange local parameters 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 two teams engaged in tug-of-wer wfrr war, each pulling in opposite directions. In fact, on the communication side, it imperative that the parameters of all clients reside in a uniform space. Moreover, to mitigate the high communication costs, major bottleneck in FL., it is beneficial to maintain these parameters within a oiwddimesiional 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 sfate-of-the-art large-scale machine learning models (such Transformers 477), high-dimensional parameter space desirable. In summary, a huge discrepancy of requirements to parameter selection exists between the communication and the learning in FL.
"This work partially supported by NSF IIS 1838627, 1837956, 1956002, 221 1492, CNS 2213701, CCF 2217003, DBI 2225775.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

To mitigate this discrepancy, we opt to 'break the rope' in this tug-of-war scenario. More precisely, we propose two-layer structure on clients: a communication layer and learning layer. As shown in Figure we denote this framework as FedSep: The communication layer connects with the server and uploads/downlodds 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 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 federated bilevel problem. The upper level problem corresponds the learning problems 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 the communication parameter. The convergence of the algorithm is guaranteed.
To demonstrate the flexibility of FedSep in in- corporating the contradictory requirements of the communication and learning, we study two real-world FL tasks. In the first task, study the communication--efiiiii
FL. task. In this task, set the dimension of the communica- tion parameter to be much smaller than the learn- ing parameter, furthermore, we decode the com- munication parameter through a LASSO prob- lem. 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 param- eter to be the parameter of the server model, and the learning parameter be 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. Fi- nally, summarize the main contributions of our work as follows:
Upinad
Communication
Lave
Deerde
Learning Loyer
Cleasi
Clisu
Figure 1: The structure of FedSep. FedSep follows the one-server mne-server-multille multiple-clients structure, in partic- ular, each client has two layers: the communication layer and the learning layer. The two layers are connected through the decode and encode opera- tions.
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 minimization problem, and therefore, the FedSep framework has 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 Efficccent FL. Numerical experiments show the superior performance
FL and Model-Heterogeneee of the methods based on FedSep.
Notations. V denotes the full gradient, V. is the partial derivative for variable x, and higher-order derivatives follow similar rules. -|- 1l2 fg-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 to K.
2 Related Works
Federated Learning. 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 [4 have been studied in the literature. Next, model-hetteoggeneous FL. is also studied to solve the device heterogeneity issue. There are two main categories of heterogeneous model methods: based and partial-training (PT) [ 6 I based. KD-based methods treat the server model as the teacher and the clients' model as the student. In contrast, PT-based methods select subset of the server model for each client, and in the aggregation phase, the server simply aggregates the updated sub-models.
K5
45
25
noowleggediisillationn (KD) [34] 22 [
Bilevel Optimization. Bilevel problem is a type of two-layer optimization problem 49 In the machine learning community. Hyper-parameter Optimization problems [28]
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 fixed number of steps and use the back- proppaaooa technique to compute the approximate hyper-gradien
7
36
40 45 More recently, single-loop
algorithms [
20 21 based on alternative updates of lower and upper level problems are proposed. However, less work studies bilevel problems under the federated setting. studies type of federated bilevel problems where the lower level problems are federated.
18
23
30
26
g
51
31
10
33 32 In particular, [33 solves the noisy label problem in FL. Note that our FedSep is formulated as a type of federated bilevel problems, however, 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.
46
problem is different from that in [46] 33). Firstly, our
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 FedSep has the one-server-muliipeeclien structure 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 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 neural network over a dataset. We assume that the learning layer on client m E [M] has parameter
g(m)) Iddm)) Note that the choice of (()) 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 E TRP. 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, abstract the decode operation as solving a minimization problem, and the encode operation be its reverse operation. As result, we have the following bilevel formulation for the FedSep framework.
M
1 M
min h(x) = M ZERP
Σ m=1
(m) = M
Σ m=1
y! = arg min g(') ((m)) ((m) (yfmm))
(1)
()m) ggd(())
In Eq. (1. M is the number of clients; f(m) (v), m E [M] denotes the local problem solved by the client m; (mm) (z, y) is the minimization problem solved in the decode operation by the client m;; is the parameter of the communication layer; ((m)) denotes the parameter of the learning layer on the client m.
Next, we propose a three-stage algorithm to solve Eq. [ As shown by Algorithm I we perform T global steps, and each step T, we randomly choose subset of clients M, to perform the training and then the server performs aggregation. The Client training is divided into three stages: Decode stage (lines 7-1,), Learning stage (lines 14-17), and Encode Stage (line 19).
Decode Stage. In Eq. ( the decode operation solves the lower optimization problem (1t) (2,8) y) to obtain (m²) We denote the decode operator on the miah client by (ec (:), thus we have: ya
(m)
(ec(m) {x]. In Algorithm I we solve ((m) (I,, y) approximately with Laee steps of stochastic gradient
3

---

Algorithm Separating Communication and Learning in FL (FedSep)
1: Input: Initial states 21; learning rates 7/ ", mini-batchsize dz. by
2: for = 1 to T do
3: Randomly sample a subset M₁ of clients;
4: for m E Mr in parallel do
5: // Decode stage, estimate Va fm] Decl")) {z};
Receive the global state zy from the server and set x(()) = 2t, and initialize yo bm); for 1 to dec do
6:
7:
8:
Randomly sample minibatch of by samples By:
9:
yi (mm end for (m)
| 7 VV (yggmm (x((m) (m)
11:11: By)
10:
Set 30Th Viac as the estimation of Declm) (z!m));
12:
// Learning stage, optimize (())
13:
for i to / do
14:
Randomly sample minibatch of by samples Bg: 0/m) -!-1 V end for
f( (1)m))
16:
// Encode stage, encode the update of the learning layer back to the communication layer;
18: Set Agim) = g(m) (((")) and compute Aa!") EEne (m)) {Ai(m)) }, where Enc {-} is
(m)
defined in Eq. (3) and we choose [B_] by.
19:
end for
20:
Fe+1 Zt [Mdl EmEM, AgA
end for
descent (Line 8-11) and use the output approximation of (m) i.e. Viace Dec (172) yz
(m)
(x((m)) se
where Dec (m)
(-) represents the approximation of the exact decoder (ec(m))...
Learning Stage. Next, we perform steps of stochastic gradient descent solve the local learning
problem (y) (Line 14- 17). This stage is similar to the local gradients in classical FL framework.
Note we get 0i ((m) as the updated learning parameter (m)) more formally, we use Agim) = (m)
0)mm) Viace to represent the update in the learning stage.
(m) (m)
Encode Stage. After the learning stage, we get the update of the local learning problem Ag(m), Suppose denote the encode operator Ene(m) then get the update of the communication parameter (m)) (m (Ag(())) (Line 19). In particular, we choose the following encode operator:
(nc( [.} =-==-- - V (m)g((m)) (x, (mm))) 9y.
(m) (zz (mm)))
where (m) (m) (x) ) is the output the exact decode operation and the second equality can be derived following the implicit function theorem under mild assumptions [16]. To reduee the computa- tion complexity, we use Neumann series to approximate the inverse operation in Eq. (2). More specif-
ically, we have the following approximation: (V (V) (x. y)) Σg (I -V (m)g(m)
(m) (z. v))"
where Q and stochastic approximation of Eq. (2) as:
constants. Since in the decode stage,
(m) Vian. get
(m)
Q
Enc
T V Wyygimi (x, vim) B, (I - TV
bng V (m) (x. vimi B2))a
(3)
Remark We use Eq. (2 as the encode operator due to the following fact about the hyper-gradient
(m(m) (x) y
(m)
(mm) (m)))
Note that suppose Aj(m) - (mm) (y(mm) which means the learning layer performs one step of gradient descent with learning rate in the learning stage of Algorithm then have: (m))
(mm)

---

EncfAg Az(m), Therefore, we update the communication parameter z such that it optimize the overall problem (x).
Server Aggregation. Finally, the server aggregates the local updates of communication parameter A m E M₁ to get the new communication parameter as shown in Line 21 of Algorithm
Difference with existing bilevel algorithms. Bilevel optimization problems are widely studied in standard
However, FedSep cannot directly apply existing algorithms. bilevel algorithm, each step of update to the upper level variable requires (approximately) solving of the corresponding lower level problem. However, in each epoch FedSep, we only solve the lower level problem once and perform multiple steps of update to the upper level variable.
the literature [6 23
3.1 Convergence Analysis
In this section, we provide the convergence analysis of Algorithm result, we first make the following assumptions.
[
Before we state the convergence
Assumption 3.1. Function (y) is possibly non-convex and g/m) (ar, y) is ju-strongly convex w.zr y for any given x.
Assumption 3.2. Function f(m) is L-smooth and has Cy-bounded gradient.
Assumption 3.3. Function (x, y) is L-smooth; V (zygg(m))( (z, y) and V Vy (z, y) are Lipschitz continuous with constants Lan and Lyz, respectively.
Assumption 3.4. We have unbiased stochastic first-order and second-orde gradient oracle with bounded variance.
Note that Assumption 3.1 Assumption 3.3 are standard assumptions used to analyze bilevel prob-
lems 23 16]. Assumption 3.4 is standard in analyzing stochastic optimization problems. For full
version of Assumption 3.4 please refer to Assumption C.I in the Appendix.
In the learning stage of Algorithm [ perform I steps of updates optimize the learning problem,
and the update has the form of: Ag(m) = 01" (m)) y0
(m)
= Σi= (
(0) (m)) Bg). therefore we
have: (m)) Enc (m))
(Agiml) E.= Enc (m) (-V f(m) (i) (m) B,)). We denote Ae(m))
(m)
Enc (m)
{V f(m))() (0. mm) BB)). Then the main effort of the proof is bounding the estimation error
(m) to the hyper-gradient (m) (x)).
of f,i
More specifically, we first show that Aa(m) estimates a term μ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 t < 1, we have
||E ( (m)) S G1, where ph -Vayg (m) (x, yland (V vog (m) (x, (mm)) (((m) (0. (m)) and
M4
EJ (m) EE[A
EJ
( i,s < G2. Gn and G2 are some constants related to Q Eq. and
1f mini-batch size.
Please refer to Proposition C.3 for the specific form of the constants in Proposition 3.5 Next, we have that μú (m) ze Vh(m)(x) as shown in the following proposition:
Proposition 3.6. Suppose Assumptions 3.2 and 3.3 hold, the following statements hold:
a) ||4 (m) (x)m)())) < 2f. (
ya (m)) 112 + 2x 1.210
yz (m)) 112, where = (
b) ((m) (x) is Lipschitz continuous in x with constant L i.e., for any given æ 01, E RP, we have [[Vh(m)) - (m))) (x1) E2| - xill, where L = (1x²).
where we denote the condition number
L/μ.
Please refer to Proposition C.2 for the proof. Proposition 3.6 shows that n (m) estimates the true hyper-gradient (x) with two types of errors: the estimation error of the decode operation (1ly)m) (m) y/ 113) and the drift of the learning process (11g! ys 11m)) where Lemma C.4 and
5

---

Lemma C.5 show bounds for these two errors. Furthermore, Proposition 3.6 b) shows that (1m)) (x) is smooth. Next, we are ready to show the convergence of Algorithm
Theorem 3.7. Suppose Assumptions
3. 3.4 hold, and run T iterations of Algorithm 7 with the
learning rates satisfy y < 1. L. 7
21L. then we have:
T
E [h
+
1
18 EE (Ar,,)l
T
27
2h(x1)
my mgLG2
412
+ x2
Thmm
byM
by
1(3r212 + ??? + 2G + 132 f1 - pr)) Co pbb
where the expectation Eg is w.I.t the noise of stochastic gradient. by. by denotes the batch size, and G1. G2 and Co 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 choose the learning rates Y
min( 2t (,T)) n
min
1,
8Ib, ML((m))
1/2
(Lh(az) C
1/3
J. and Ml 21L'
then we have:
TG3
T
1/2
1/3
B (((V((x))1
Σ
| (1))) =0
K
+
+ G
T
21
T
T
T2
1=1
where G (x) 7())((+)))))))) (1 (1- / fune, Cn and C.y, are some constants.
To reach stationary point, need to run Algorithm I with T s ((x))=>)))))) number of iterations,
furthermore, we need Q ( log( Idee (r log )). Note that this matches the iteration complexity of standard FL algorithms (50) up some logarithm factors (Q and Idec). Note that
Corollary 3.8 shows that the following term converges O: [z))] 21 E'n E|| ElEg [ In fact, for the first term, it shows that 2a reaches to stationary point of the bilevel problem h(z), meanwhile, since we have;
Be A
1 M
Σ
VV (myg(m)) (x,, (me)
(I V
yygg(m) (m)) (mm)) )) (m)) (m))))
Therefore, if ( (zz 1) is lower-bounded and | V (m)) (xr, 3)ll is upper- aoper-bounded, we also get
the learning parameter ii (m) converges to the stationary point of the local learning problem ((m) (y))
Application of FedSep to Real-world FL Problems
In FedSep, the separation of the communication layer and the learning layer makes it able solve various challenges in Federated Learning. In this section, we apply FedSep to solve two challenging problems in Federated Learning: communication-efficiene FL and model-hettrogennous FL.
4.1 Communication--ffciinnt Federated Learning
Communication-cost is one of the major bottlenecks in Federated Learning due slow connections between clients and the server. In Federated Learning, compression is commonly employed to mitigate communication costs, where clients compress the local updates before transferring to the server. A common compressor is choosing the Top-K coordinates, however, this approach is only good at reducing the upload communication cost. Since clients have different Top-K coordinates, the aggregated updates are often dense. More complicated approaches can save both upload and download communication cost, such as the Count-sketch based compressor 42. In fact, we can develop a simple yet effective approach to reduce the communication cost based on FedSep.

---

In our FedSep framework, suppose we have the learning parameter & E RR, and the communication parameter E Rt. We choose p < to have communication efficient federated learning algorithm. More specifically, we consider the following formulation:
min 1 Σ WERP M
C(0(T)) (D(m)) s.t. 6l(²)
arg min
[| 1e0l-wll - :|10||
(4)
BER
m=
where D(m denotes the training distribution of the mea client; L(-) denotes the loss function. In particular, S(m) Rexd (p < d) is random sketch matrix whose coordinates sampled from Gaussian distribution. We choose the quadratic optimization problem (LASSO) as the decode function.
Eq. (4) is special case Eq. [ 0 corresponds to the learning parameter y. w corresponds to the communication parameter C(0; D("")) corresponds to the learning problem (1))(y) and the decoding problem LASSO problem. We can solve it through Algorithm
The LASSO problem of the decode operation involves a non-smooth L regularization term and we [3]. For the encode operation, we have:
solve it with the standard proximal gradient method
Enc (m) [.] -.olm -
U (1- (()(mm)) s(m))()
(5)
where I is the identity matrix and U Di + +(S(m)) an indicator operator, which outputs if the absolute value of the input is greater than 7.8. Please refer Appendix A.1 for more details of the proximal gradient method and the encode operator. Remark 2. If we set B in Eq. 4 i.e. we remove the sparsity constraints, both encode and decode operators have explicit solutions. More specifically, ((S(m))T& 5(m)) - (5 and the encode operator Enc(() {-} = ((m) ((5(m))"' -1 If we choose ((m) S for all m [M], we get linear compressor.
((m))!(")) where (,,())
have the decode operator (ec(m)) {-} ==
4.2 Model-HHterogenneon Federated Learning
In practical Federated Learning applications, the scale the model is inherently limited by the on-device resources of the participating clients, which often exhibit significant diversity. As result, we can choose different scale of models based on the available resources of each individual client. One widely-used approach is the sub-model extraction, i.e. each client selects part of the server model to perform local training. Different strategies can be used to do extraction: fixed [I]] (2]. random [6] and roll [I]. In contrast, we provide data-dependent approach select the sub-models. As shown by the following formulation:
M
min Σ (()(m)) ( (m)) WERF M
(6)
m=1
s.t. g(mm) = (m) @ w, (m)) arg min wi ((m)) + (R(T(R), (m))
(E(111)
where D(m)) denotes the training distribution; L(-) denotes the loss function and 0 and w are the learning parameter. Note that denotes the parameter of the full model and f is the parameter of the sub-model. In partieular, denote mask vector (m)) whose value is from f0, 1}. @ represents the coordinate-wise multiplication.
[m] For the decode operation, we have e!Th) aw @ w. So we first need to find an optimal mask (m) instead solving the complicated integer programming in Eq. (6), we solve the following relaxed continuous problem:
(m) arg minC(a @ w; ((")) (R)( (m)) T(a) oE(0.1]p
Eq. 77) includes two parts. C(a w; (!m)) measures the loss of the extracted sub-model over D (m)
The second part R(T (m)) ( is called resource loss 15 Tiol denotes FLOPS of
val
7

---

Test Accuracy w.r.t Communication Rate for FedSep and other baseline methods for MNIST Dataset. The first two plots show results under the LLD case, and the last two plots show results for the Non-L.LD case. p in the second and the fourth plot is the dimension of the communication parameter. The local learning steps are set as =
Figure 2:
the server model and T(a) denotes FLOPS of the sub-model. Therefore, p(m) [0. con- trols the size of the sub-model. Following [[5], choose R(-) as: p( Log(maz(T(aa)) h/o)) the loss value small and meanwhile, satisfies the FLOPS constraint. Note that this contrast to the data-independent sub-model selection methods [11 1.
In summary, choose the 'best' sub-model which keeps
Next for the encode operator Enc(.). J. we follow the definition of Eq. 2 to have:
Enc(m) {.} V., v( v (a @ w Diag {a ( [wa V (m))
(m)
Where the last equality follows from the chain rule. Finally, combine Eq. 7 and Eq. 8 we can solve Eq. (6) with Algorithm] 11
Remark For the encode operator Eq. 8 we have (nc(m)) (.) == - -Diagfal ( by omitting the
second term in Eq. (8, where A0 denotes the updates to the sub-model 0. As shown in Line 21 of
Algorithm 1 we have the update of the server model Aw - [.M. EmEM, "h Diaglas" (mm) HAO In
other words, the server simply aggregates the updates of the sub-model, this the aggregation rule widely used in the partial training literature [1T 1.
5 Numerical Experiments
In this section, we perform numerical experiments to validate the efficacy of FedSep in solving communication-effiienn FI. and wodel--hetengeenous FL discussed in Section 4 The code is writ- ten with Pytorch (39], and the Federated Learning environment simulated via Pttorch.Dsstrbuut Package. We used servers with AMD EPYC 7763 64-core CPU and 8 NVIDIA V100 GPUs to our experiments.
5.1 Experiments for Communication--Ffiin Federated Learning
In this section, we provide numerical experiments for the communication-eficieet FL. task discussed Section 4.1 We first introduce the experimental settings.
Dataset. We consider MNIST 229 and CIFAR-IO (27] in our experiments. We create the federated version of these datasets by evenly distributing the training samples among 10 clients. We consider an L.I. setting and a Non-I.I.D setting. For the I.I.D setting, each client has data samples from all 10 classes, while for the Non-I.L.D setting, each client only has data samples from classes.
Model. For experiments related to MNIST, we use a four-layer convolutional neural network. The first three layers contain 32 kernels of size 3x3 and the fourth layer contains 32 kernels of size 2x2. For the experiments related to CIFAR-10 10, we increase the number of kernels to 64. The total number of parameters of the model is around 10a. For the sketch matrix ((m) our FedSep, we let all clients choose the same sketch matrix, and we test the dimension of the communication parameter p E {10, 100, 500, 1000, 2000, 5000, 10000, 20000).
Baselines and Hyper-parameters For baselines, we first compare with the uncompressed base- line (the communication parameter and the learning parameter are the same dimension); Next we compare with the local Top-K compressor (with error feedback [4] and the Count-Sketch compressor Please refer the Appendix for the hyper-parameter selection FedSep and baselines.

---

parameterize the mask with neural Table 1: Test accuracy comparison between FedSep with other model-hetterogeneous FL baseline
methods, High data heterogeneity represents = for CIFAR-10 and = 20 for CIFAR-100;
Lower data heterogeneity represents K 5 for CIFAR-10 and K 50 for CIFAR-100.
High Data Heterogeneity
Low Data Heterogeneity
Method
CIFAR-10
CIFAR-100
CIEAR-10
CIEAR-100
FedDF
34
73.81 (# 0.42)
31.87 (# 0.46)
76.55 (# 0.32)
37.87 (# 0.31)
KD-based
DS-FL
65.27 (# 0.53)
29.12 (# 0.51)
68.44 (# 0.47)
33.56 (# 0.55)
Fed-ET 9
78.66 0.31)
35.78 0.45)
81.13 (# 0.28)
41.58 (# 0.36)
HeteroFF [1
63,90 (# 2.74)
52.38 (# 0.80)
73.19 (# 1.71)
57.44 (# 0.42)
PT- PT-based
Federated Dropout
46.64 (# 3.05)
45.07 0.07)
76.20 (# 2.53)
46.40 (# 0.21)
ZeroFL [4]
64.61 2.18)
51.39 0.45)
83.31 (# 0.78)
53,62 (+ 0.51)
FedDST 5
67.65 (# 1.27)
54.21 (# 0.34)
84,57 (# 0.28)
54.97 (= 0.44)
Flash 2
67.08 (+ 1.46)
54.92 (+ 0.29)
84.61 (+ 0.37)
55.04 (+ 0.32)
FedRolex 1
69. (# 1.50)
56.57 (# 0.15)
84.45 (+ 0.36)
58,73 (# 0.33)
FedSep (Ours)
71.13 (= 0.94)
58.16 (# 0.25)
84.61 (+ 0.37)
61.41 (= 0.29)
Homogeneeus (smallest)
38.82 (#: 0.88)
12.69 (#: 0.50)
46.86 (t: 0.54)
19. 19,70 (+: 0.34)
Homogencous (largest)
75.74 (# 0.42)
60.89 (# 0.60)
84.48 (# 0.58)
62.51 (# 0.20)
The experimental results are summarized in For MNIST (Figure 3 in the Appendix for CIFAR- 10). As shown by the figures, FedSep outperforms other baselines. In particular, our FedSep achieves much better performance in the nigh-compression-rat regime and can get similar performance as other baselines in the low-compression-rae rate regime. Please refer to the Appendix for more ablation studies.
5.2 Experiments for Hedee--eetsoggenoous Federated Learning
In this section, we provide numerical experiments for the task of model-heterogeenous FL task discussed in Section 4.2 First, we introduce the experimental settings.
Dataset. We consider CIFAR-I0 and CIFAR- (27] in our experiments. We create Federated datasets by evenly distributing images among clients; we have 100 clients in our experiments. We create data heterogeneity by controlling the number of classes K client can have [1]. For CIFAR-10, we test E {2, 4, 5, 8). while for CIFAR-100, we test K E (5. 10, 20, 50}. Note that smaller values of K mean higher heterogeneity of the data. For FedSep. the validation and training set are the same.
Model. We choose R8 17 in the experiments. Following the setting in [Ι we replace the batch normalization layer with static batch normalization and add scalar module after each convolution layer. Instead of using the coordinate mask as defined in Eq. (6), we select kernels in each convolutional layer of ResNet- 18; furthermore,
network H HNO) [15). i.e. a = (o));; finally, we reuse the masks and only update the masks every two epochs, this stabilizes the training. This reduces the dimension of that need to optimize. In experiments, the size of the submodel (Eq. (6) of a client is randomly chosen from 0., 0.25. 0.125, 0.0625). where the ratio p w.z. the server model. For our FedSep, after optimizing Eq. (7). we select the top p percent kernels of each convolutional layer by the value of the leaned mask a, and then only the selected kernels to perform the training. In the encode operation, we only evaluate the first term of Eq. (8) as stated in Remark 3
Baselines and flper-parameters. We compare with both state-of-the-art Knowledge Distillation- based methods: FedDF [34). DS-FL (22) and Fed-ET [9, and Partial Training Based methods:
HeteroFL [11. ZeroFL [1), Federated Dropout (6), FedDST (5), Flash [2) and FedRolex [I].
Fjord 19 gets similar performance as HeteroFL, so we do not include in the comparision.
For ZeroFL, FedDST and Flash, we consider their heterogeneous-moode version by varying the compression rate among clients.
The experimental results are summarized in Table I As shown in the table, our FedSep gets comparable performance with the Knowledge Distillation based methods, which needs additional public data, furthermore, FedSep outperforms the partial training based methods, including the state-of-the-art FedRolex ( method. This result shows the effectiveness of the adaptive sub-model extraction strategy used by our FedSep. Our method chooses the most important part of the model

---

at each step for clients, thus our method converges faster than other data-inde iata-indeeendent methods. For more experimental results, please refer the Appendix.
6 Conclusion
In this work, we propose a novel federated learning framework, i.e. FedSep with separated commu- nication and learning components. Based on the observation that communication and learning set opposite requirements to the parameter, we let clients to have two layers: communication layer and a learning layer. We formulate the decode operation from the communication layer to the learning layer as solving minimization problem, therefore, FedSep has bilevel structure. We propose an efficient algorithm to solve this bilevel problem and also prove that the algorithms converge to a stationary point with rate ( Finally, we apply FedSep to solve the communication--fficcee. FL task and the model heterogeneous FL task. Numerical experiments show the superior performance of our FedSep over various baselines.
References
[1] S. Alam, L. Liu, M. Yan, and M. Zhang. Fedrolex: Model-hettroggeneous federated learning with rolling sub-model extraction. arXiv preprint arXiv:2212.01545: 2022.
[2] S. Babakniya, S. Kundu, S. Prakash, Y. Niu, and S. Avestimehr. Federated sparse training: Lot- tery aware model compression for resource constrained edge. arXiv preprint arXiv:2208. 13092, 2022.
[3] A. Beck and M. Teboulle. A fast iterative shrinkage-thresholdingg algorithm for linear inverse problems. SIAM journal on imaging sciences, 2(1):183-2)) 202, 2009.
[4] Y. Bengio. Gradient-bast -assed optimization of hyperparameters. Neural computation, 12(8): 1900, 2000.
[5] S. Bibikar, H. Vikalo, Z. Wang, and X. Chen. Federated dynamic sparse training: Computing less, communicating less, yet learning better. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 6080-6088, 2022.
[6] S. Caldas, J. Konečny, H. B. McMahan, and A. Talwalkar: Expanding the reach of federated learning by reducing client resource requirements. arXiv preprint arXiv:1812.072 10, 2018.
[7] 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-1288 IEEE, 1999.
[8] T. Chen, Y. Sun, and W. Yin. A tinlee-tmeescae stochastic bilevel optimization method. arXiv preprint wxXv:22022.06711, 2021.
[9] J. Cho, A. Manoel, G. Joshi, R. Sim, and D. Dimitriadis. Heterogeneous ensemble knowledge transfer for training large models in federated learning. arXiv preprinl w:Xiv:2204.12700 2022.
[10] 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 urXiv:2201.13400 2022.
[11] E. Diao, J. Ding, and V. Tarokh. Heterofl: Computation and communication efficient federated learning for heterogeneous clients. arXiv preprint arXiv:2010001264, 2020.
[12] C. B. Do, C.-S. Foo, and A. Y. Ng. Efficient multiple hyperparameter learning for log- lineai models. In NIPS, volume 2007, pages 377-384. Citeseer, 2007.
[13] J. Domke. Generic methods for optimization based modeling. In Artificial Intelligence and Statistics, pages 318-326. PMLR, 2012.
[14] L. Franceschi, M. Donini, P. Frasconi, and M. Pontil. Forward and reverse gradient-basee based hyperparameter optimization. Proceedings of the 34th International Conference on Machine Learning- Volume 70, pages 1165-1173. JMLR. org, 2017.
10

---

[15] S. Gao, F. Huang, W. Cai, and H. Huang. Network pruming via performance maximization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 9270-9280,, 2021.
[16] S. Ghadimi and M. Wang. Approximation methods for bilevel programming. arXiv preprint 1rXiv:1802.022444 2018.
[17] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770- 778, 2016.
[18] 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.
[19] S. Horvath, S. Laskaridis, M. Almeida, I. Leontiadis, S. Venieris, and N. Lane. Fjord: Fair and accurate federated learning under heterogeneous targets with ordered dropout. Advances in Neural Information Processing Systems, 34:12876-12889, 2021.
[20] F. Huang and H. Huang. Biadam: Fast adaptive bilevel optimization methods. arXiv preprint arXiv:2106. /396, 2021.
[21] Huang and H. Huang. Enhanced bilevel optimization via bregman distance. arXiv preprint arXiv:2107.12300 2021.
[22] S. Itahara, T. Nishio, Y. Koda, M. Morikura, and K. Yamamoto. Distillation-based semi- supervised federated learning for communication-efice private data. IEEE Transactions on Mobile Computing, 22(():191-205 2021.
collaborative training with non-iid
[23] K. Ji, J. Yang, and Liang. Provably faster algorithms for bilevel optimization and applications to meta-learning. arXiv preprint arXiv:2010.07962. 2020.
[24] S. P. Karimireddy, S. Kale, M. Mohri, S. J. Reddi, S. U. Stich, and A. T. Suresh. Scaf- fold: Stochastic controlled averaging for on-device federated learning. arXiv preprins arXiv:1910.06378, 2019.
[25] S. P. Karimireddy, Q. Rebjock, S. Stich, and M. Jaggi. Error feedback fixes signsgd and other gradient compression schemes. In International Conference on Machine Learning, pages 3252-3261. PMLR, 2019.
[26] Khanduri, S. Zeng. M. Hong, H.-T. Wai, Z. Wang, and Z. Yang. A near-optimal algorithm for stochastic bilevel optimization via couble-momentuue arXiv preprint arXiv:2102.07367 2021.
[27] A. Krizhevsky, G. Hinton, et Learning multiple layers of features from tiny images. 2009.
[28] Larsen, L. K. Hansen, C. Svarer, and M. Ohlsson. Design and regularization of neural networks: the optimal use of a validation set. In Neural Networks for Signal Processing VI. Proceedings of the 1996 IEEE Signal Processing Society Workshop, pages 62-71. IEEE, 1996.
[29] Y. LeCun, C. Cortes, and C. Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann. lecun. com/exalb/mnist, 2, 2010.
[30] Li, B. Gu, and H. Huang. Improved bilevel model: Fast and optimal algorithm with theoretical guarantee. arXiv preprint arXiv:2009.0066 2020.
[31] J. Li, B. Gu, and H. Huang, A fully single loop algorithm for bilevel optimization without hessian inverse. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 7426-7434, 2022.
[32] Li. F. Huang, and H. Huang. Communication--fficii federated bilevel optimization with local and global lower level problems. arXiv preprint arXiv:2302.06700) 2023.
[33] J. Li, J. Pei, and H. Huang. Communication-efiielee robust federated learning with noisy labels. Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 914-924, 2022.

---

[34] T. Lin. L. Kong, S. U. Stich, and M. Jaggi. Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems, 33:2351-2363, 2020.
[35] Y. Lin, S. Han, H. Mao, Wang, and W. J. Dally. Deep gradient compression: Reducing the communication bandwidth for distributed training, arXiv preprint arXiv: 1rXiv:171..01887, 2017.
[36] D. Maclaurin, D. Duvenaud, and R. Adams. maadient-based hyperparameter optimization through reversible learning. In International Conference on Machine Learning, pages 2113- 2122, 2015.
[37] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas, Communication--fficinn learning of deep networks from decentralized data. In Artificial Intelligence and Statistics, pages 1273-1282. PMLR, 2017.
[38] K. Nandakumar, Ratha, S. Pankanti, and Halevi. Towards deep neural network training on encrypted data. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition Workshops, pages 0-0, 2019.
[39] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury. G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems, pages 8024-8035, 2019.
[40] F. Pedregosa. Hyperparameter optimization with approximate gradient. arXiv preprint 1rXiv:1602.02355, 2016.
[41] X. Qiu, J. Fernandez-No Bernandez-Marques, P. Gusmao, Y. Gao, Parcollet, and N. D. Lane. Ze- rofl: Efficient on-device training for federated learning with local sparsity. arXiv preprint arXiv:2208. 02507, 2022.
[42] D. Rothchild, A. Panda, E. Ullah, N. Ivkin, I. Stoica, V. Braverman, J. Gonzalez, and R. Arora. Fetchsgd: Communication-effiiene federated learning with sketching. In International Confer- ence on Machine Learming, pages 8253-8265, PMLR, 2020.
[43] A. Shaban, C.-A. Cheng, N. Hatch, and B. Boots. Truncated back-propagation for bilevel optimization. arXiv preprint 1010::/810.0066 2018
[44] S. U. Stich. Local sgd converges fast and communicates little. arXiv preprint arXiv:1805.0976) 2018.
[45] S. U. Stich, J.-B. Cordonnier, and M. Jaggi. Sparsified sgd with memory. arXiv preprint arXiv:1809.07555 2018.
[46] D. A. Tarzanagh, M. Li, C. Thrampoulidis, and Oymak. Fednest: Federated bilevel, minimax, 2022.
and compositional optimization. arXiv preprint
[47] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, E. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[48] W. Wen, C. Xu. F. Yan. C. Wu, Y. Wang. Y. Chen, and H. Li. Terngrad: Temary gradients to reduce communication in distributed deep learning. arXiv preprint arXiv:1705.07877 2017.
[49] R. A. Willoughby. Solutions of ill ill-posed problems (an tikhonov and vy arsenin). SIAM Review, 21(2):266, 1979.
[50] B. Woodworth. The minimax complexity distributed optimization. arXiv preprint AXXi:2119.005344 2021.
[51] J. Yang, K. Ji, and Y. Liang. Provably faster algorithms for bilevel optimization. arXiv preprint arXiv:2106. 04692, 2021.
12


Reviews:
[{"review_id": "o61dlTvQkn", "reviewer": "Reviewer_WQf3", "summary": "The paper uses bilevel optimization in a federated learning setting. A unique decomposition of communication and learning has been identified, which has applications for reducing communication overhead and supporting heterogeneous models. The theoretical convergence guarantee has been presented and experiments also show the advantage of the proposed method.", "questions": "- What is the additional computation overhead compared to standard FL algorithms such as FedAvg?\n- What are the key technical (mathematical) challenges and novelties?\n- In the experiments in Section 5.1, biased compressors such as top-K usually only works well with error feedback [a]. Is error feedback used in the baseline methods? In addition, are the first and third plots in Figure 2 results at convergence? Theoretically, all methods (when using error feedback in baselines) should converge to the optimal solution when the experiments are run long enough. More explanation is needed on why a much worse accuracy is obtained when the compression rate is high. Is it a matter of hyper parameter tuning? It may be better to plot the accuracy against the amount of communication (e.g., in the number of bytes transmitted, *not* in communication rounds), to compare the accuracies of different methods and different compression rates at the same amount of communication.\n- In Section 5.2, the advantage of having heterogeneous models is not quite clear. \n\n[a] Stich, Sebastian U., and Sai Praneeth Karimireddy. \"The error-feedback framework: Better rates for SGD with delayed gradients and compressed communication.\" arXiv preprint arXiv:1909.05350 (2019).", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The use of bilevel optimization as a decomposition into communication and learning in a FL setting is interesting. \n- The framework is general and supports multiple application scenarios of FL.\n- Theoretical convergence bound is provided.", "weaknesses": "- The use of encoder/decoder structure may incur additional computation compared to common FL algorithms. It would be nice to quantify the amount of such additional computation.\n- The extension of common bilevel optimization algorithms to include multiple update steps (as pointed out at the end of page 4) seems to be somewhat similar to the idea of local updates in FedAvg / local SGD algorithms. It would be nice to identify what are the key technical challenges and novel solution techniques in this extension. \n- There is space for improvement in the experiments (see details below).\n- The writing could be improved to highlight the usefulness of encoder/decoder in the context of FL, in early parts of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for spending time reviewing our paper. Below are our responses to your questions:\n\nResponse to Concerns:\n\n**A(W1: Additional computational cost)**: Compared to FedAvg, FedSep needs to perform extra decode and encode operations. For the decode operation, we need to perform $I_{dec}$ of gradient descent steps to optimize the decode problem (Line 7-12 of Algorithm 1), while for the encode operation, we need to perform $Q$ Hessian-vector products (Eq.3). Since the estimation error of both decode and encode operations have linear convergence, we can choose small values of $I_{dec}$ and $Q$ in practice, thus is negligible. Furthermore, to incorporate the communication-learning trade-offs, most existing approaches needs some extra encode/decode operations. For example, FetchSGD [1] is a communication-efficient FL method and uses Count-sketch to compress the local updates before communicating with the server. \n\n**A(W2: Technical novelty and difficulty)**: Our Algorithm 1 is much more efficient compared to a naive application of the bilevel optimization algorithm. Suppose we use a standard bilevel optimization algorithm to solve FedSep, then we perform multiple rounds of the following steps 1-4: step 1: we decode the communication parameter to get the learning parameter (solve the inner optimization problem); step 2: we compute the gradient of the learning objective w.r.t. the learning parameter; step 3: we compute the hyper-gradient w.r.t. the communication parameter through the encode operation; step 4: we use gradient descent to update the communication parameter. \n\nThis algorithm is actually a direct application of the FedAvg to a bilevel optimization problem, which is inefficient due to performing the decode (step 1) and encode (step 3) operations multiple times. **Our idea is to optimize the learning objective at the learning parameter space**. As shown by Algorithm 1, we perform the following steps: step 1: same as the naive approach; step 2: we perform multiple gradient descent steps to optimize learning objective at the learning parameter space; step 3: we map the learning parameter update back to the communication parameter space through the encode operation; step 4: same as the naive approach. **Compared to the naive approach, we only perform decode/encode operation once.**\n\nThe convergence analysis is also challenging. If we follow the naive approach, we can get accurate estimation of the hyper-gradient (gradient w.r.t to the communication parameter), and the convergence analysis is a straightforward application of FedAvg analysis. However, since we perform multiple gradient descent steps at the learning parameter space, we would have a biased hyper-gradient estimation. Then we need to bound this hyper-gradient estimation error carefully (see Lemma C.6 for more details). Finally, in corollary 3.8, we show that the communication parameter $x$ reaches to a stationary point of the bilevel problem (in Eq.1), and the learning parameter $y^{(m)}$ converges to the stationary point of the local learning problem $f^{(m)}(y)$.\n\n**A(W3: Experimental setting)**: See our response to **Q3** below.\n\n**A(W4: Usefulness of encoder/decode)**: The encoder/decoder are bridges to connect the communication parameter and learning parameter. Selecting the appropriate decode/encode operation is central to applying FedSep in addressing various challenges within Federated Learning (FL). As shown by the communication-efficient FL task and model-heterogenous task.\n\nResponse to Questions:\n\n**A(Q1: Additional computational cost)**: Please see response to the **W1** above.\n\n**A(Q2: Technical novelty and difficulty)**: Please see response to the **W2** above.\n\n**A(Q3: Experimental setting)**:  We include the error feedback for Top-K method. \n\nWe compare our FedSep with other baselines using the following schema in the first and third plot of Figure 2: First, we run the uncompressed baseline with a given number of global epochs (until convergence), and we denote its communication budget as $P$. Note that the uncompressed method has a compression rate of 1. **For other methods, we vary hyper-parameters of each method under a given compression rate and report the best accuracy**. Take the Top-K method as an example, under the compression rate 10, we run Top-K with different values of K until reaching the communication budget of $0.1P$, and report the best accuracy achieved.\nIn summary, the compression rate is defined by the ratio between the amount of bits transmitted with the total amount of communication budget ($P$). By using the this comparison scheme, we can **compare the rate limit of each method under a given communication budget**, note that [1] also uses a similar scheme for comparison. Furthermore, by reading the first and third plot of Figure 2 from right to left, we can see that the accuracy indeed increases when the amount of communication increases.\n\n**A(Q4: Advantages of having heterogeneous models)**: Training heterogeneous models is used to incorporate the different resource limit of clients, some clients have the capacity to train large model, while other clients can only train a small model. In particular, we implement one way of achieving heterogeneous model with FedSep: data-driven sub-network selection. As shown by Table 1, our FedSep outperform other sub-network selection methods, and achieve comparable accuracy with the Homogeneous (large) baseline, where all clients train the full-size model.\n\n\n\n**References**\n\n[1] D. Rothchild, A. Panda, E. Ullah, N. Ivkin, I. Stoica, V. Braverman, J. Gonzalez, and R. Arora.453\nFetchsgd: Communication-efficient federated learning with sketching. In International Confer-454\nence on Machine Learning, pages 8253–8265. PMLR, 2020.\n"}, {"review_id": "eSHKiANIrw", "reviewer": "Reviewer_Dwjg", "summary": "This paper asserts that the tasks of communication and learning are at odds in federated learning. As such, a new approach is suggested that separates these tasks. The approach comprises an encode-decode operation, where decoding is cast as an optimization problem. The overall learning task is formulated as a bilevel optimization problem. Standard first-order gradient-based approaches are then employed to solve the bilevel optimization problem. The structure of the proposed algorithm mirrors that of typical FL algorithms. A convergence bound is proven for this algorithm. Some applications pertaining to communication-efficiency and model-heterogeneity in FL are also discussed. ", "questions": "No questions other than the ones above.", "limitations": "I couldn't find a clear discussion of the limitations. ", "rating": 3, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The particular \"separation\" framework suggested in this paper appears to be novel; I have not seen anything exactly like this in earlier FL work.\n\n- One advantage of the proposed bilevel approach is that each agent can train a separate local model $y^{(m)}$ of its own. This allows for personalization in the face of data-heterogeneity.", "weaknesses": "I have several major concerns ranging from the motivation to the utility of the proposed approach. Let me elaborate on them below. \n\n- The entire paper is based on the premise that communication and learning are conflicting goals in federated learning. However, this statement is never formalized at any point in the paper. The only discussion regarding the tension between communication and learning shows up in lines 25-30, which don't convey anything concrete.  \n\n- Continuing with the above point, I also missed the motivation for the specific bilevel formulation. By now, there are several approaches that guarantee communication-efficiency in RL (see Refs [R1]-[R5] below) using techniques ranging from quantization, sparsification, compressed sensing, etc. It stands to reason that compressing information about models/gradients will naturally come at the cost of performance in learning. So the tension here is evident, and has been well-explored/quantified in the papers I mentioned earlier. In particular, sending less bits of information slackens the rate of convergence, and one can potentially try to investigate the fundamental limits on the rate needed to achieve a desired level of accuracy; see, for instance, [R5]. However, there is no clear discussion of why the FedSep bilevel idea is any significant improvement over any of these MANY existing schemes. \n\n[R1] Federated learning with compression: Unified analysis and sharp guarantees, Haddadpour et al., AISTATS 21\n\n[R2] Linear convergence in federated learning: Tackling client heterogeneity and sparse gradients, Mitra et al., NeuRIPS 21\n\n[R3] Optimizing the communication-accuracy trade-off in federated learning with rate-distortion theory, Mitchell et al., arXiv 2022\n\n[R4] Communication-Efficient Federated Learning through Importance Sampling, Isik et al., arXiv 2023\n\n[R5] Differentially quantized gradient methods, Yin et al., IEEE Transactions on Information Theory, 2022\n\n- Regarding model-heterogeneity, some of the references I alluded to earlier do account for heterogeneous loss functions across agents. Moreover, if the goal is to allow for personalized models, then why not use any of the existing schemes for personalization in FL (of which, there are many)? See, for instance, ideas based on Moreau Envelopes and MAML in [R6] and [R7], and representation learning in [R8]. If one cares about both personalization and communication-efficiency, I can imagine simply using these ideas in tandem (in meaningful ways). \n\n[R6] Personalized Federated Learning with Moreau Envelopes, Dinh et al, NeuRIPS 2020\n\n[R7] Personalized Federated Learning: A Meta-Learning Approach, Fallah et al., NeuRIPS 2020\n\n[R8] Exploiting Shared Representations for Personalized Federated Learning, Collins et al, ICML 21\n\nThe overarching point I am trying to make here is that each of the key considerations in FL that the authors allude to (communication-efficiency, heterogeneity, personalization, etc), have several existing principled algorithmic solutions. It wasn't apparent to me at all why there is a compelling need to depart from these existing approaches.\n\n- In addition to the motivation, several parts of the paper are somewhat vaguely written. For instance, in Eq. (1), the meaning of the object $g^{(m)}$ is not explained. The encoding operator in Eq.~(2) isn't clear to me either. What is this operation and how does it compare with any of the standard encoding techniques (say for instance, standard scalar and vector quantizers, or sparsifying mechanisms like Top-k)? How many bits are needed to perform this encoding? What is the error caused by this encoding? Is this an unbiased encoding mechanism? No intuition is provided at all about any of these crucial points, making it hard to draw meaningful comparisons with existing schemes.\n\n- The difference with existing federated bilevel optimization algorithms in lines 153-157 is also quite terse. The discussion did not come across as anything fundamental. It wasn't clear to me why the bilevel algorithm proposed can't be analyzed by simply adapting ideas from other existing FL bilevel algorithms. \n\n- The main convergence result in Theorem 3.7 is hard to parse. I was left wondering about several key questions: (i) How does the compression scheme (encoding-decoding) affect the rate of convergence? (ii) How does this trade-off compare with the existing known bounds for compression in FL? (iii) How does the effect of heterogeneity in the agents' loss functions manifest in the bounds? Does this dependence match with those known for federated bilevel optimization? \n\nAlso, the iteration complexity seems to have a $\\kappa^5$ dependence on the condition number $\\kappa$. This seems much larger than what one typically obtains. Thus, I am not convinced about the tightness of the bounds either. \n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for spending time reviewing our paper. Below are our responses to your questions:\n\n**A(W1: Formal discussion of communication/learning conflict)**: Please refer to the answer for **Q1 in the overall response.**\n\n**A(W2: Comparison with existing compression FL methods)**: Please refer to the answer for **Q2 in the overall response.**\n\n**A(W3: Comparison with existing personalized FL models)**: Firstly, FedSep provides a general framework to deal with the communication-learning dilemma in FL. [1]-[3] can be viewed as a special case of FedSep when we choose proper decode functions: For the Moreau Envelopes based Personalized FL [1], we set both the decode optimization problem ($h^{(m)}(x)$) and the local training objective ($f^{(m)}(x)$) as a regularized loss function with l2-norm; for the MAML-based method [2], we set the decode problem as one step gradient descent; while the representation learning [3] method corresponds to a degenerated case where the communication parameter (global representation) and learning parameter (client-specific heads) are not connected through a decode problem. Note that [1]-[3] are designed for the data heterogeneity of FL, our FedSep can also incorporate the model-heterogeneity FL as discussed in the Section 4.2. In model-heterogeneous FL, we can choose different scale of models based on the available resources of each individual client. \n\nNext, although various aspects of the communication-learning dilemma have been studied in the literature, these methods are proposed in **isolation**. FedSep provides the first unified analysis framework for them using the theory of bilevel optimization. Furthermore, our FedSep can be used to design novel algorithms  as demonstrated by the communication-efficient algorithm and model-heterogenous algorithm in Section 4 of the manuscript.\n\n**A(W4: Discussion of the encoder operation)**: $g^{(m)}$ is the minimization problem solved in the decode operation by the client m;\n\nAbout the encode operation Eq.2, it is the **inverse operation of the decode operation**. Recall that the decode operation is defined as $y^{(m)}_x = Dec(x)$, where $y^{(m)}_x$ is the minimizer of $g^{(m)}(\\cdot, x)$. For the encode operation, it measures how the communication parameter is changed when the minimizer $y^{(m)}_x$ changes, i.e. $\\nabla_x y^{(m)}_x = Enc(\\cdot)$. Eq.2 shows an explicit form of $\\nabla_x y^{(m)}_x$ and can be derived following the implicit function theorem (a standard result in bilevel optimization). In section 4.1, we show an instantiation of the encode/decode operations in the context of communication-efficient FL: the decode function is a LASSO problem that solves the high-dimensional learning parameter from a low dimensional communication parameter with a sparse constraint, while the encode operation is a non-linear mapping as shown in Eq.5 in the manuscript. Finally, The dimension of the communication parameter can be chosen arbitrarily, however, we would see a coding-accuracy trade-off. Please refer to our answer to **Q3 in the overall response** for a more detailed analysis of this trade-off.\n\n\n**A(W5: Difference with existing federated bilevel optimization algorithms)**: Our Algorithm 1 is much more efficient compared to a naive application of the bilevel optimization algorithm. Suppose we use a standard bilevel optimization algorithm to solve FedSep, then we perform multiple rounds of the following steps 1-4: step 1: we decode the communication parameter to get the learning parameter (solve the inner optimization problem); step 2: we compute the gradient of the learning objective w.r.t. the learning parameter; step 3: we compute the hyper-gradient w.r.t. the communication parameter through the encode operation; step 4: we use gradient descent to update the communication parameter. \n\nThis algorithm is actually a direct application of the FedAvg to a bilevel optimization problem, which is inefficient due to performing the decode (step 1) and encode (step 3) operations multiple times. **Our idea is to optimize the learning objective at the learning parameter space**. As shown by Algorithm 1, we perform the following steps: step 1: same as the naive approach; step 2: we perform multiple gradient descent steps to optimize learning objective at the learning parameter space; step 3: we map the learning parameter update back to the communication parameter space through the encode operation; step 4: same as the naive approach. **Compared to the naive approach, we only perform decode/encode operation once.**\n\nThe convergence analysis is also challenging. If we follow the naive approach, we can get accurate estimation of the hyper-gradient (gradient w.r.t to the communication parameter), and the convergence analysis is a straightforward application of FedAvg analysis. However, since we perform multiple gradient descent steps at the learning parameter space, we would have a biased hyper-gradient estimation. Then we need to bound this hyper-gradient estimation error carefully (see Lemma C.6 for more details). Finally, in corollary 3.8, we show that the communication parameter $x$ reaches to a stationary point of the bilevel problem (in Eq.1), and the learning parameter $y^{(m)}$ converges to the stationary point of the local learning problem $f^{(m)}(y)$.\n\n**A(W6: Discussion of the convergence results)**: Please refer to the answer for **Q3** in the overall response.\n\n**References**\n\n[1] Personalized Federated Learning with Moreau Envelopes, Dinh et al, NeuRIPS 2020\n\n[2] Personalized Federated Learning: A Meta-Learning Approach, Fallah et al., NeuRIPS 2020\n\n[3] Exploiting Shared Representations for Personalized Federated Learning, Collins et al, ICML 21"}, {"review_id": "Ja4y7fey8D", "reviewer": "Reviewer_G158", "summary": "The authors proposed a two-layer federated learning framework called FedSep, with one layer for communication and another layer for learning. The two layers are connected through decode/encode operations. Furthermore, the authors proposed an efficient algorithm to solve FedSep by treating it as a bilevel optimization problem, and showed convergence results which match those of the standard framework. The FedSep can incorporate Communication-Efficient FL and Heterogeneous-Model FL.", "questions": "(Q) The authors claimed that Theorem 3.7 implies that $\\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 ]||$ converges to zero. Because the right-hand side of the inequality in Theorem 3.7 contains constant, this claim is not straightforward to me, can the authors elaborate?", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "(S1) The proposed framework, FedSep, provided a general framework for resolving the tug-of-war by separating the communication and learning of federated learning into two layers. \n\n(S2) The algorithm was explained in detail and the theoretical results are solid. The experiments results supports the claims of the authors well.", "weaknesses": "(W) The framework rely on the analyticity and strong convexity of the second level problem, which may not be true in many cases. For example, for problem (equation 4) formulated in Sec 4.1, the analyticity and strong convexity does not seem to hold. Although I still believe some convergence results can be derived.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your positive comment. Below are responses to your concerns and questions:\n\n**W: The framework rely on the analyticity and strong convexity of the second level problem.**\n\n**A**: In the bilevel optimization literature, the Non-convex-strongly-convex assumption (including the smoothness assumptions) is widely used. These assumptions ease the analysis of the bilevel optimization problem: strong convexity of the lower problem makes the hyper-gradient well-defined, as an analytical form of the hyper-gradient requires the Hessian to be positive definite.  As for the more general non-convex lower level optimization problem, it is much more challenging to analyze, and we leave the analysis of this case as a future work \n\n\n**Q: how the left-hand side of Theorem 3.7 converges to 0?**\n\n**A**: One approach to achieve the convergence is choosing diminishing learning rate as stated in Corollary 3.8. In our updated version of the manuscript, we will move this discussion after the Corollary 3.8 to avoid confusion."}, {"review_id": "8YumKkcnzv", "reviewer": "Reviewer_A2g4", "summary": "This paper proposes a two-layer federated learning (FL) framework (FedSep) by separating the communication and learning parameters. By a bilevel optimization formulation, FedSep enjoys a convergence guarantee. In addition, two settings, communication-efficient FL and model-heterogeneous FL, are solved in FedSep framework.", "questions": "see above", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. FedSep is an interesting and novel setting.\n2. Theoretical convergence analysis is provided, showing the sublinear rate.\n3. FedSep is applied to communication-efficient FL and model-heterogeneous FL.", "weaknesses": "I have the following concerns and questions:\n\n1. The convergence rate is sublinear, but it does not show any speedup in terms of the clients' number and local steps. So it does not exactly match the convergence of standard FL algorithms. Any possible improvement? Or the statement in the paper should be modified accordingly (e.g., line 75).\n\n2. In the model-heterogeneous FL, the validation data set is also used in the training process both for problem formulation and experiments. I believe this is a concern and unfair for the direct comparison with other algorithms", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your positive comment. Below are responses to your questions:\n\n**Q1: is it possible to show speedup in terms of the clients' number and local steps**\n\n**A**: In Corollary 3.8, we focus on achieving the sub-linear rate. Indeed, it is possible to achieve the linear speedup in terms of the number of clients $M$ and the number of local steps $I$. More specifically, compared to corollary 3.8, we can set $\\gamma$ to be a smaller value: $\\gamma = \\min(\\frac{1}{2L},(\\frac{1}{C_{\\gamma}MIT})^{1/2})$, \nthen we have \na convergence bound of $O(\\frac{1}{(MIT)^{1/2}})$, \nwhich achieves linear speed up w.r.t. $M$ and $I$. Recall that $\\gamma$ denotes the learning rate used by the decode stage (Line 7 -12 in Algorithm~1), smaller $\\gamma$ leads to more decode steps (larger $I_{dec}$). But since we assume the lower optimization problem $g$ be strongly convex (thus linear convergence), the additional cost is negligible.\n\n**Q2 In the model-heterogeneous FL, the validation data set is also used in the training process both for problem formulation and experiments. I believe this is a concern and unfair for the direct comparison with other algorithms.**\n\n**A**: We want to clarify that all methods (both our FedSep and other baselines) are trained over the same training set. While for our FedSep, we use a small subset of the training data to do the sub-net selection. In other words, the $D_{val}$ in Eq.6 is a subset of the training set in experiments. Therefore, it is a fair comparison with other baselines."}]

Author rebuttal:
Thanks all the reviewers for their time and effort. Below we provide responses to questions related to the problem formulation, motivation and interpretation of the convergence theorem:

**Q1: Can you formalize the conflict between communication and learning?**

**A**: Formally, the conflicts between communication and learning can be expressed by the following federated optimization problem with constraints: $\underset{x}{\min } h(x) \coloneqq \frac{1}{M}\sum_{m=1}^M L(x; D^{(m)}), \text{s.t. } \text{Comm}(x) \leq C_1, \text{Cap}(x) \leq C_2$. $L(x; D^{(m)})$ denotes the training objective: we fit a model parameterized with $x$ over the dataset $D^{(m)}$, $L$ denotes a loss objective. For the constraints, $\text{Comm}(\cdot)$ denotes the communication constraint, e.g. the bit-rates transferred between the server and the client is upper bounded, while $\text{Cap}(\cdot)$ denotes the computation capacity constraint, e.g. GPU memory size of clients constrains the largest model that can be trained.  For simplicity, we assume both $\text{Comm}(\cdot)$ and  $\text{Cap}(\cdot)$ only depends on the dimension of the parameter $x$. It is straightforward to identify the conflict between communication and learning from this formulation: Firstly, to satisfy the communication constraint, we need to set $x$ in a low-dimensional space, however,  a high-dimensional parameter is desired for optimizing the objective $L(x; D^{(m)})$, furthermore, the capacity limit $C_2$ is decided by the smallest capacity of all clients, i.e. $C_2 = min (C^{(m)}_2, m\in [M])$, where $C^{(m)}_2$ denotes the capacity constraint of client $m$. As a result, clients with larger capacity are under-utilized. One intuitive approach to mitigate the above mentioned issues is by using different sets of parameters to satisfy the two constraints, which is indeed the communication parameter and the learning parameter proposed in the FedSep framework.

**Q2: Why do we need FedSep if several approaches exists for the communication-efficient FL.**

**A**: Compared to existing works in the literature, our FedSep framework (with a bilevel formulation) **provides a novel and different perspective** to tackle the communication-efficiency challenge in FL. 
More specifically, existing works often view compression operations (such as quantization and sparsification) as introducing errors into the optimization process, as a result, they either employ error feedback techniques for corrections or leverage information theory tools to determine the optimal encoding strategy within communication constraints. In contrast, FedSep distinguishes between communication parameters and learning parameters. Rather than focusing on compression error analysis as existing approaches, we perform convergence analysis through the lens of bilevel optimization.

As shown in the convergence analysis (Theorem 3.7 and Corollary 3.8), we have that the communication parameter $x$ reaches to a stationary point of the bilevel problem (in Eq.1), and the learning parameter $y^{(m)}$ converges to the stationary point of the local learning problem $f^{(m)}(y)$. In the context of communication-efficient FL, we choose the decode function $g^{(m)}$ to be identical across all clients, therefore, the learning parameter $y^{(m)}$ is also identical across clients. Combining with the convergence result, we show that $y \coloneqq y^{(m)}, m \in [M]$ is a stationary point of the classical FL optimization problem $\frac{1}{M}\sum_{m=1}^M f^{(m)}(y)$. To the best of our knowledge, this is the first convergence result for communication parameter (equivalently, the compressed parameter in the literature).

**Q3: How to interpret the convergence result in Theorem 3.7**

**A**: Firstly, **FedSep shows a Coding-Accuracy trade-off**. As shown in Theorem 3.7 (Corollary 3.8), the encode/decode operation affects the convergence in two folds: firstly, $G_2$, the variance of the encode operator, is a coefficient of the $\frac{1}{\sqrt{T}}$ term; $\tilde{G}$, the bias of encode/decode operations, is an additive term. To control the variance term, we can increase the batch-size $b_x$, and to control the bias term, we can increase the number of decode steps $I_{dec}$ and the number of encode steps $Q$. Since the bias error of both decode and encode operations have linear convergence, we can choose small values of $I_{dec}$ and $Q$ in practice. This is consistent with the Accuracy-Communication trade-off in the compression FL literature. In fact, higher communication rate makes the decode minimization problem more ill-posed (e.g. the decode function in Eq.4 in the manuscript), thus needs more decode/encode steps to reach a given error.


Next, our FedSep matches the standard $O(\epsilon^{-2})$ rate as compression FL [1] for the non-convex function. For the encoding/decoding related dependence: [1] has a linear dependence over the compression error, while our FedSep has a linear dependence over the variance term $G_2$ and an additive dependence of the bias term $\tilde{G}$. 

Finally, the convergence rate of FedSep matches that of the federated bilevel optimization method FedNest [2] with a rate of $O(\kappa^5\epsilon^{-2})$. Note that the $\kappa$ is the condition number of the decode function $g^{(m)}$, which is different from the condition number of a strongly-convex objective in classical FL. The convergence bound does not rely on the heterogeneity of client's loss function, as we choose the global learning rate $\eta_g = O(\frac{1}{I})$ to control the error caused by local updates.

**References**

[1] Federated learning with compression: Unified analysis and sharp guarantees, Haddadpour et al., AISTATS 21

[2] Tarzanagh, Davoud Ataee, et al. "Fednest: Federated bilevel, minimax, and compositional optimization." International Conference on Machine Learning. PMLR, 2022.

Meta-review: This work explores a bilevel optimization technique in federated settings that aims to reduce communication costs while maintaining model utility. Reviewers all agreed that the problem formulation was interesting and appreciated the empirical results and convergence guarantees. There were concerns about (1) better motivating/explaining the tension that exists between communication/learning, and (2) more thoroughly explaining how FedSep compares to prior work in personalized and communication-efficient FL. The authors/reviewers had a significant back & forth discussion including discussing some of these ideas in detail; I would recommend the authors carefully incorporate these discussions in their revision.
Decision: Accept (poster)