Differentiated Aggregation to Improve Generalization in Federated Learning
Authors: Peyman Gholami
Hulya Seferoglu
Venue: Accepted by TMLR ()

Abstract: This paper focuses on reducing the communication cost of federated learning by exploring generalization bounds and representation learning. We first characterize a tighter generalization bound for one-round federated learning based on local clients’ generalizations and heterogeneity of data distribution (non-iid scenario). We also characterize a generalization bound in R-round federated learning and its relation to the number of local updates (local stochastic gradient descents (SGDs)). Then, based on our generalization bound analysis and its interpretation through representation learning, we infer that less frequent aggregations for the representation extractor (typically corresponds to initial layers) compared to the head (usually the final layers) leads to the creation of more generalizable models, particularly in non-iid scenarios. We design a novel Federated Learning with Adaptive Local Steps (FedALS) algorithm based on our generalization bound and representation learning analysis. FedALS employs varying aggregation frequencies for different parts of the model, so reduces the communication cost. The paper is followed with experimental results showing the effectiveness of FedALS. Our codes are available for reproducibility.

Published in Transactions on Machine Learning Research (10/2025)
Differentiated Aggregation to Improve Generalization in Federated Learning
Peyman Gholami
pghola24ruic.ed edu
Department of Electrical and Computer Engineering University of Illmois Chicago
Hulya Seferoglu
Faubya&euic... edu
Department of Electrical and Computer Engineering University of Mlinois Chicago
Reviewed on OpenReview: https: // wptp///eepprv.ee...mmee/e. net/ formff id=F5hgp01Ced
Abstract
This paper focuses on reducing the communication cost of federated learning by exploring generalization bounds and representation learning. We first characterize a tighter general- ization bound for one-round federated learning based on local clients' generalizations and heterogeneity of data distribution (non-iid scenario). We also characterize a generalization bound in R-round federated learning and its relation to the number of local updates (local stochastic gradient descents (SGDs)). Then, based on generalization bound analysis and its interpretation through representation learming, we infer that less frequent aggrega- tions for the representation extractor (typically corresponds to initial layers) compared to the head (usually the final layers) leads to the creation of more generalizable models, par- ticularly in non-iid scenarios. We design a novel Federated Learning with Adaptive Local Steps (FedALS) algorithm based on our generalization bound and representation learning analysis. FedALS employs varying aggregation frequencics for different parts of the model, so reduces the communication cost. The paper is followed with experimental results showing the effectiveness of FedALS. Our codes are available for reproducibility.
1 Introduction
Federated lcarning advocates that multiple clients collaboratively train machine learning models under the coordination of parameter server (centrral aggregator) (McMahan el al. 2016) by leveraging all clients' computational capacities. Despite its promise, federated learning suffers from high communication costs between clients and the parameter server. In particular, exchanging machine learning models between clients and the parameter server is costly, especially for large models, which typical in today's machine learning applications (Koneený et al. 2016 Zhang et al.) 2013: Barnes et al. 2020 Braverman et al. 2015 Zibaeirad et al. 2024). Furthermore, the uplink bandwidth of clients may be limited, time- saine-varying and expensive. Thus, there an increasing interest in reducing the communication cost federated learning via (i) multiple local updates, also known as "Local SGD" (Stich) 2018: Stich & Karimireddy 2019 Wang Joshi 2018). (ii) pruning Wangni et al. (2017); Jiang et al. (2022), (iii) quantization Bernstein et al. (2018):
Reisizadeh et al. (2020), etc. In this paper, we take a dramatically different and complementary approach: Aggregating the different parts of machine learning model at different frequencies in federated learning setup.
The primary purpose of communication in federated learning is to periodically aggregate local models to reduce the consensus distance among clients. This practice helps maintain the overall optimization process on a trajectory toward global optimization. It is important note that when the consensus distance among clients becomes substantial, the convergence rate reduces. This occurs as individual clients gradually veer towards their respective local optima. without being synchronized with the models from other clients. This

---

Published in Transactions on Machine Learning Research (10/2025)
issue is amplified when the data distribution among clients is non-iid. It has been demonstrated that the consensus distance is correlated to (i) the randomness in each client's 's own dataset, which causes variation in consecutive local gradients, as well as (ii) the dissimilarity in loss functions among clients due to non- fidness (Stich & Karimiroddy/ 2019 Gholami & Sefcroglu 2024). More specifically, the consensus distance at iteration t is defined as K >>= 00: 0x.ell?, where 0, K 2k=1 Ōk.t; K is the number of clients, Øg.t is the local model at client k at iteration 6, and |1-11² is squared l2 norm. Note that the consensus distance goes to zero when global aggregation is performed at each communication round. This makes the communication of models between clients and the parameter server crucial, but this introduces significant communication overhead. This paper aims to reduce federated learning communication overhead through the following contributions.
Contribution 1: Improved Generalization Error Bound. The generalization error of learning model is de- fined as the difference between the model's empirical risk and population risks. (We provide a mathematical definition in Section 3. Existing approaches for training models mostly minimize the empirical risk or its variants. However, a small population risk is desired, showing how well the model performs in the test phase as it denotes the loss that occurs when new samples are randomly drawn from the distribution. Note that a small empirical risk and a reduced generalization error correspond to a low population risk. Thus, there is an increasing interest in establishing an upper limit for the generalization error and understanding the underlying factors that affect the generalization error. The generalization error analysis is also impor- lant to quantitatively assess the generalization characteristics of trained models, provide reliable guarantees concerming their anticipated performance quality, and design new models and systems.
In this paper, we offer a tighter genoralization bound compared to the state of the art Barnes et al. (022 Yagli et al.] (2020); Sun et al.] (2023) for federated learning, considering local clients generalizations and non-iidness (i.e., heterogeneous data distribution across the clients).
Contribution II: Representation Learning Interpretation. Recent studies have demonstrated that the con- cepl. of representation learning is promising approach to reducing the communication cost of federated learning (Collins et al.) 2021). This is achieved by leveraging the shared representations in all clients' datasots. For example, let us consider a federated learning application for image classification, where differ- ent clients have datasets of different animals. Despite each client having a different dataset (one client has dog images, another has cat images, etc.), these images usually have common feat such as an eye/car shape. These shared features, typically extracted in the same way for different types of animals, require consistent layers of a neural network to extract them, whether the animal is dog a cat. As result, these layers demonstrate similarity (i.e., less variation) across clients even when the datasets are non-iid. This implics that the consensus distance for this part of the model (feature extraction) is likely smaller. Based on these observations, our key idea is to reduce the aggregation frequency of the layers that show high similarity, where these layers are updated locally between consecutive aggregations. This approach would reduce the communication cost of federated learning as some layers are aggregated, hence their parameters are exchanged, less frequently. The next example scratches the surface of the problem for a toy example.
We consider a federated learning setup of five clients with a central parameter server to train a ResNel-20 [He el ul., (2015) on a heterogeneous partition of CIFAR-10 dataset (Krishevsky) (009). We use Federated Averoging (FedAvg) (MeMahan et al.) 2016) as an aggregation algorithm since it is the dominant algorithm in federated learning. We applied FedAng with 50 local steps prior to each averaging step, denoted = 50. Non-üidnsss is introduced by allocoting 2 classes to each client. Finally, we evaluate the quantity of the average consensus dislance for each model layer during the oplimization in Pig. 0. It is cleur that the initial layers have smaller consensus distance as compared to the final layers. This is due to initial layers' role in extracting representations from input data and their higher similority across clients.
Example 1.
The above example indicates that initial layers show higher similarity, so they can be aggregated less fre- quently. Additionally, several empirical studies (Reddi et. al.] 2021 Yu et al., 2020) show that federated learning with multiple local updates per round learns a generalizable representation and is unexpectedly successful in non-iid settings. This view supported by recent thcoretical work showing that as you go deeper into a network, the way it represents information changes in a structured way. For example, Li & SSmpoliisky (2021) and Zavatone- Veth et al. (2022) have found that a layer's representation is blend of
2

---

Published in Transactions on Machine Learning Research (10/2025)
2.00
Feskvg// 3-50 /non
Aw Awerage consensus detance
1.50
1.25
1.00
0.75
0.25
0.0
7.5
20.0
Figure 1: Average consensus distance over time for different layers during ResNet-20 training with FedAvg on CIFAR-10 (5 clients, non-iid, 2 classes per client) shows that early layers have lower consensus distance.
the input data's structure and the learning task's structure. Critically, they showed that the influence of the task becomes stronger in deeper layers, meaning early layers focus more on general input features while later layers adapt to focus on what's needed to solve the task. Also Bordelon & Pehlevan 20222, by doing study on learning dynamics, confirmed this, showing that deeper layers change their representations more significantly during training to ultimately align with the final goal. Taken together, these papers provide a strong thcoretical basis for the idea that networks learn by gradually shifting their focus from representing the input data in early layers to representing the task solution in later layers. These studies encourage us to delve deeper into investigating how local updates and model aggregation frequency for the model's representation extractor and the head affect model generalization.
In this paper, based on our improved generalization bound analysis and our representation learning inter- pretation of this analysis, we showed for the first time that employing different frequencies of aggregation, i.e., the number of local updates (local SGDs), for the representation extractor (typically corresponding to initial layers) and the head (final label prediction layers), leads to the creation of more generalizable models particularly in non-iid scenarios.
Contribution III: Design of FedALS. We design a novel Federated Learning with Adaptive Local Steps (FedALS) algorithm based on our generalization error bound analysis and its representation learning inter- pretation. FedALS employs varying aggregation frequencies for different parts of the model.
Contribution IV: Evaluation. We evaluate the performance of FedALS using deep neural network model
ResNet-20 20 for CIFAR-10, CIFAR-100 (Krizhevsky) 2009). SVHN (Netzer et al. 2011 and MNIST (Lecun ct al.
1998) datasets. We also estimate the impact of FedALS on large langnage models (LLMs) in fine-tuning Zhhang et al., 2022) on the Multi-Genre Natural Language Inference (MultiNLI) corpus (Williams
OPT
25M
el al.
2018). We consider both iid and non-iid data distributions. Experimental results confirm that FedALS outperforms the baselines in terms of accuracy in non-iid setups while also saving on communication costs
across all sctups.
2 Related Work
There has been increasing interest in distributed learning recently, largely driven by Federated Learning. Several studies have highlighted that these algorithms achieve convergence a global optimum or sta- tionary point of the overall objective, particularly in convex non-convex scenarios (Stich & Karimireddy
2019 Stich 2018 Gholami & Seferoglu, 2024 Lian et al. 2018 Kairouz et al. 2019). However, it is widely accepted that communication cost is the major bottleneck for these techniques in large-scale optimization applications (Konecný et al., 2016 Lin et. al.] 2017]. To lackle this issue, two primary strategies are put forth: the utilization of mini-bated parallel SGD, and the adoption of Local SGD. These approaches aim enhance the equilibrium between computation and communication. Woodworth al. (2020b a) attempt theoretically capture the distinction to comprehend under what circumstances Local SGD outperforms minibatch SGD.
3

