Communication Efficient Federated Learning with Differentiated Aggregation

Authors: peyman  gholami
hulya  seferoglu
Venue: ICML_cc_2024_Workshop_WANT
Published: 2024-05-28 01:07:39

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 our representation learning interpretation of this analysis, we show for the first time that less frequent aggregations, hence more local updates, for the representation extractor (usually corresponds to initial layers) leads to the creation of more generalizable models, particularly for 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.

Main paper: 
 introduction: 
 Federated learning advocates that multiple clients collaboratively train machine learning models under the coordination of a parameter server (central aggregator) (McMahan et al., 2016) . This approach has great potential for preserving the privacy of data stored at clients while simultaneously leveraging the computational capacities of all clients. Despite its promise, federated learning still suffers from high communication costs between clients and the parameter server.In a federated learning setup, a parameter server oversees 1 Department of Electrical and Computer Engineering, University of Illinois Chicago, Chicago, Illinois, United States. Correspondence to: Peyman Gholami <pghola2@uic.edu>. Learning (WANT@ICML 2024) . a global model and distributes it to participating clients. These clients then conduct local training using their own data. Then, the clients send their model updates to the parameter server, which aggregates them to a global model. This process continues until convergence. Exchanging machine learning models is costly, especially for large models, which are typical in today's machine learning applications (Konecný et al., 2016; Zhang et al., 2013; Barnes et al., 2020; Braverman et al., 2015) . Furthermore, the uplink bandwidth of clients may be limited, time-varying and expensive. Thus, there is an increasing interest in reducing the communication cost of federated learning especially by taking advantage of multiple local updates also known as "Local SGD" (Stich, 2018; Stich & Karimireddy, 2019; Wang & Joshi, 2018) . The crucial questions in this context are (i) how long clients shall do Local SGD, (ii) when they shall aggregate their local models, and (iii) which parts of the model shall be aggregated. The goal of this paper is to address these questions and reduce communication costs without hurting convergence.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 to 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 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 own dataset, which causes variation in consecutive local gradients, as well as (ii) the dissimilarity in loss functions among clients due to non-iidness (Stich & Karimireddy, 2019; Gholami & Seferoglu, 2024) . More specifically, the consensus distance at iteration t is defined as 1 K K k=1 ∥ θtθ k,t ∥ 2 , where θt = 1 K K k=1 θ k,t , K is the number of clients, θ k,t is the local model at client k at iteration t, and ∥ • ∥ 2 is squared l 2 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 the communication overhead of federated learning through the following contributions.Contribution I: Improved Generalization Error Bound. The generalization error of a learning model is defined 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 important to quantitatively assess the generalization characteristics of trained models, provide reliable guarantees concerning their anticipated performance quality, and design new models and systems.In this paper, we offer a tighter generalization bound compared to the state of the art (Barnes et al., 2022; Yagli et al., 2020; Sun et al., 2023) for one-round federated learning, considering local clients' generalizations and non-iidness (i.e., heterogeneous data distribution across the clients). Additionally, we characterize the generalization error bound in R-round federated learning.Recent studies have demonstrated that the concept of representation learning is a 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' datasets. For example, let us consider a federated learning application for image classification, where different 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 features such as an eye/ear 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 a dog or a cat. As a result, these layers demonstrate similarity (i.e., less variation) across clients even when the datasets are noniid. This implies 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.Example 1. We consider a federated learning setup of five clients with a central parameter server to train a ResNet-20 (He et al., 2015) on a heterogeneous partition of CIFAR-10 dataset (Krizhevsky, 2009) . We use Federated Averaging (FedAvg) (McMahan et al., 2016) as an aggregation algorithm since it is the dominant algorithm in federated learning. We applied FedAvg with 50 local steps prior to each averaging step, denoted as τ = 50. Non-iidness is introduced by allocating 2 classes to each client. Finally, we evaluate the quantity of the average consensus distance for each model layer during the optimization in Fig. 1 . It is clear 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 similarity across clients.The above example indicates that initial layers show higher similarity, so they can be aggregated less frequently. 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. These studies encourage us to delve deeper into investigating how local updates and model aggregation frequency affect the model's representation extractor in terms of its generalization.In this paper, based on our improved generalization bound analysis and our representation learning interpretation 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 interpretation. 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 for CIFAR-10, CIFAR-100 (Krizhevsky, 2009) , SVHN (Netzer et al., 2011) , and MNIST (Lecun et al., 1998) datasets. We also estimate the impact of FedALS on large language models (LLMs) in fine-tuning OPT-125M (Zhang et al., 2022) on the Multi-Genre Natural Language Inference (MultiNLI) corpus (Williams et 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 setups. 
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 to a global optimum or a stationary point of the overall objective, particularly in convex or 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 tackle this issue, two primary strategies are put forth: the utilization of mini-batch parallel SGD, and the adoption of Local SGD. These approaches aim to enhance the equilibrium between computation and communication. Woodworth et al. (2020b; a) attempt to theoretically capture the distinction to comprehend under what circumstances Local SGD outperforms minibatch SGD.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 between 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., 2023) . One prominent example is the Stochastic Con-trolled Averaging algorithm (SCAFFOLD) (Karimireddy et al., 2019) , which addresses the node drift caused by noniid characteristics 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 FedAvg and incorporating multiple Local SGD per round exhibits unexpected effectiveness when subsequently fine-tuned for individual clients in non-iid FL 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 each node for downstream 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) considered this problem and gave upper bounds on the expected generalization error for FL in iid setting in terms of the local generalizations of clients. This work demonstrates an improved dependence of 1 K on the number of nodes. Motivated by this work, we build our research foundation by analyzing generalization in a non-iid setting and use the derived insights to introduce FedALS, aiming to enhance conventional machine learning generalization.Note that FedALS differs from exploiting shared representations for personalized federated learning, as discussed in Collins et al. (2021) . In FedALS, we do not employ different models on different clients, as seen in personalized learning. Our proof demonstrates that increasing the number of local steps enhances generalization in the standard (single-model) federated learning setting. 
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 k has a local datasetS k = {z k,1 , ..., z k,n k }, where z k,i = (x k,i , y k,i ) is drawn from a distribution D k over X × Y ,where X is the input space and Y is the label space. We consider X ⊆ R d and Y ⊆ R . The size of the local dataset at node k is n k . The dataset across all nodes is defined as S = {S 1 , ..., S K }. Data distribution across the nodes could be independent and identically distributed (iid) or non-iid. In iid setting, we assume that D 1 = ... = D K = D holds. On the other hand, non-iid setting covers all possible distributions and cases, where D 1 = ... = D K = D does not hold.We assume that M θ = A(S ) represents the output of a possibly stochastic function denoted as A(S ), where M θ : X → Y represents the learned model parameterized by θ. We consider a real-valued loss function denoted as l(M θ , z), which assesses the model M θ based on a sample z.We first define an empirical risk on dataset S asEQUATION)where K is an arbitrary distribution over nodes to weight different local risk contributions in the global risk. Specifically, K(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.e.,K(k) = 1 K for all k. R S k (M θ ) is the empirical risk for model M θ on local dataset S k . We further define a population risk for model M θ as R(M θ ) = E k∼K R k (M θ ) = E k∼K,z∼D k l(M θ , z), (2)where R k (M θ ) is the population risk on node k's data distribution. Now, we can define the generalization error for dataset S and function A(S ) as∆ A (S ) = R(A(S )) -R S (A(S )).(3)The expected generalization error is expressed asE S ∆ A (S ), where E S [•] = E {S k ∼D n k k } K k=1 [•]is used for the sake of notation convenience.We consider a federated learning scenario with K nodes/clients and a centralized parameter server. The nodes update their localized models to minimize their empirical risk R S k (M θ ) on local dataset S k , while the parameter server aggregates the local models to minimize the empirical risk R S (M θ ). Due to 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 r of FedAvg, each node k trains its model M θ k,r = A k,r (S k ) locally using the function/algorithm A k,r . The local models M θ k,r are transmitted to the central parameter server, which merges the received local models to aggregated model parameters θr+1 = Â(θ 1,r , ..., θ K,r ), where Â is the aggregation function. In FedAvg, the aggregation function calculates an average, so the aggregated model is expressed asEQUATIONSubsequently, 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 at round r as θr . Specifically, upon receiving θr , node k computesθ k,r,t+1 = θ k,r,t - η |B k,r,t | i∈B k,r,t ∇l(M θ k,r,t , z k,i ) (5)for t = 0, . . . , τ -1, where τ is the number of local SGD steps, θ k,r,0 is defined as θ k,r,0 = θr , η is the learning rate, B k,r,t is the batch of samples used in local step t of round r in node k, ∇ is the gradient, and | • | shows the size of a set. Upon completing the local steps in round r, each node transmits θ k,r = θ k,r,τ to the parameter server to calculate θr+1 as in (4).Our approach for analyzing the generalization error bounds for federated learning, by specifically focusing on FedAvg, uses representation learning, which we explain next.We consider a class of models that consist of a representation extractor (e.g., ResNet). Let θ be the model M θ 's parameters. We can decompose θ into two sets: ϕ containing the representation extractor's parameters and h containing the head parameters, i.e., θ = [ϕ, h]. M ϕ is a function that maps from the original input space to some feature space, i.e., M ϕ :R d → R d ′, where usually d ′ ≪ d. The function M h performs a low complexity mapping from the representation space to the label space, which can be expressed asM h : R d ′ → R. For any x ∈ X , the output of the model is M θ (x) = (M ϕ • M h )(x) = M h (M ϕ (x)). For instance, if M θ is a neural network, M ϕrepresents several initial layers of the network, which are typically designed to extract meaningful representations from the neural network's input. On the other hand, M h denotes the final few layers that lead to the network's output. 
improved generalization bounds: 
 In this section, we derive generalization bounds for FedAvg based on clients' local generalization performances in a general non-iid setting for the first time in the literature. First, we start with one-round FedAvg and analyze its generalization bound. Then, we extend our analysis to R-round FedAvg.In the following theorem, we determine the generalization bound for one round of FedAvg.Theorem 4.1. Let l(M θ , z) be µ-strongly convex and L-smooth in M θ . M θ k = A k (S k ) represents the model obtained from Empirical Risk Minimization (ERM) algorithm on local dataset S k , i.e., M θ k = arg min M n k i=1 l(M, z k,i), and M θ = A(S ) is the model after one round of FedAvg ( θ = E k∼K θ k ). Then, the expected generalization error isE S ∆ A (S ) ≤ E k∼K LK(k) 2 µ E S k ∆ A k (S k ) Expected local generalization (6) + 2 L µ K(k) E S δ k,A (S )Expected non-iidnessE S k ∆ A k (S k ) Expected local generalization 1 2 , where δ k,A (S ) = R S k (A(S )) -R S k (A k (S k ))indicates the level of non-iidness at client k for function A on dataset S .Proof: The proof of Theorem 4.1 is provided in Appendix A of the supplementary materials. □ 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 I: Sampling K clients with replacement based on distribution K, followed by averaging the local models with equal weights.Case II: Sampling K clients without replacement uniformly at random, then performing weighted averaging of local models. Here, the weight of client k is rescaled to K(k)K K . The generalization error results in these cases are affected by substituting 1 K and K(k)K K instead of K(k) in (6) for cases I and II, respectively. The detailed proof is provided in Appendix C.Discussion. Note that there are two terms in the generalization error bound: (i) local generalization of each client that shows more 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 a factor of K(k) 2 for the first term, which is the sole term in the iid setting. For example, in the uniform case (K(k) = 1 K ), we will observe an improvement with a factor of 1 K 2 for the iid case. This represents an enhancement compared to the state of the art (Barnes et al., 2022; Sun et al., 2023; Yagli et al., 2020) , which only demonstrates a factor of 1 K . As a result, after the averaging process carried out by the central parameter server, the generalization error is reduced by a factor of K(k) 2 in iid case.On the other hand, we do not see a similar behavior in noniid case. In other words, the expected generalization error bound does not necessarily decrease with averaging. These results show why FedAvg works well in iid setup, but not necessarily in non-iid setup. This observation motivates us to design a new federated learning approach for non-iid setup. The question in this context is what should be the new federated learning design. To answer this question, we analyze R-round generalization bound in the next section.In this setup, after R rounds, there is a sequence of weights { θr } R r=1 and the final model is θR . We consider that at round r, each node constructs its updated model as in ( 5) by taking τ gradient steps starting from θr with respect to τ random mini-batches Z k,r = {B k,r,t } τ -1 t=0 drawn from the local dataset S k . For this type of iterative algorithm, we consider the following averaged empirical riskEQUATION)The corresponding generalization error, ∆ F edAvg (S ), is1 R R r=1 E k∼K E z∼D k l(M θ r , z)- 1 |Z k,r | i∈Z k,r l(M θ r , z k,i ) . (8)Note that the expression in (8) differs slightly from the end-to-end generalization error that would be obtained by considering the final model M θ R and the entire dataset S . More specifically, ( 8) is an average of the generalization errors measured at each round, similar to (Barnes et al., 2022) ). We anticipate that the generalization error diminishes with the increasing number of data samples, so this generalization error definition yields to 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.Theorem 4.3. Let l(M θ , z) be µ-strongly convex and Lsmooth in M θ . Local models at round r are calculated by doing τ local gradient descent steps (5), and the local gradient variance is bounded by σ 2 , i.e., E z∼D k ∥∇l(M θ , z) -E z∼D k ∇l(M θ , z)∥ 2 ≤ σ 2 . The aggregated model at round r, M θ r , is obtained by performing FedAvg, and the data points used in round r (i.e., Z k,r ) are sampled without replacement. Then the average generalization error,EQUATION)where ), number of local steps in each round; τ . We note that (9) also depends on K (more specifically K), but this dependence is similar to the discussion we had for one-round generalization, so we skip it here.A = Õ C(M θ ) |Z k,r | + σ 2 µτ + L µ , B = Õ E {Z k,r } K k=1 δ k,A ({Z k,r } K k=1 ) + σ 2 µτ + L µ , Õ 
representation learning interpretation of r-round generalization bound: 
 The complexity of the model class and the number of samples and local steps used in every round are crucial to minimizing the generalization error bound especially in non-iid case in (9) where the generalization error bound is loose in comparison to iid setup.Some common complexity measures in the literature include the number of parameters (classical VC Dimension (Shalev-Shwartz & Ben-David, 2014)), parameter norms (e.g., l 1 , l 2 , spectral) (Bartlett, 1997) , or other potential complexity measures (Lipschitzness, Sharpness, . . . ) (Neyshabur et al., 2017; Dziugaite & Roy, 2017; Nagarajan & Kolter, 2019; Wei & Ma, 2019; Norton & Royset, 2019; Foret et al., 2021) . Independent from a specific complexity measure, a model in representation learning can be divided into two parts: (i) M ϕ , which is the representation extractor, and (ii) M h , a simple head which maps the representation to an output. The complexities of these parts follow C(M h ) ≪ C(M ϕ ).Our key intuition in this paper is that we can reduce the aggregation frequency of M ϕ , which leads to a larger τ and |Z k,r |, hence smaller generalization error bound according to (9). 1As seen, there is a nice trade-off between aggregation fre-1 We do not reduce the aggregation frequency of M h as its complexity, so its contribution to generalization error, is small.Input: Initial model {θ k,0,0 = [ϕ k,0,0 , h k,0,0 ]} K k=1 , learning rate η, number of local steps for the head τ , adaptation coefficient α.1: for Round r in 0, ..., R -1 do 2:for Node k in 1, ..., K in parallel do 3:for Local step t in 0, ..., τ -1 do 4:Sample the batch B k,r,t from D k .5:θ k,r,t+1 = θ k,r,t - η |B k,r,t | i∈B k,r,t ∇l(M θ k,r,t , z k,i ) 6:if mod (rτ + t, τ ) = 0 then 7:h k,r,t ← 1 K K k=1 h k,r,t 8:else if mod (rτ + t, ατ ) = 0 then 9:ϕ k,r,t ← 1 K K k=1 ϕ k,r,t 10: θ k,r+1,0 = θ k,r,τ 11: return θR = 1 K K k=1 θ k,R,0quency of M ϕ and population risk. In the next section, we design our Federated Learning with Adaptive Local Steps (FedALS) algorithm by taking into account this trade-off. Remark 5.1. We note that the aggregation frequency of M ϕ cannot be reduced arbitrarily, as it would increase the empirical risk. It is proven that the convergence rate of the optimization problem of ERM in a general non-iid setting for a non-convex loss function is Koloskova et al., 2020 ). Here T is the total number of iterations, i.e., T = τ R.O τ T + τ T 2 3 + 1 √ T ( 
fedals: federated learning with adaptive local steps: 
 Theorem 4.3 and our key intuition above demonstrate that more local SGD steps (less aggregations at the parameter server) are necessary for representation extractor M ϕ as compared to the model's head M h to reduce generalization error bound. This approach, since it will reduce the aggregation frequency of M ϕ , will also reduce the communication cost of federated learning.The main idea of FedALS is to maintain a uniform generalization error across both components (M ϕ and M h ) of the model. This can be achieved if τ M ϕ is set larger than τ M ϕ , where τ M denotes the number of local iterations in a single round for the model M while τ M ϕ and τ M h are the corresponding number of local iterations for M ϕ and M h , respectively. Following this approach, we designed FedALS in Algorithm 1.FedALS in Algorithm 1 divides the model into two parts: (i) the representation extractor, denoted as M ϕ , and (ii) the head, denoted as M h . Additionally, we introduce the parameter α = ratio. Note that this ratio depends on C(M ϕ ) and C(M h ), and determining these values is not straightforward. 
experimental results: 
 In this section, we assess the performance of FedALS using ResNet-20 as a deep neural network architecture and OPT-125M as a large language model. We treat the convolutional layers of ResNet-20 as the representation extractor and the final dense layers as the model head. For OPT-125M, we consider the first 10 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 for the LLM. The experimentation was conducted on a network consisting of five nodes alongside a central server. For image classification, we utilized a 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 -4 . 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 each experiment by multiplying and dividing the learning rate by powers of two, stopping 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, we include the average and standard deviation error bars.In this section, we allocate the dataset to nodes using a noniid 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 τ = 5 and α = 10 in FedALS, it becomes apparent that aggregation and communication costs are reduced as compared to FedAvg with the same τ 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 contributes to an overall performance enhancement. Thus, our approach in FedALS is beneficial for both communication efficiency and enhancing model generalization performance simultaneously.The results for the iid setting are presented in Fig. 2b, 3b .In order to obtain these results, the data is shuffled, and then evenly divided among nodes. We note that in this situation, the performance improvement of FedALS is negligible. This is expected since there is a factor of 1 K 2 in the generalization in this case, ensuring that we will have nearly the same population risk as the empirical risk. Therefore, the deciding factor here is the optimization of the empirical risk, which is improved with a smaller τ as discussed in Remark 5.1. Thus, the improvement of the generalization error using the FedALS approach is negligible in this setup. Additionally, since SCAFFOLD requires the transmission of control variables to the central server, which is of the same size as the models themselves, it results in approximately twice the communication overhead when compared to FedAvg.Let us consider Fig. 2a , 3a to notice that in real-world deep learning situations, FedALS enhances performance 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 1 .As shown in Table 2 , it becomes evident that when we increase α from 1 (FedAvg), we initially witness an enhancement in accuracy owing to improved generalization. However, beyond a certain threshold (α = 10), further increment in α ceases to contribute to performance improvement. This is due to the adverse impact of a high number of local steps on optimization performance indicated in Remark 5.1.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.In Table 3 , we 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. 
acknowledgments: 
 This work was supported in part by ARL under Grant W911NF-2120272, and in part by NSF under Grant CCF-1942878, Grant CNS-2148182, and Grant CNS-2112471. 
conclusion: 
 In this paper, we first characterized generalization error bound for one-and R-round federated learning. One-round generalization bound is tighter than the state of the art. Based on our improved generalization bound analysis and our representation learning interpretation of this analysis, we showed for the first time that less frequent aggregations, hence more local updates, for the representation extractor (usually corresponds to initial layers) leads to the creation of more generalizable models, particularly for non-iid scenarios. This insight led us to develop the FedALS algorithm, which centers around the concept of increasing local steps for the initial layers of the deep learning model while conducting more averaging for the final layers. The experimental results demonstrated the effectiveness of FedALS in heterogeneous setups. 
 
 Appendix: We first state and prove the following lemma that will be used in the proof of Theorem 4.1. Lemma A.1 (Leave-one-out). [Expansion of Theorem 1 in (Barnes et al., 2022) ]EQUATIONProof. We haveEQUATIONAlso, observe thatEQUATIONPutting 11, and 13 together, and by the definition of the expected generalization error, we getEQUATIONIn the following lemma, we establish a fundamental generalization bound for a single round of ERM and FedAvg. (Theorem 4.1). Theorem A.2. Let l(M θ , z) be µ-strongly convex and L-smooth in M θ , M θ k = A k (S k ) represents the model obtained from Empirical Risk Minimization (ERM) algorithm on local dataset S k , i.e., M θ k = arg min M n k i=1 l(M, z k,i ), and M θ = A(S ) is the model after one round of FedAvg ( θ = E k∼K θ k ). Then, the expected generalization error,E {S k ∼D n k k } K k=1 ∆ A (S ), is bounded by E k∼K LK(k) 2 µ E {S k ∼D n k k } ∆ A k (S k ) Expected local generalization +2 L µ K(k) E {S k ∼D n k k } K k=1 δ k,A (S )Root of expected non-iidnessEQUATIONwhereδ k,A (S ) = R S k (A(S )) -R S k (A k (S k ))indicates the level of non-iidness for client k in function A on dataset S .Proof. We again considerS ′ k = (z ′ k,1 , ..., z ′ k,n k ), where z ′ k,i is sampled from D k .Let also define S (k) = (S 1 , ..., S ′ k , ..., S K ). Based on Lemma A.1, we can express the expected generalization error asEQUATION)Based on L-smoothness of l(M θ , z) in M θ , we obtainEQUATION)where ⟨•, •⟩, ∥ • ∥ 2 indicate Euclidean inner product, and squared L2-norm. Note that (17) holds due toEQUATION)We can bound expectation of the inner product term on the right-hand side of (17) using Cauchy-Schwarz inequality asEQUATIONwhere ( 19) is true because on the right we have the absolute value. (20), and ( 21) are based on Cauchy-Schwarz inequality. Now Let's find an upper bound for k) )∥ 2 that appears on the right-hand side of both (17), and (21). We obtainE {S k ,S ′ k ∼D n k k } K k=1 ∥A(S ) -A(S (EQUATIONwhere ( 22) proceeds by observing that A(S (k,i) ) varies solely in the sub-model derived from node k, diverging from A(S ), and this discrepancy is magnified by a factor of K(k) when averaging of all sub-models. ( 23) holds due to the µ-strongly convexity of l(M θ , z) in M θ which leads to µ-strongly convexity of R S k (M θ ) and the fact thatEQUATION(24), ( 25) are based on local empirical and population risk definitions.E k∼K,{S k ,S ′ k ∼D n k k } K k=1 ∥∇ 1 n k n k i=1 l(A(S (k) ), z ′ k,i)∥ 2 on the right-hand side of (21). Note thatEQUATIONEQUATION( θ = E k∼K θ k ). Then, the expected generalization error,E {S k ∼D n k k } K k=1 ∆ A (S ), is bounded byE k∼K E {S k ∼D n k k } K k=1 2LK(k) 2 µ ∆ A k (S k ) + 2ϵ k (S k ) (38) + 8L µ K(k) E {S k ∼D n k k } K k=1 δ k,A (S ) + ϵ k (S k ) E {S k ∼D n k k } K k=1 ∆ A k (S k ) + 2ϵ k (S k ) , where ϵ k (S k ) = R S k (A k (S k )) -R S k (A * (S k )).Proof. All the steps are exactly the same as in the proof of theorem A.2 except for the two steps below:First, we use a new upper bound for k) )∥ 2 that appears on the right-hand side of both (17), and (21). We haveE {S k ,S ′ k ∼D n k k } K k=1 ∥A(S ) -A(S (EQUATIONwhere in (40) A * (S k ) is the ERM on S k . ( 41) is based on the following inequality.EQUATION)Secondly, we derive a new bound forE k∼K,{S k ,S ′ k ∼D n k k } K k=1 ∥∇ 1 n k n k i=1 l(A(S (k) ), z ′ k,i )∥ 2 on the right hand side of (21). We get ∥∇ 1 n k n k i=1 l(A(S (k) ), z ′ k,i )∥ 2 ≤ 2L 1 n k n k i=1 l(A(S (k) ), z ′ k,i ) - 1 n k n k i=1 l(A * (S ′ k ), z ′ k,i ) (48) ≤ 2L R S ′ k (A(S (k) )) -R S ′ k (A * (S ′ k )) (49) ≤ 2L R S ′ k (A(S (k) )) -R S ′ k (A k (S ′ k )) + R S ′ k (A k (S ′ k )) -R S ′ k (A * (S ′ k )) (50) ≤ 2L δ k,A (S (k) ) + R S ′ k (A k (S ′ k )) -R S ′ k (A * (S ′ k )) . (51) Let's define ϵ k (S k ) = R S k (A k (S k )) -R S k (A * (S k )). Putting ( 17) into ( 16) and considering (21) we getEQUATION+ 8L µ K(k) E {S k ∼D n k k } K k=1 δ k,A (S ) + ϵ k (S k ) E {S k ∼D n k k } K k=1 ∆ A k (S k ) + 2ϵ k (S k ) ,where in (52) we have applied (46), and (51). This completes the proof. Now, we prove Theorem 4.3 as follows.Theorem B.2. Let l(M θ , z) be µ-strongly convex and L-smooth in M θ . Local models at round r are calculated by doing τ local gradient descent steps (5), and the local gradient variance is bounded by σ2 , i.e., E z∼D k ∥∇l(M θ , z) - E z∼D k ∇l(M θ , z)∥ 2 ≤ σ 2 .The aggregated model at round r, M θ r , is obtained by performing FedAvg, and the data points used in round r (i.e., Z k,r ) are sampled without replacement. The average generalization error bound is Lemma C.1. Let {x i } K i=1 denote any fixed deterministic sequence. Assume P is derived by sampling K clients with replacement based on distribution K followed by an equal probability on all sampled clients, i.e., K(k) = 1 K . ThenE S ∆ F edAvg (S )≤ 1 R R r=1 E k∼K 2LK(k) 2 µ A + 8L µ K(k)(AB) 1 2 where A = Õ C(M θ ) |Z k,r | + σ 2 µτ + L µ , B = Õ E {Z k,r } K k=1 δ k,A ({Z k,r } K k=1 ) + σ 2 µτ + L µ , ÕEQUATIONProof.EQUATIONLemma C.2. Let {x i } K i=1 denote any fixed deterministic sequence. Assume P is derived by sampling K clients without replacement uniformly at random followed by weighted probability on all sampled clients as K(k) = K(k)K K . ThenEQUATIONProof.E P E k∼ K x k = E P K K K i=1 K(i)x i = K K K i=1 E P K(i)x i = K E P K(i)x i = K 1 K k i=1 K(i)x i = E k∼K x k (68) E P E k∼ K K(k)x k = E P K 2 K2 K i=1 K 2 (i)x i = K 2 K2 K i=1 E P K 2 (i)x i = K 2 K E P K 2 (i)x i = K 2 K 1 K k i=1 K 2 (i)x i (69) = K K E k∼K K(k)x k E P E k∼ K K2 (k)x k = E P K 3 K3 K i=1 K 3 (i)x i = K 3 K3 K i=1 E P K 3 (i)x i = K 3 K2 E P K 3 (i)x i = K 3 K2 1 K k i=1 K 3 (i)x i (70) = K 2 K2 E k∼K K 2 (k)x kSo based on lemmas C.1, and C.2, it becomes evident that the expectation of the participation gap in (62) becomes zero for both two methods, i.e., E P R(A(S )) -R K(A(S )) = 0.(71)The expected generalization error, E {S k ∼D n k k } K k=1 ∆ A (S ),will be just the expectation of the second term in 62 that can be bounded using Lemma A.2 by for Node k in 1, ..., K in parallel do for Node k in 1, ..., K in parallel do If we take the expectation of (72) with respect to P, and taking into account Lemma C.1, we get the generalization bound for method 1 asEQUATIONEQUATION)For Scheme 2, we can obtain the generalization bound in the same way by taking the expectation of (72) with respect to P and considering Lemma C.2. We get the generalization bound for Method 2 as:EQUATIONIn this section, we have listed our implementation of FedAvg (Algorithm 2), SCAFFOLD (Karimireddy et al., 2019) (Algorithm 3), and FedALS+SCAFFOLD (Algorithm 4). In Algorithm 3, we observe that in addition to the model, SCAFFOLD also keeps track of a state-specific to each client, referred to as the client control variate c k,r . It is important to recognize that the clients within SCAFFOLD have memory and preserve the c k,r and K k=1 c k,r values. Additionally, when c k,r consistently remains at 0, SCAFFOLD essentially becomes equivalent to FedAvg.Algorithm 4 demonstrates the integration of FedALS and SCAFFOLD. It is important to observe that in this algorithm, the control variables are fragmented according to different model partitions as distinct local step counts exist for different parts. for Node k in 1, ..., K in parallel do  


Decision: We thank the authors for their time and contribution to WANT and we are pleased to share that after the reviewing process the paper has been accepted. Congratulations! We encourage the authors to consider reviewers' feedback for the improvement of the camera-ready version. We hope to see you in person at the workshop and brainstorm on efficient training research together!