---

Published in Transactions on Machine Learning Research (10/2025)
Local SGD appears to be more intuitive compared to minibatch SGD, as it ensures progress towards the optimum even in cases where workers are not communicating and employing a mini-batch size that is too large may lead to a decrease in performance (Lin et al. 2017). However, due to the fact that individual gradients for each worker are computed at distinct instances, this technique brings about residual errors. As a result, a compromise arises belween reducing communication rounds and introducing supplementary errors into the gradient estimations. This becomes increasingly significant when data is unevenly distributed across nodes. There are several decentralized algorithms that have been shown to mitigate heterogeneity (Karimireddy et al.] 2019 Liu et al.; ) One prominent example is the Stochastic Controlled Averaging algorithm (SCAFFOLD) (Karimireddy et al. 2019), which addressos the node drift caused by non-iid characteristies of data distribution. They establish the notion that SCAFFOLD demonstrates a convergence rate at least equivalent to SGD, ensuring convergence even when dealing with highly non-iid datasets.
However, despite these factors, multiple investigations (Reddi et al.) 2021 Yu et al., 2020; Lin et al. 2020: Gu et al., 2023), have noted that the model trained using FodAvg and incorporating multiple Local SGD per round exhibits unexpected effectiveness when subsequently fine-tuned for individual clients in non-iid Federated learning setting. This implies that the utilization of FedAvg with several local updates proves effective in acquiring a valuable data representation, which can later be employed on cach node for down- stream tasks. Following this line of reasoning, our justification will be based on the argument that the Local SGD component of FedAvg contributes to improving performance in heterogeneous scenarios by facilitating the acquisition of models with enhanced generalizability.
An essential characteristic of machine learning systems is their capacity to extend their performance to novel and unseen data. This capacity, referred to as generalization, can be expressed within the framework of statistical learning theory. There has been a line of research to characterize generalization bound in FL Wang & Ma (2023); Mohri et al. (2019). More recently Barnes et al. (2022); Sun et al. (2023) Yagli et al. (2020); Seeidggran et al. (2024) considered this problem and gave upper bounds on the expected generalization error for FL in iid sctting in terms of the local goneralizations of clients. These studics demonstrate a к dependency on the number of nodes. Motivated by this, our research focuses on analyzing generalization 1 in iid setting. We then use the derived insights to introduce FedALS, aiming to enhance generalization in non-iid sctup.
that improves the dependency to K²
Our work differentiates itself from personalized federated learning work Collins el al. 2021 in way that our work focuses on a general federated learning setup, where a global model is trained collectively while personalized federated learning advocates training a model from each client's 's perspective. Thus, the results are strikingly different; They show that the heads should be trained locally, while our work FedALS shows that heads should be aggregated more frequently. They also consider only linear models, while our work is generic. [oolins et al. (2022) analyze the behavior of FedAvg in multi-task linear regression with common reprosentation, focusing on the convergence behavior of the algorithm. They prove faster convergence rate when more than one local step is deployed per round. However, they do not address generalization in the context of statistical learning, which is a crucial aspect of our work.
3 Background and Problem Statement
3.1 Preliminaries and Notation
We consider that we have K clients/nodes in our system, and each node has its own portion of the dataset. For example, node i has a local datasct S. = (=k,1,..., Zk.ny where Zk,; = (æt,, yk.:) is drawn from a distribution Dr over X' x y), where X' is the input space and y the label space. We consider X' C Rd and 3 C R. The size of the local dataset at node k is n4- The dataset across all nodes is defined (S (S1,... Sx) J. Data distribution across the nodes could be independent and identically distributed (iid) S = or non-iid. In iid setting, we assume that D2 covers all cases where this equality does not hold.
= Dx = D holds. On the other hand, non-iid setting
We assume that Mg- = A(S) represents the output of a possibly stochastic function denoted as A(S), where Me X + y represents the learned model parameterized by 0. We consider a real-valued loss function denoted 1(Mo, 2); which assesses the model Mo based on a sample 2.
4

---

Published in Transactions on Machine Learning Research (10/2025)
3.2 Generalization Error
We first define an empirical risk on dataset S
1 T2Ak 1=1
=Bk
(Ma, 2k,5);
(1)
where K. is arbitrary distribution over nodes to weight different local risk contributions in the global risk. Specifically, K(() represents the contribution of node k's loss in the global loss. In the most conventional case, it is usually assumed to be uniform across all nodes, i.c., K({k) = к for all k. Rs, (Mo) is the empirical risk for model Ma on local dataset Sk. We further define a population risk for model Ma as
R(Ma) = Ex (Mx) = ERNKLNND (Mo
(2)
where R& (Mo) the population risk on node k's data distribution.
Now, we can define the generalization error for dataset. S and function A(S)
AA R(A(S)) Rs(A(S)).
(3)
The expected generalization error is expressed as LEs ASS)), where Es[[] = E [D , is used for the sake of notation convenience.
3.3 Federated Learning
We consider a federated learning scenario with K nodes/client and a centralized parameter server. The nodes update their localized models to minimize their empirical risk Rs, ( (Mo) on local dataset Ss, while the parameter server aggregates the local models to minimize the empirical risk (MM) Due connectivity and privacy constraints, the clients do not exchange their data with each other. One of the most widely used federated learning algorithms is FedAvg (McMahan et al. 2016), which we explain in detail next.
At round of FedAvg, each node k trains its model Me..- = (Sx) locally using the function/all algorithhm A The local models Me..- are transmitted to the central parameter server, which merges the received local models to aggregated model parameters θ 0,41 A((0 .... 0x,r), where Á is the aggregation function. In FedAvg, the aggregation function calculates an average, so the aggregated model is expressed as
0 ,+1 Ex~x OR
(4)
Subsequently, the aggregated model is transmitted to all nodes. This process continues for R rounds. The final model after R rounds of FedAvg is A(S).
The local models are usually trained using stochastic gradient descent (SGD) at each node. To reduce the communication cost needed between the nodes and the parameter server, each node executes multiple SGD steps using its local data after receiving an aggregated model from the parameter server. To be precise, we have the aggregated model parameters round 0,.- Specifically, upon receiving 0,, node k computes
Ok.r.tt1 = Ox.r, |E,rtt iEBr.r.
1)
Σ
VI(Mo Zk.i)
(5)
for = 0,..., 1, where T is the number of local SGD steps, 0k.r.0 is defined as 0x. θ, is the learning rate, Bk.r, is the batch of samples used in local step of round r in node k, V is the gradient, and |. shows the size of a set. Upon completing the local steps round r, each node transmits Og,r O to the parameter server to calculate 0,41 as in 14)
3.4 Representation Learning
Our approach for analyzing the generalization error bounds for federated learning, specifically focusing on Avd, uses representation learning, which explain next.
5

---

Published in Transactions on Machine Learning Research (10/2025)
We consider class of models that consist of representation extractor (c.g., ResNet). Let 0 be the model Mo's is parameters. We can decompose 0 into two sets: & containing the representation extractor's parameters and h containing the head parameters, i.e., 0 = ( h]. Mg is a function that maps from the original input space to some feature space, i.c., Ms R Rd where usually d' < d. The function Mr performs low complexity mapping from the representation space to the label space, which can be expressed MA : Rd' + R.
For any a E X, the output of the model is M²(x) = (Mr, M4)(x) = Mn (Mo For instance, Me is a neural network, Ma represents several initial layers of the network, which are typically designed extract meaningful representations from the neural network's imput. On the other hand, Mn, denotes the final few layers that lead to the network's output.
4 Improved Generalization Bounds
In this section, we derive generalization bounds for FedAvg based on clients' local generalization performances in general non-iid setting for the first time in the literature. First, start with one-round FedAvg and analyze its generalization bound. Then, extend our analysis to R-round FedAvg.
4.1 One-Round Generalization Bound
In the following theorem, we determine the generalization bound for one round of FedAvg.
Theorem 4.1. Let z) be p-strongly convex and L-smooth in Ma- Me, = As(Ss)) represents the model obtained from Empirical Risk Minimization (ERM) algorithm on local dataset Sk, i.c., Mo. = arg minM E [(M, zk.), and Ma A(S) is the model after one round of FedAvg (0 = Renx On). Then, the expected generulization error, Es (SS),, is upper bounded with
L
+2 EXCC)
Ex-K
Es. Au (St) Expected local generolizaation
Es (S) Expeected won-iidness Expected local generulization
Es. A (Sa)
(6)
where ( Rs, (A(S)) Rs, (AR(Sk) shows the level of non-iidness at client for function A on dataset S.
Proof: The proof of Theorem 4.1 in Appendix
Remark 4.2. We note that this theorem and its proof assume that all clients participate in learning. The other scenario is that not all clients participate in the learning procedure. We can consider the following two cases when not all clients participate in the learning procedure.
Case 1: Sampling & clients with replacement based on distribution K, followed by averaging the local models with equal weights
Case II: Sampling & clients without replacement uniformly at random, then performing weighted averaging of local models. Here, the weight of client rescaled K(A)
K
The generalization error results in these cases are affected by substituting & and K(A)
instcad of KC(k)
for cases and II, respectively. The detailed proof is provided in Appendix R
Discussion. Note that there are two terms in the generalization error bound: (i) local generalization of each client that shows
generalizable local models lead to a better generalization of the aggregated model, (ii) non-iidness of each client which deteriorates generalization. Theorem 4.1 reveals factor of IC())² for the first term, which is the sole (erm in the iid setting. For example, in the uniform case (KC()) = )) we will observe an improvement with a factor of K2 for the iid case. This represents an enhancement compared to the state of the art [Barnes al. 2022 Sun et al. 2023 Yagli et al. 2020), which only demonstrates a factor of + As a result, after the averaging process carried out by the central parameter server, the generalization error is reduced by factor of k² in iid case. This is interesting because one would normally expect an improvement

---

Published in Transactions on Machine Learning Research (10/2025)
of K based the linear increase in the collective dataset, but this shows additional improvement of as well.
On the other hand, we do not see similar behavior in non-iid case. In other words, the expected general- ization error bound does not necessarily decrease with averaging. These results show why FedAvg works well in iid setup but not necessarily non-iid setup. This observation motivates us to design a new federated learning approach for non-iid setup. The question in this context what the new federated learning design should be. To answer this question, analyze R-round generalization bound in the next section.
4.2 Round Generalization Bound
In this setup, after R rounds, there is a sequence of weights {6, [ and the final model is 2 We consider that at round r, each node constructs its updated model as in (5) by taking 7 gradient steps starting from
O, with respect to T random mini-batcher Zu.r = Ui drawn from the local dataset Sk. For this type of iterative algorithm, we consider the following averaged empirical risk
R
E
Σ
1(Ma, =n.))
(7)
R
Zkarl
ieZz.r
The corresponding generalization error, A FedAAg (S), is
R
E BEN BaND,(MM, 2) 1
EiMMo zx.i)
(8)
R
Note that the expression in differs slightly from the end-to-end generalization error that would be obtained by considering the final model Møg and the entire dataset S. More specifically, (S) is an average of the generalization errors measured at each round, similar to Barnes al (2022 ). We anticipate that the generalization error diminishes with the increasing number of data samples, so this generalization error definition yields a more cautious upper limit and serves as a sensible measure. The next theorem characterizes the expected generalization error bounds for R-Round FedAvg in iid and non-iid settings.
Thcorem 4.3. Let I(Mo 2) be p-stroogly conver and L-smooth in Mg. Local models at round are cal- culated by doing т local gradient descent steps 5 and the local gradient variance is bounded by o², i.c., Ez-D. [[V[ x - Ez~Dz VI(Me, x]]2 - 0². The aggregated model at round r, Mo., is obtained by per- forming FedAug, and the data points used in round r (L.c., Zurr are sampled without replacement. Then the average generalization error, E, A FeadAg (S), is upper bounded by
(2LK((k)2
8L ((K)(AB) A +
R
Ex-K
V P
(9)
where,
A A=ō (v (((M)
o²
+
NT
B = Ō Eignn.
S0.A Zz 1K
μT 2
Ô hides constants and
factors, and C(Me) shows the complerity of the model class of Me.
Proof: The proof of Theorem is in Appendix C
The generalization error bound in 19 depends on the following parameters: (i) number of rounds; R, (ii)
number of samples used in every round; [Zr.rl, (iii) the complexity of the model class; C(Me), non-iidness; br.A(1Za, number of local steps in each round; т. We note that (g) also depends on K (more
specifically XC. but this dependence is similar to the discussion we had for one-round generalization, so we skip it here.
7

---

Published in Transactions on Machine Learning Research (10/2025)
Algorithm FedALS
Input: Initial model {0x,0.0 = [04.0.0. hx.0,,l) learning rate 7,, number of local steps for the head adaptation coefficient a.
1: for Round in O,.... R 1 do
2:
for Node k in I,.., K in parallel do
3: for Local step in O,... 10-110 do
4: Sample the batch RR, from Dx-
5: 0k.r.t+1 = 0krr. Be EieBe... Vie kk,s) if mod (rT +t,)) = then
7
Send he, lo the server.
Receive the aggregated head: hR. end if
Di=1 hgr
K
10: if
mod (rT + t, ar) then
11:
Send 6k,r, to the server.
Receive the aggregated Representation extractor: Øk,r, t X EK Pk. 13: end if
12:
14: end for
15: 0 = RR.r.T
16: end for
17: end for
18: return OR = Σk=1 Og,R,0
5 FedALS: Federated Learning with Adaptive Local Steps
The number of samples and local steps in each round are crucial for minimizing the generalization error bound, particularly in non-iid scenarios, as described in (9 where the error bound is looser compared to the iid setup. Our key insight in this paper is that reducing the aggregation frequency, which increases both T and 1 can lead to a smaller generalization error bound according to 9
Increasing I also increases the non-iidness among nodes, defined as:
(n..) Rz (Mn Mτ) Rz (May Mœx),
where Mn o MR represents the aggregated model A((ZR, ())), and Ma, o Mæn represents the local model As (Z This happens because the model moves closer to the local optimum of node k's loss, thereby reducing Rz,,, (Ma, Mş.)-
In the context of representation learning, discussed in the Introduction, the representation extractors M and Mrr are much more similar compared to the task-specific heads M,, and Mr. Therefore, increasing the number local steps for the representation extractor does not significantly increase non-iidness, unlike increasing local steps for the head.
The main idea behind FedALS steps while maintaining low level of non-iidness. This can be achieved if TM, is set larger than TM₁ where TMa and TMp, the corresponding number of local iterations for Ms and Mn, respectively. FedALS in Algorithm 1 divides the model into two parts: (i) the representation extractor, denoted as Mş, and (fi) the head, denoted as Mn. Additionally, we introduce the parameter TM an adaptation coefficient, which can be regarded as hyperparameter for estimating the true ratio.
improve generalization of the model by increasing the number of local
6 Experimental Results
In this section, we assess the performance of FedALS using ResNet-20 as a deep neural network architecture and OPT-125M as large language model. We treat the convolutional layers of ResNet-20 as the repre- sentation extractor and the final dense layers the model head. For OPT-125M, consider the first. 10
8

---

Published in Transactions on Machine Learning Research (10/2025)
Frided
24
2.
PodiL SCAPTOLS
Sentor 107
(a) non-iid
(b) iid
Figure 2: Training ResNet-20 20 on SVHN.
1.4
1.3
512
Rond 0.5
(a) non-iid
(b) iid
Figure 3: Fine-luning OPT-125M on MultiNLI.
layers of the model as the representation extractor. We used the datasets CIFAR-10, CIFAR-100, SVHN, and MNIST for image classification and the Multi-Genre Natural Language Inference (MultiNLI) corpus conducted a network consisting of five nodes alongside central
for the LLM. The experimentation server. For image classification, we utilized batch size of 64 per node. SGD with momentum was employed as the optimizer, with the momentum set to 0.9, and the weight decay to 10 For the LLM fine-tuning, we employed a batch size of 16 sentences from the corpus, and the optimizer used was AdamW. In all the experiments, to perform a grid search for the learning rate, we conducted cach experiment by multiplying and dividing the learning rate by powers of two, slopping each experiment after reaching a local optimum learning rate. We repeat each experiment 20 times and present the error bars associated with the randomness of the optimization. In every figure, include the average and standard deviation crror bars. Detailed experimental setup is provided in Appendix E of the supplementary materials.
6.1 FedALS in non-iid Setting
In this section, we allocate the dataset to nodes using non-iid approach. For image classification, we initially sorted the data based on their labels and subsequently divided it among nodes following this sorted sequence. In MultiNLI, we sorted the sentences based on their genre.
In this scenario, we can observe in Fig. 2a 3a the anticipated performance improvement through the incorporation of different local steps across the model. By utilizing parameters T = 5 and a 10 in FedALS, it becomes apparent that aggregation and communication costs are reduced compared to FedAvg with the same 7 value of 5, This implies that the initial layers perform aggregation at every 50 iterations. This reduction in the number of communications is accompanied by enhanced model generalization stemming from the larger number of local steps in the initial layers, which contribules to an overall performance enhancement. Thus, our approach in FedALS is beneficial for both communication efficiency and enhancing model generalization performance simultaneously.
6.2 FedALS in iid Setting
The results for the iid setting are presented in Fig. 2b 3b In order obtain these results, the data is shuffled, and then evenly divided among nodes. We note that in this situation, the performance improvement
9

---

Published in Transactions on Machine Learning Research (10/2025)
Table 1: Test Performance for 5 nodes FL; Accuracy after training ResNet-20 and test loss after Sane-tunig OPT-125M in iid and non-iid settings with and 10.
FedAvg
FedALS
SCAFFOLD FedALS SCAFFOLD
Model/ Dataast
I1D
Non-IID
IID
Nou-IID
Non-IID
Non-IID
ResNed-20/SVHH
0.948 + 0.002
0.701 0.033
0.954 0.002
0.812 + 0.021
0.718
0.040
0.810 + 0.027
0.876 + 0.009
0.465
0.007
0.522 0.002
0.451 0.037
0.071
0.512 0.010
0.600 + 0.007
0.418
0.014
0.612 0.005
0.486 + 0.022
0
0.482 + 0.013 0.011
ResNet-22/ ///NISTT
0.991 + 0.000
0.797 + 0.064
0.991 + 0.000
0.821 + 0.036
0.812 0.055
0.783 + 0.118
OPT-125M/Multiie
1.250 + 0.000
1.284
0.002
1.248 = 0.001
1.277 + 0.002
1.284
1.278 + 0.002
of FedALS is negligible. This is expected since there is a factor of K in the generalization in this case, ensuring that we will have nearly the same population risk as the empirical risk.
6.3 Compared to and Complementing SCAFFOLD
Karimiroddy et al. (2099 introduced an innovative technique called SCAFFOLD, which employs some control variables for variance reduction address the issue of "client-drift" in local updates. This drift. happens when data is heterogeneous (non-iid), causing individual nodes /cleess converge towards their local optima rather than the global optima. While this approach is significant thcorctical advancement in achieving independence from loss function disparities among nodes, it hinges the assumption of smoothness in the loss functions, which might not hold true for practical deep learning problems in the real world. Additionally, since SCAFFOLD requires the transmission of control variables to the central server, which is of the same size as the models themselves, results in approximately twice the communication overhcad when compared FedAvg.
Let consider Fig. 2a 3a to notice that in real-world deep learning situations, FedALS enhances per- formance significantly, while SCAFFOLD exhibits slight improvements in specific scenarios. Moreover, we integrated FedATS and SCAFFOLD to concurrently leverage both approaches. The results of the test accuracy in different datasets are summarized in Table
6.4 Complementing SCAFFOLD
Karimireddy et al. 2019 introduced an innovative lechnique called SCAFFOLD, which employs some control variables for variance reduction to address the issue of "client-drift" in local updates. This drift happens when data is heterogeneous (non-iid), causing individual nodes/cliente to converge towards their local optima rather than the global optima. Methods like SCAFFOLD are complementary to FedALS. They can be applied concurrently to gain the dual benefits of improved communication efficiency from FedALS while further combating the effects of non-iid data. SCAFFOLD incurs approximately double the communication overhead of FedAvg because it requires sending control variables that are equal in size to the model parameters. However, this overhead can be mitigated by applying FedALS, which improves SCAFFOLDD's communication efficiency.
Let consider Fig. 2a 3a to notice that in real-world deep learning situations, FedALS enhances per- formance significantly, while SCAFFOLD exhibits slight improvements in specific scenarios. Moreover, we integrated FedALS and SCAFFOLD to concurrently leverage both approaches. The results of the test accuracy in different datasets are summarized in Table
6.5 The Role of a and Communication Overhead
As shown in Table 2 it becomes evident that when we increase a from (FedAvg), we initially witness an enhancement accuracy owing to improved generalization. However, beyond a certain threshold (a 10), further increment in ceases to contribute to performance improvement. This is due to the adverse impact of a high number of local steps the non-iidness. The trade-off we discussed in the earlier sections is evident in this context. We have also demonstrated the impact of FedALS on the communication overhead in this table.
10

---

Published in Transactions on Machine Learning Research (10/2025)
Table 2: The accuracy and communication overhead per client after training ResNet-20 in non-iid setting
with r 5 and variable a.
DATASET
VALUE OF
# OF COMMUNICATED
SVHN
CIFAR-10
PARAMETERS
0,7010 + 0.0330
0.4651 + 0.0071
2.344B
5
0.8107 = 0.0278
0.5201 + 0.0302
0.473B
10
0.8117 + 0.0214
0.5224 + 0.0365
0.239
25
0.7201 = 0.0549
0.3814 + 0.0641
0.099
50
0.6377 + 0.0520
0.2853 + 0.0641
0.052
100
0.5837 + 0.0715
0.2817 + 0.032
0.029
Table 3: Different Combinations of Φ, h for training ResNet-20 in non-iid setting with 5, 10.
DATASET
VALUE OF
SVHN
CIFAR-10
CIFAR-100
20
0.6991 + 0.0160
0.4383 + 0.0423
0.4781 = 0.0123
16
0.7112 0.0471
0.4687
0.0111
0.4782 + 0.0087
12
0.6760 + 0.0474
0.4125 + 0.0283
0.4249 = 0.0143
8
0.6381 + 0.0428
0.3779 = 0.03451
0.4085 + 0.0094
4
0.6339 + 0.0446
0.3730 + 0.0310
0.4183 = 0.0108
0.6058 0.0197
0.4013 0.0308
0.3880 + 0.0305
6.6 Different Combinations of h,
In Table 3 have presented the results of our experiments, illustrating how different combinations of & and h influence the model performance in FedALS. The parameter L indicates the number of layers in the model considered as the representation extractor (ø), while the remaining layers are considered as h. We observe that for ResNet-20, choosing & to be the first. 16 layers and performing less aggregation for them seems to be the most effective option.
7 Conclusion
In this paper, we characterized the generalization error bound for one- and R-round federated learning, showing that the one-round bound is lighter than the current state of the art Our analysis, combined with a representation learning perspective, revealed that less frequent aggregations, resulting in more local updates for the initial layers, lead more generalizable models, espocially non-iid scenarios. This insight inspired the development of the FedALS algorithm, which increases local steps for the initial layers while performing more averaging for the final layers. Experimental results demonstrated FedALS's effectiveness in heterogeneous setups.
Acknowledgments
This work was supported in parts by the Army Research Lab (W911NF2420172), and the National Science Foundation (CCF-1942878, CNS-2148182, CNS-2112471).
References
Leighton Pate Barnes, Yanjun Han, and Ayfer Ozgur. Lower bounds for learning distributions under com- munication constraints via fisher information. Journal of Machine Learning Research, 21 (236):1-3) 2020.
11

---

Published in Transactions on Machine Learning Research (10/2025)
URL http://ppp..orr///pppps/////////
Leighton Pate Barnes, Alex Dytso, and H. Vincent Poor. Improved information-thooretic generaliza- tion bounds for distributed, federated, and iterative learning. Enlropy, 24, 2022. URL https://app.
semanticscholar orgg/oorpusD::246333555
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Aziwadenesheli, and Animashree Anandkumar. signSGD: Compressed optimisation for non-convex problems. In Jennifer Dy and Andreas Krause (eds.), Proceed- ings of the 35th International Conference on Machine Learning. volume 80 of Proceedings of Machine Learning Research, pp. 560-569 PMLR, 10-15 Jul 2018. URL https://proceeddingg mlr paa.r.rressv/00 bernstein1@a.htmr
Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks, 2022. URL https: http://xxvvv...gaaba/ wtgg////////.....// 09653
Mark Braverman, Ankit Garg, Tengyu Ma, Huy L. Nguyen, and David P. Woodruff. Communication lower bounds for statistical estimation problems via distributed data processing inequality. Pro-
ceedings of the forty-eighth annual ACM symposium on Theory of Computing, 2015. URL https:
7/api. //api.semantao www.stiismmttiiitne...e. www.eerrroooooo wrg/CorppsID:1172222
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations
for personalized federated learning. In International Conference on Machine Learning. 2021. URL https:
//pp. www.smmiccccchlarr...... htt.://ar....rooorruaaaa02111120
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Fedavg with fine tuning: Local
updates lead to representation learning, 2022. URL htttps://axxii. http://rxxw...o//bbb//b/aaaa 13692
Peyman Gholami and Hulya Seferoglu. Digest: Fast and communication efficient decentralized learning with local updates. IEEE Transactions on Machine Learning in Communications and Networking, pp. 1 1, 2024. doi: 10.1109 10.1109/TMLCN.... TTMCC.
Xinran Gu, Kaifeng Lyu, Longbo Huang, and Sanjeev Arora. Why (and when) does local SGD generalize better than SGD? In The Eleventh International Conference on Learning Representations, 2023. URL
https://openreviewt htt.://wwww....e.oooo////stttaaa
Kaiming He, X. Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. 2016
IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2015.
Yuang Jiang. Shiqiang Wang, Victor Valls, Bong Jun Ko, Wai-Han Lee, Kin K. Leung, and Leandros Tassiulas. Model pruming enables elficient federated learning on edge devices, 2022. URL https://
arxiv.org wttp:orr//a/a////ooeo 19...999.229 12326
Peter Kairouz, H. B. McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary B. Charles, Graham Cormode, Rachel Cummings, Rafael G. L. D'Oliveira, Salim Y. El Rouayheb, David Evans, Josh Gardner, Zachary Garrett, Adrià Gascón, Badih Ghazi, Phillip B. Gibbons, Marco Gruteser, Zaid Harchaoui, Chaoyang He, Lie He, Zhouyuan Huo, Ben Hutchinson, Justin Hsu, Martin Jaggi, Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Konecný, Aleksandra Korolova, Farinaz Koushanfar, Oluwasanmi Koyejo, Tancrède Lepoint, Yang Liu, Prateek Mittal, Mehryar Mohri, Richard Nock, Ayfer Özgür, R. Pagh, Mariana Raykova, Hang Qi, Daniel Ramage, Ramesh Raskar, Dawn Xiaodong Song, Weikang Song, Sebastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramèr, Pranoeth Vepakomma, Jianyu Wang, Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and open problems in federated learning. Found. Trends Mach. Learn., 14:1 210, 2019. URL org Pttg/CorpusID:.. :ttp//CoopuDID:22222666
atctp::/aai..seaatiicce
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J. Reddi, Sebastian U. Stich, and Ananda Thcertha. Suresh. Scaffold: Stochastic controlled averaging for foderated learning. In Inter- national Conference on Machine Learning, 2019. URL https::// wwwanticschola.tt...omoeaaaaaaaa wttpr//sscccooaooooo//orrpaaaaaa org/CorpusID:
214069261
12

---

Published in Transactions on Machine Learning Research (10/2025)
Jakub Konecný, H. B. McMahan, Felix X. Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. ArXin, abs/1610.05492, 2016. URL
Stapp:://ap..smmatiicss wtw.:oaaaa.rrroooou....o. 14999259
Alex Krizhevsky, Learning multiple layers of features from tiny images. 2009.
Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-basee learning applied to document recognition. Proceedings of the IEEE, 86(11)22278 2324, 1998. doi: 10.1109/5.7200 726799..
Qianyi Li and Haim Sompolinsky. Statistical mechanics of deep linear neural networks: The backpropagating kernel renormalization. Physical Review X, 11(3), September 2021. ISSN 2160-3308. doi: 10.1103/ physrevx. 11 11.03100.. ..30099. URL http://dx doi org/ 10. 1103/PhysRev/ 11 0110599
Xiangru Lian, Wei Zhang, Ce Zhang, and Ji Liu. Asynchronous decentralized parallel stochastic gradient descent. In Jennifer G. Dy and Andreas Krause (eds.), ICML, volume 80 of Proceedings of Machine Learning Research, pp. 3049 3058. PMLR, 2018
Tao Lin, Sebastian U. Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't use large mini-batches, use local sgd. In International Conference on Learning Representations, 2020. URL httpss://peereevies net/ Borym
Yujun Lin, Song Han, Huizi Mao, Yu Wang, and William J. Dally. Deep gradient compression: Reducing the communication bandwidth for distributed training. ArXiv, abs/1712.01887, 2017.
semanticscholar. wtwpiiiccccoaaaooo/ooooooooo org//oorpusD:33799222
Yue Liu, Tao Lin, Anastasia Koloskova, and Sebastian U. Stich. Decentralized gradient tracking with local steps. ArXiv, abs/2301.01313, 2023. URL https://api set.atttesssblll.r....///roooapa
H. B. McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera Arcas. Communication- efficient learning of deep networks from decentralized data. In International Conference on Artificial Intelligence and Statistics, 2016. URL
14955348
Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. CoRR,
abs/1902.00116, 2019. URL
Yuval Netzer, Tao Wang, Adam Coates, A. Bissacco, Bo Wu, and A. Ng. Reading digits in natural im- ages with unsupervised feature learning. 2011. URL https://api semanticacholar org/ CorpusID: 16852518
Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Representations, 2021. URL http:///peenereeeeeeeeeeoorrraea
Amirhosscin Reisizadch, Aryan Mokhtari, Hamed Hassani, Ali Jadbabaic, and Ramtin Pedarsani. Fedpaq: A communication--efiiieee federated learning method with periodic averaging and quantization. In Silvia Chiappa and Roberto Calandra (eds.), Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pp. 2021- 2031. PMLR, 26 28 Aug 2020. URL httpa///poooeddngggg.ppeeee/seee mlr http:///poccoooogoaaaa.eeeseeeer
Milad Sefidgaran, Romain Chor, Abdellatif Zaidi, and Yijun Wan. Lessons from generalization error analysis of federated learning: You may communicate less often!, 2024. URI. https://arxiw.orr//bb/
05862
Sebastian U. Stich. Local sgd converges fast and communicates little. ArXíe, abs/1805.09767, 2018. URL http::://pp...sematticcccchhhbb
httpaa/acisccc.aooooooooooooiiaa wrg./orppsII:: :4:0:pppsuI::3944111 43964415
Scbastian U. Stich and Sai Pranceth Karimireddy. The error-foedback framework: Better rates for sgd
with delayed gradients and compressed communication. ArXiv, abs/1999.05350, 2019. URL https:
7/api. www.wsecccccccaaa.ccoo.m wtt.:/cccc......oooooopuuaaaaaad
13

---

Published in Transactions on Machine Learning Research (10/2025)
Zhenyu Sun, Xiaochun Niu, and Ermin Wei. Understanding generalization of federated learning via stability: Helerogeneity matters, 2023.
Leslie G. Valiant. theory of the learnable. Commun. ACM, 27:1134 1142, 1984. URL https petpss//apti //api.
semanticscholar Ctapolaarrrrr/ooopuaueee Drg/CorpusID:59712
Jianyu Wang and Gauri Joshi. Cooperative sgd: A unified framework for the design and analy-
sis of communicationn-efiieel sgd algorithms. ArXio, Jbs/1808.07576, 2018. URL https://api.
semanticscholar. hwgpe/ohoarroooo/oooooooooo Peap//CopisD:200000 2200 52078222
Mingze Wang and Chao Ma. Generalization error bounds for deep neural networks trained by sgd, 2023.
Jianqiao Wangni, Jialei Wang, Ji Liu, and Tong Zhang. Gradient sparsification for commmicationn-eliiient 09854.
distributed optimization, 2017. URL http:///axwvvoo../.e./e
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume 1 (Long Papers). pp. 1112 1122. Association for Computational Linguistics, 2018. URL http:://cllwee org/
1101
Blake E. Woodworth, Kumar Kshitij Patel, and Nathan Srebro. Minibatch vs local sgd for heterogencous dis- tributed learning. ArXiv, 0os/22006004335 2020a. URL http::://ppi.senaaaiiiceee wtt.:/ooollo..roor/orsrss/lt
org/
219531445
Blake E. Woodworth, Kumar Kshitij Patel, Sebastian U. Stich, Zhen Dai, Brian Bullins, H. B. McMahan, Ohad Shamir, and Nathan Srebro. Is local sgd better than minibatch sgd? ArXiv, abs/2002.07883, 2020b. URL https://apl swmanticscoolarr httpr//ccccccc.ooooooooooossioia 2111111488
Semih Yagli, Alex Dytso, and H. Vincent Poor. -hlormation-theoretiic bounds on the generalization error and privacy leakage federated learning. In 2020 IEEE 21st International Workshop on Signal Processing Ad- vances in Wireless Communications (SPAWC), pp. 1 5, 2020. doi:
Tao Yu, Eugene Bagdasaryan, and Vitaly Shmatikov. Salvaging federated learning by local adaptation. semanicschoollr
ArXis, 20s/2000.04758, 2020. URL http:s//ppp...staatttcoooooooora org/CorrpusD:22108888
Jacob Zavalone- Veth, Abdulkadir Canalar, Benjamin Ruben, and Cengiz Pehlevan. Asymplotics of representation learning in finite bayesian neural networks*. Journal of Statistical Mechanics: Theory and
Experiment, 2022(11):11400g) November 2022. ISSN 1742-5468. doi: 10.1088/1742-544
URL
http://áx.dow doi. wttp://0.118//4444466////3333330
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. Opt: Open pre-trained transformer language models. ArXis, Abs/2205.0106x 2022. URL https::/apt. swmanticscholar........ org/
Corpus.DD:288962222
Yuchen Zhang, John C. Duchi, Michael I. Jordan, and Martin J. Wainwright, Information-theorre lower bounds for distributed statistical estimation with communication constraints. In NIPS, 2013. URL https:
//api. semanticccocoll.. wwtp:ooopo//:. 11929529
Arastoo Zibacirad, Farnoosh Koleini, Shengping Bi, Tao Hou, and Tao Wang. A comprehensive survey on the security of smart grid: Challenges, mitigations, and future research opportunities, 2024. URL
https: //tr:v/.//ivvvvv. httr:::/arxxv.oor//bbb/// wwtp:oo//////////////
14

---

Published in Transactions on Machine Learning Research (10/2025)
A Appendix
B Proof of Theorem 4.1
We first state and prove the following lemma that will be used the proof of Theorem 4.1
Lemma B.1 (Leavo-onc-out). [Expansion of Theorem in Barnes et al. 2022
Let S (2 =knn), where zkii is sampled from Dg. Denote S (S,,.., Sk: k..,.. Sk).
Then
E BIS AA(S) BK
I(A(S), -
(10)
1=1
Proof. We have
E [S R(A(S)) E (S. ((A(S),
(11)
Also, observe that
E B(S Rs(A(5)) = R 1 (S
((A(5), 2n21
(12)
Thk
1=1
Eg. (S.
((A(S()),, =/.1)
(13)
)am
Putting 11 and 13 together, and by the definition of the expected generalization error, we get
E DS As(S) = (S..K((5..5--")))
- -(
(14)
In the following lemma, we establish a fundamental generalization bound for a single round of ERM and FedAvg. (Theorem 4.1).
Theorem B.2. Let I(Mo, 2) be p-strongly conver and L-smooth in Me. Me. = As(ss) represents the model obtained from Empirical Risk Minimization (ERM) algorithm on local dataset Sk. i.e., Me. = arg min D 1(M, Zk.;), and M3 = A(S) is the model after one round of FedAug (0 = E 0x). Then, the expected generulization error, Es D AA(S), is bounded by
T P
E BISEDDT AA (Sx) +2 T K(K)
E
E EISEND HK (S)
E
BISENDEY AA, (Sk)
Expected local generalization
Root of expected non-sidness Root of expected local generalisation
(15)
where (5)) on dataset S.
Rs. Rs.. (A(S)) R.s. (Au())
indicates the level of non-tidness for client in function A
Proof. We again consider S% 1.... 2k where zh, sampled from DR. Let also define (kk)
(Si, Sk...., Sk). Based on Lemma B.1 we can express the expected generalization error
((A(S), -((A))))). =4x)))
E BS As(S) (A..,
(16)
n s=1
15

---

Published in Transactions on Machine Learning Research (10/2025)
Based on L-smoothness of (Mo, 2) in Me, we obtain
((,) - -(A(((t))) zke)))
(V 1 Σ
+ - (|A(5)) - A((
(17)
where 1--1 indicate Euclidean inner product, and squared L2-norm. Note that 17
holds due to
f(y) S f(x) + (Vf()),, y - x) +
2 ly - =
(18)
We can bound expectation of the inner product term on the right-hand side of inequality as
17 using Cauchy Schwarz
nh
E (S..
V n
((A(St)).
2(,)) A(S!A)))
< BRAK B (SS 1221
1VV
(A(SSM)) =.,), A(S
(19)
Thk
< BBn B( Huml
( 2,D1 [[A(S) - A(S(
(20)
< BREx BiS..S(-D
IV
(A(SS))) = 118)) ||A(S) A(S()))12
(21)
Tkk
where 19 is true because on the right we have the absolute value.
20
and
21
are based on
Schwarz incquality.
- (M())-((S(H))) that appears on the right-hand side Now Let's find an upper bound for S( of both
17 and
21 We obtain
E ((
|A(S) A(S
= -E( 1._
(((K) - As(S111
(22)
2K((k)2 P
< (S
Rs (Ar(5)) - Rs; (Ac(S)))
(23)
2A((k222 11
mE
([AA(())) =l.) ((AAc(5)), 2h.))
D
(24)
Thk
2(((k) 11
A (S))
(25)
2K((k)²
(Sk),
(26)
1
where 22 proceeds by observing that A(S(k.0) varies solely in the sub-model derived from node k, diverging from A(S), and this discrepancy is magnified by factor of K(k)) when averaging of all sub-models. (23)
holds due to the u-strongly convexity of I(Mo. 2) in Mo which leads to p-strongly convexity of Rs, (Mo)
and the fact that As(S)) is derived from the local FRM, ie.., As(S)) = arg min
(Ees 1(M, 2k,) and
v Rsg (Ax (S))) Note that if f is p-strongly convex, we get
f(x) f(g) + - x vl < (V)()), -9)-
(27)
16

---

Published in Transactions on Machine Learning Research (10/2025)
24
25
based on local empirical and population risk definitions.
Now we bound (S..
V E
((A(S(R),, =/, on the right- side of 21.
Note that
|V
Σ ((A(S(M)
21 2k
22 2L
1 nk
((A(S()) ),
1
Σ
nK
(AA =(,))
(28)
n
TE
==1
- 2L
Rs; (A((()))) Rs; (AA (5x)))
(29)
where 28 is obtained using the fact that for any L-smooth function f, we have
|V/(()]I & 21( f²),
(30)
and the fact that As())) is derived from the local ERM, i.e., As())) arg min > ((M, xx,)) based on the definition local empirical risk.
29 is
Putting into 166 and considering (1 we get
E (S k (A(S)
L (|A(S)) - A(5
S BA-K
E D( 1im
(31)
+
B (
(A(S())) =1,)11 ( ||A(S) - A(S
LKCKK) (Sx)
< BK-x Es Pmm
(32)
2/C(A:)2
+V B S.. -D' 1mm
2L (A(((1))) -Rs; (Aa(5))) E BISEND
A (Sx)
H
LIC(K)2 (Sk)
$ BAnK
Bs 18
(33)
+V E BSSNDD =mm 2L (A(()) Rs. (An(Sw)) EISENDRN lkms
21((A:)2
A (Sk)
TK((K)2
To
< &BAnK
E
(Sk) +2 / BS)
AA, (Sx)
1
(34)
(35)
where in
322
have applied (26) and 29
33
proceeds by considering that
E s
S D --
Rs; Rs; (A(S()))) Rs; (Aa())))
E -E
Rs. (A(S)) Rs. (Ax()))
(36)
In 334 we have used the definition (S) = Rs. (A(S)) - Rsa (Ax())) and provides the following upper bound for 1E (SSU-DN)) Aa(S),
This completes the proof
E (K)
L / /R (S) BSSDDDR AA. (Sx)
E
AA. (Sx) 2/
(37)
11
17

---

Published in Transactions on Machine Learning Research (10/2025)
C Proof of Theorem 4.3
Here we provide an identical thcorem as Theorem B.2 except that instead of ERM, multiple local stochastic gradient descent steps are used as the local optimizer.
Theorem C.1. Let 1(Mo z) be u-strongly conver and L-smooth in Me. Mo, A (Sk) represents the model obtoined by doing multiple local steps as in 5 on local dalaset Sk, and Mg A(S) is the model after one round of FedAug (0 = E xk). Then, the expected generalization error, I9 (sSSWD"*1 1e AA(S), is bounded
2LK(K)2
+ (5x))))
E
E B 1.
(38)
1
8L x(())) ( (6u.(1))) + c((ss) E B(SeND sem
(AL((S)) + 2cx(5k))
+
where (x() (Sx) - Rs. (Sx).
Proof. All the steps are exactly the same as in the proof of theorem B.2 except for the two steps below:
First, we use a new upper bound for DE [A(S) - ((S that appears on the right-hand
side of both 17 and 21
We have
E (S. (A(S) A(S
E [S. )5_ ((k)) An
(39)
= (SS..),))) W k(A)) (SA - (S)) + ( - (s2211
(40)
=E -E DD "h""" v 2 2 (IA)(S)) - A'(S) ((A(S)) - AA((5))
(41)
4K((K)²
S E DD 1k
(AL( Rs, (S)) + RS: (Ax( - RS; (A
(42)
4K((k)2 P 4K((k)2
= (S D... yem
Rs; (At(S))) Rs₁ (Ax (S.)) + 2255 (At(S))) 2Rs; (AA
(43)
Th
((Ax =1,)) + +RRs (Ax(S)) -2Rs (A*
Tkk j=1
(44)
4K(K)2
(AA( 22 (Ax( -25; (Ar(()))
SS HE
(45)
AKC(1:)2
= EES pe
+2 (S,) + 2Rs. (Sx) - 2Rs. (Ar(S)))
(46)
where 40
(S&) the ERM on St.
is based on the following inequality.
2 <
la,11
(47)
i=1
18

---

Published in Transactions on Machine Learning Research (10/2025)
Secondly, we derive new bound for BAN.K,(S. VV 1 mi side of We get
E ((A(S()) = the right hand
(A(S(M)) =
2L
(A((S())) ),
Σ (SAr))),, =h.u)
(48)
Th
1
$ 2L
Rs; (A(((())) Rs, (s,,(A()5)))
(49)
$ 2L
(s Rss; (AA (S%)) + Rs, (Av(S)) - Rs; (AA
(50)
$ 2L
(s(A()( Rs; (Sx)) - Rs; (s,)
(51)
Let's define Ex (S&) Rss. (As(s))) Rss (A+). Putting
17 into 16
and considering
21
we get
B (SLDD (A(S) S B
B(S |A(S) A()
E BI 1mm1
((((S()), 2 12 2 E (s.. ||A(S) A(S()))12
2LK(())2
(AA( + 2cx(
$ B
BISINDAA
(52)
4KC(K)2
E BISENDRHEE +V E (/ 1 2L (5(0())) + (x(())
(Aa) 2C 2cx(Sk)
2LK(K)2
(AL(SA)) 2Ck(Sk)
Bank BIS
(53)
SL + x(k)
Bis (om(()) + (S((s))
B(S
(Sx) + 22x 2cx((k))
where in 52 we have applied
46 and
51
This completes the proof.
Now, we prove Theorem 4.3 as follows.
Theorem C.2. Let 1(Mo,2) be j-strooggy conver and L-smooth in Me. Local models at round calculated by doing local gradient descent sleps (5), and the local gradient variance is bounded by o², i.c., Ez-D, [| 2) - Ez-D, VI(Mo, 2)[l < o². The aggregated model at round Ma.- is obtained by performing FedAug, and the data points used in round (i.e., Zk.r) are sampled without replacement. The average generulization error bound is
R 2LK(K)2
8L
Eg A FedAug (S)
Σ
ERnK
A+ +
(AB)
R
where A= = -6(V (CMM) + 6² [Ze,rl pr) logarithmic factors, and C(Mo) shows the complerity of the model class of Mo-
B= B B O 8k. [k o² Ö hides constants and poly-
19

---

Published in Transactions on Machine Learning Research (10/2025)
Proof. Based on the definition, have
K
E (S 1k-1
APedAug = D5 1h-1
1 2 BanK R
BawD, 1
Σ
1Ad0, =k.1)
(54)
1Zwirl IE?N.
1
E
Bink km
ENNDE 1(Mo,
Σ :
(55)
R
(Z
(Zx,
IZurrl ieZN..
1
R
E
(Z 1
Aa(tZR,)kk))
(56)
1
2LK(()}
(Aal ZC (Zxi))
BR
E
(57)
R
1
+1
8L
(((R)
V
(Znnnolferrl Zxo + txx (
(f (Ann (2k,r) + cx (z.r)
1
2LK(())2
C(Mo)
0²
Ex
(58)
R r=1
NuT
8L +V IC(R))
(k + JTT (( C(Mo)
0² 1AT
E.
Z
where in 166 A represents one-round FedAvg algorithm. In 57 we have used Theorem C.1 In
58
we
have used the conventional statistical learning theory originated with Leslie Valiant's probably approximately correct (PAC) framework Valiant 1984). We have also applied the optimization convergence rate bounds
the literature Stich & Karimireddy (2019). Note that Ô hides constants and logy-loga factors.
D Partial Client Participation Setting
We first define empirical risk for the partial participation distribution i on dataset S, where Supp(KC) # (I,..., K} and [Supp(K)) = K < K, as
(Mo) = =E Rs, (Mo) = E Th
(MMo 2k,s),
(59)
where IC is an arbitrary distribution on participating nodes that is part of all nodes, and Rs. (Me) is the empirical risk for model Mo on local dataset Sg. We further define partial population risk for model Me
R H& (Mo) = E (x(M)) = EE-R, (Me, z),
(60)
where (M)) the population risk on node k's data distribution.
Now, we can define the generalization error for dataset. S and function A(S)
AS R(A(S)) R2 (A(S)) ((A((S))) - + R (A(S)) R
(61)
(62)
Participation gap
Out-of-sample gap
The expected generalization error is expressed as Eg.-DA**k=11 (S)).. Note that the second term in which is related to the difference between in-sample and out-of-sample loss, can be bounded in the same way as in Theorem 4.1 and Theorem 4.3 The first term is associated with the participation of not all clients. In the following, we demonstrate that under certain conditions, this term would be zero in expectation.
62
We assume there is meta-distribution P supported all distributions k
20

---

Published in Transactions on Machine Learning Research (10/2025)
Lemma D.1. Let { denote any fized deterministic sequence. Assume P is derived by sampling K clients with replacement based on distribulion K. followed by an equal probability on all sompled clients, i.e., R(k)= = k Then
Ep E Ik = BRAK Lki Ep R R(k)zx = R
BR Lk. Ep E (K))) = R2
B Ik.
(63)
Proof.
k
1
Ep E knR 2k Ep
I1
Ep Ij
Ep 2; B Tk
(64)
i=1
K
K
1
Ep B k(kk).x+ = Rp
Ep
Ep
B
(65)
R
K
K
K
==1
i=1
Ep E Bank (k(xx)x = Ep
Ep
Ep 2j
Ex Ik
(66)
R33
R3
R22
R²
Lemma D.2. Let (xi) denole any fixed delermánislic sequence. Assume P is derived by sampling R clients without replacement uniformly at random followed by weighted probobility all sampled clients R(k) K(R)K Then K
K B K(K)aak Ep E R²(x)0x K
K² EanK K2 (k) (67)
Ep E th = Ex-x The Ep E R(k)ax
K2
Proof.
K R
Ep E vb Ep
K(i);;
K
Ep =())
K KE) K 1
Σ
K(i)a; E F4
K
K
K
(68)
2
K2 K
K² -
K² 1
Ep Ex R(k)0z Ep
Σ
(i)xx
Σ
Ep 12
12 (1)x;
Ep
K3
K2 K² (69)
R2
R2
K
K K
K (B
R
K3
K3 K Re
K3 (i)x;
K3
Ep Ex (k))k)² = Ep
Rs
Σ
x))))
Σ
Ep (1)x,
R2
Ep
Ep
=
R2 K
K K3) (70)
K²
E K2 (2(
= R²
So based on lemmas D.I and D.2 it becomes evident that the expectation of the participation gap in becomes zero for both two methods, i.e.,
62
Ep H(A(S)) - ={((A()))) = O.
(71)
The expected generalization error, D 1 An(S), be just the expectation of the second term in 62 that can be bounded using Lomma B.2 by
L(((K)²
L [R(H) B (S) (SSWDDj AA. (Sx)
Bx
E (SSNDDN (Sx) + +2
(72)
1r
1
21

---

Published in Transactions on Machine Learning Research (10/2025)
If take the expectation of 72 generalization bound for method as
with respect to P, and taking into account Lemma D.1
get the
I. R Lik2 (S (Sa) + 3
To
BRn
B (S (5) AA (Sx)
(73)
For Scheme 2, we can obtain the generalization bound in the same way by taking the expectation of 72 with respect lo and considering Lemma D.2 We get the generalivation bound for Method as:
IK(K)222 /
I.
I. K K
E
E B(SEND (Sa) +2
(B (E AA, (Sx)
K2
(74)
E
Detailed Experimental Setup
E.1
Image Classification
The details are specified in Table
Table 4: Default experimental settings for the image classification training
Dataset
CIFAR-10, CIFAR-100
Krizhevsky
2009),
SVHN
Netzer et al. 201
MNIST
[ecunn et al. 1998
Architccture
RosNet--2 He et al.
2015
Training objective
Cross entropy
Test objective
Top-1 accuracy
Number of clients Data distribution
IID (shufiled and split),
Non-IID (sorted based on labels then split)
Local Steps
(unless explicitly specified)
Adaptation coefficient
10 (unless explicitly specified)
Representation Extractor
Convolutional layers
(unless explicitly specified)
Head
Final dense layers
(unless explicitly specified)
Optimizer
SGD with momentum
Batch size
64 per client
Momentum
0.9 (Nesterov)
Learning rate
Constant 17 after tuning
Number of Iterations
10 for IID and 2 104 for non-IID
Weight decay
10-4
Repetitions
20, with varying seeds
Reported metric
Mean and standard deviation of the aggregated model's test. ac- curacies over the last 5 rounds
22

---

Published in Transactions on Machine Learning Research (10/2025)
E.2 Large Language Model Fine-tuning
The details are specified in Table
Table Default experimental settings for the large language model inn-tuning
Datasct
Multi-Genre Natural Language Inference (MultiNT.I)
corpus
Architecture
OPT-125M He et al. 2015
Training objective
Cross entropy
Test objective
Top-1 accuracy
Number of clients
5
Data distribution
IID (shuffled and split),
Non-IID (sorted based on genre then split)
Local Steps
5
Adaptation coefficient a
10
Representation Extractor Head
First 10 attention layers
Final 2 attention layers
Optimizer
AdamW
Batch size
16 sentences per client
Adam $1
0.9
Adam B2
0.999
Adam
10-8
Learning rate
Decaying 100 +10
where the itcration number and ry is tuned
Number of Iterations
10 for IID and 2 x 10 for non-IID
Weight decay
10-4
Repetitions
20, with varying seeds
Reported metric.
Mean and standard deviation of the aggregated model's test
curacies over the last 5 rounds
E.3 papper-prrmmttrs and tuning details
We independently tuned the learning rates for both image classification and large language model (LLM) fine-tunin Ene-tuning experiments for each dataset, algorithm, and data heterogeneity scenario (both iid and non-iid), For image classification, the learming rate was tuned, while for LLM fino-tuning, we specifically tuned the initial learning rate. The luning process followed these steps:
Begin with an initial value of 0.01 and perform the experiment several times, each time using a different random seed.
Perform a grid search to find the optimal learning rate. Start by multiplying or dividing the initial value by powers two, Test both larger and smaller learning rates, running the experiment multiple times with different random scods. The best lcarning rate the one that, on average, produces the best results. If this learning rate is between two others that give worse results, then you have found the optimal value.
23

---

Published in Transactions on Machine Learning Research (10/2025)
E.4
Algorithms Used in the Experiments
In this section, provide detailed list of our implementation of
FedAvg (Algorithm
SCAFFOLD Karimireddy et al. 2019 (Algorithm 3
Combined method FedALS+SCAFFOLD (Algorithm
Algorithm [ demonstrates the implementation FodAvg. This is a widely used algorithm in federated learning that performs model averaging afler local updates from each client.
Moving on Algorithm we observe that SCAFFOLD introduces an additional mechanism beyond just the model updates. Specifically, SCAFFOLD keeps track of state that is specific to each client, referred to as the client control variate, denoted by Ck,r- This control variate is a key feature of SCAFFOLD, which helps to correct the client drift by utilizing these client-specific states. It's important to note that the clients within SCAFFOLD maintain memory and persistently store the values of Ch.r and Σk= Ck.r- This stored information plays a crucial role in the algorithm's 's functionality. Additionally, is worth recognizing that when the value of Ct.r consistently remains at zero, the SCAFFOLD algorithm essentially reduces FedAvg. This highlights the close relationship between the two algorithms and emphasizes the special role of the control variates in differentiating SCAFFOLD from FedAvg.
Finally, Algorithm showcases the integration of FedALS with SCAFFOLD, presenting hybrid approach. It is important to observe that in this combined algorithm, the control variables are not handled in a uniform manner but are instead fragmented according to the different partitions the model. This fragmentation occurs because there are distinct local step counts for different parts of the model, leading to a more complex update mochanism compared the standard SCAFFOLD approach.
Algorithm 2 FedAvg
Input: Initial model (0,.10]k Learning rale v), and number of local steps т.
Output: 0R
1: for Round r in L,...., R do for Node 1,... K in parallel do for Local step in O,..,, do Sample the batch BR from Dr.
0K. O end for 0R.++100 end for end for
Ba.rol LieBu... VI(Mo Rk.i)
Ek=1 ORnT
10: return 0R- = K
1 Ek= OR.R.T
24

---

Published in Transactions on Machine Learning Research (10/2025)
Algorithm SCAFFOLD
Input: Initial model (0k,100)== Initial control variable (Ct,) learning rate n, and number of local
steps т.
Output: 0R
for Round r in 1, ..... R do for Node k in 1, K in parallel do for Local step in 0,...T- do Sample the batch BR from DA. Ok.r,++1= 0k,,, n7( Bk.ne EieBan VI( Vi(Monne Zk,e) okr + =k=1 Ck.r)
6: end for
Ek Ek=1 Ckrt + 1 (0k.r.0 kk.r.)
7:
Ck.r+1 Ck.r R
8: O 0k.
K k= OK.
end for
10: end for
11: return On R Ek- RR,
Algorithm FedALS SCAFFOLD
Input:
Initial model 10 Lo, ho. | Initial control variable {cx,1 = [cf, d ( learning
rate v7, number of local steps for the head model T, adaptation coefficient a.
Output: 0R
1: for Round r in 1,..., R do
2: for Node in L,.... K in parallel do do
for Local step in 0,..., Sample the batch BRrr from Dr.
5: :R. OR,rT - 17( Bk.nel EieB VI(Mo Zk,e) - Cx. + K Ex CR.r)
if mod (rT r) = then
de +- Chr K 5 Ckr + np (ha,r.0 hvna))
hk.r. + K -x=1 hx.rt
end if
10: if mod (rT 1, aT) 0 then
11:
kr t- -ch * Ex c
LK
nar (
( mod (rr4t-07,7) Paan))
ΦR K > Øk.ra
12:
13:
end if
14:
end for
15:
Ck.++1 Ck.r
16:
0kr+110 = 0k.n.
17:
end for
18:
end for
19: return 0R * Ex=1 OR,
25


Reviews:
[{"review_id": "g2rXQZ7xUN", "reviewer": "Reviewer_MdNz", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: The paper's core contribution is FedALS, a novel and intuitive algorithm that improves both accuracy and communication efficiency in non-iid federated learning. The key insight—to aggregate the model's initial \"representation extractor\" layers less frequently than its final \"head\" layers—is well-motivated by both a theoretical generalization bound analysis (Section 4) and compelling empirical evidence showing lower consensus distance in early layers (Figure 1). The experimental results in Table 1 and Table 2 strongly validate this approach, demonstrating that FedALS achieves higher final accuracy while simultaneously reducing communication overhead.\n\nLimitations: The primary weakness is the paper's contradictory theoretical justification for its core idea. In Section 5, the authors claim that increasing the number of local steps (τ) tightens their generalization bound in Theorem 4.3, which motivates the FedALS algorithm. However, the same theorem includes a non-iidness term (δk,A) that, as the authors acknowledge, worsens with increased τ. The paper fails to formally analyze or resolve this critical trade-off within its own theoretical framework, which undermines the claim that the algorithm's design is directly and rigorously derived from the provided bounds.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Comment 1: The paper's theoretical motivation is contradictory. The authors claim in Section 5 that increasing local steps (τ) improves the generalization bound in Theorem 4.3, yet they also state it increases non-iidness (δk,A), a term that worsens the same bound. The paper must provide a rigorous analysis that formally reconciles this tension to justify its core theoretical claim.\n\nComment 2: The method for splitting the model into a \"representation extractor\" and a \"head\" is presented as a critical design choice in Algorithm 1 but is treated as an arbitrary hyperparameter in the experiments (Section 6 and Table 3). For FedALS to be a generalizable method, the authors must provide a principled, systematic procedure for determining this architectural split, rather than relying on model-specific ad-hoc tuning.\n\nComment 3: The experimental comparison to SCAFFOLD in Section 6.3 is misleadingly framed as competitive rather than complementary. The fact that \"FedALS + SCAFFOLD\" achieves the best performance in Table 1 strongly indicates the two methods address orthogonal problems. The paper should reframe this discussion to highlight this synergy, positioning FedALS as a communication-efficiency method that can be combined with variance-reduction techniques.\n\nComment 4: The empirical analysis of the key hyperparameter α in Table 2 is disconnected from the paper's own theoretical framework. The authors observe that performance degrades for α > 10 due to non-iidness but miss the crucial opportunity to connect this finding back to the generalization bound in Theorem 4.3. A stronger analysis would empirically estimate the bound's components to show how the theory explains this observed performance peak.\n\nComment 5: The core motivation for FedALS—that initial layers have lower consensus distance—is only empirically validated for the ResNet-20 model in Figure 1. The paper then applies the same logic to fine-tuning an OPT-125M model in Section 6 without providing any evidence that its initial layers exhibit the same property. This unsubstantiated assumption weakens the justification for applying FedALS to LLMs.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "N/A", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "mVTHvth1eE", "reviewer": "Reviewer_cj4e", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\nTheoretically sound derivation of a new FL method that decreases generalization error while also decreasing the communication overhead is a crucial objective in FL.\n\nWeaknesses:\nMore experiments are needed to verify the proposed method empirically. In particular, it is not clear whether the method will work on a larger set of clients and how it behaves when the fraction of clients used in each round for training varies. Note, this is commonly analyzed in FL. Moreover, there is no information on whether the method can compete with other popular global FL methods, e.g. FedOpt [1] etc.\n\n[1] Reddi, Sashank, et al. \"Adaptive federated optimization.\" arXiv preprint arXiv:2003.00295 (2020).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a novel and tighter generalization bound for FL in the non-iid. setting and, together with empirical findings from the field of representation learning, propose a new FL algorithm that reduces communication efficiency while improving performance (test loss and accuracy). In particular, they find that their derived generalization bound can be decreased by decreasing the aggregation frequency of the feature extractor, as compared to the classification head. This, in turn, increases the number of samples and local SGD steps terms in the bound, while keeping the increase of the non-iidness term comparatively low, which ultimately decreases the derived bound.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Add more baselines, e.g.: FedProx [2], FedOpt [1].\n- Test the method on a larger number of clients, e.g., 100, as is commonly done [3]. \n- As an ablation study experiment, vary the fraction of clients used for training in each round, e.g. 10 out of 100, randomly sampled. \n- As an ablation study experiment, vary \n- While in Figure 2 you plot (from left to right): \nTraining global loss (non-iid), Test accuracy (non-iid), Training global loss (iid.), Training accuracy (iid.), in Figure 3 you plot: Training global loss (non-iid), Test loss (non-iid), Training global loss (iid.), Training accuracy (iid.). Why? Please plot, from left to right: \nTraining global loss (non-iid), Test accuracy (non-iid), Training global loss (iid.), Training accuracy (iid.) for both the CNN and the LLM. if you want to add test loss (non-iid), add another column.\n- Tables 2 & 3 include ablation studies for L and alpha only for the CNN. Please run the same ablation study experiments with the LLM.\n- For the non-iid experiments, add the commonly used Dirichlet label split method [4] for your experiments and real-life FL splits, such as FEMNIST [5].\n\n[1] Reddi, Sashank, et al. \"Adaptive federated optimization.\" arXiv preprint arXiv:2003.00295 (2020).\n\n[2] Li, Tian, et al. \"Federated optimization in heterogeneous networks.\" Proceedings of Machine learning and systems 2 (2020): 429-450.\n\n[3] McMahan, Brendan, et al. \"Communication-efficient learning of deep networks from decentralized data.\" Artificial intelligence and statistics. PMLR, 2017.\n\n[4] Yurochkin, Mikhail, et al. \"Bayesian nonparametric federated learning of neural networks.\" International conference on machine learning. PMLR, 2019.\n\n[5] Caldas, Sebastian, et al. \"Leaf: A benchmark for federated settings.\" arXiv preprint arXiv:1812.01097 (2018).\n\n[6] Kairouz, Peter, et al. \"Advances and open problems in federated learning.\" Foundations and trends® in machine learning 14.1–2 (2021): 1-210.", "broader_impact_concerns": "If it turns out that the method doesn’t work for a larger number of clients, a limitations section should be added that explains that the method is only applicable to the cross-silo FL setting [6] with full client participation.\n\n\n\n[6] Kairouz, Peter, et al. \"Advances and open problems in federated learning.\" Foundations and trends® in machine learning 14.1–2 (2021): 1-210.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "LUXbPsfnLS", "reviewer": "Reviewer_7HKh", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "1) The paper carefully separates the contributions coming from non-IIDness to demonstrate that the previous error bounds 1/K improve to 1/K^2 when the dataset across all nodes is uniformly distributed.\n\n- In Eq.9, the abstract model complexity's appearance obscures the bound's practical applications. It would be helpful to demonstrate the tightness of the bound, at least for synthetic data scenarios.\n\n2) The FedALS algorithm is empirically well-motivated by the observation that early layers should learn global features rather than head layers, which are more task-specific. The algorithm is easy to implement and shows superior performance to its predecessors.\n\n- The bounds depend on strong assumptions on convexity and smoothness, which are not guaranteed in deep neural nets. While FedALS in demonstrated examples performs better than previous methods, it is unclear whether this will hold in other applications. It would be helpful to present a numerical test of how predictive the theory is for at least synthetic data.\n\n- The representation learning perspective is mainly justified based on empirical observations. I believe this narrative can be strengthened by theoretical studies that demonstrate that the representations in earlier layers align with the input data and final layers with the task structure. Here are a few works that I know of:\n    - https://arxiv.org/pdf/2012.04030\n     - https://arxiv.org/abs/2106.00651\n    - https://arxiv.org/abs/2205.09653\n\n- Beyond expensive hyperparameter searches, is there a principled way to choose $\\alpha$ based on theory? What are some factors that $\\alpha$ strongly depends on?\n\nOverall, I believe this is a solid paper, both deriving stronger bounds for federated learning and introducing a novel algorithm. It provides valuable insight into federated learning and is clearly written. I think a deeper and clearer discussion of its theoretical assumptions, experimental guidance, and its limitations would strengthen the paper further.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper conducts a detailed analysis of the communication cost of federated learning algorithms, specifically FedAvg, in terms of generalization. Based on their analysis, the authors observe that the representations across clients learn more shared features in early layers than final layers, and accordingly propose a new algorithm (FedALS) which adaptively reduces the aggregation frequency of local weights for those layers. Through extensive experimentation, the authors demonstrate that FedALS leads to better generalization with reduced communication overhead.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See above.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)