Year: 2021

Abstract:
In Federated Learning, a common approach for aggregating local models across clients is periodic averaging of the full model parameters. It is, however, known that different layers of neural networks can have a different degree of model discrepancy across the clients. The conventional full aggregation scheme does not consider such a difference and synchronizes the whole model parameters at once, resulting in inefficient network bandwidth consumption. Aggregating the parameters that are similar across the clients does not make meaningful training progress while increasing the communication cost. We propose FedLAMA, a layer-wise model aggregation scheme for scalable Federated Learning. FedLAMA adaptively adjusts the aggregation interval in a layer-wise manner, jointly considering the model discrepancy and the communication cost. The layer-wise aggregation method enables to finely control the aggregation interval to relax the aggregation frequency without a significant impact on the model accuracy. Our empirical study shows that FedLAMA reduces the communication cost by up to 60% for IID data and 70% for non-IID data while achieving a comparable accuracy to FedAvg.

heading: 1 INTRODUCTION
text: In Federated Learning, periodic full model aggregation is the most common approach for aggregating local models across clients. Many Federated Learning algorithms, such as FedAvg (McMahan et al. (2017)), FedProx (Li et al. (2018)), FedNova (Wang et al. (2020)), and SCAFFOLD (Karimireddy et al. (2020)), assume the underlying periodic full aggregation scheme. However, it has been observed that the magnitude of gradients can be significantly different across the layers of neural networks (You et al. (2019)). That is, all the layers can have a different degree of model discrepancy. The periodic full aggregation scheme does not consider such a difference and synchronizes the entire model parameters at once. Aggregating the parameters that are similar across all the clients does not make meaningful training progress while increasing the communication cost. Considering the limited network bandwidth in usual Federated Learning environments, such an inefficient network bandwidth consumption can significantly harm the scalability of Federated Learning applications.
Many researchers have put much effort into addressing the expensive communication issue. Adaptive model aggregation methods adjust the aggregation interval to reduce the total communication cost (Wang & Joshi (2018); Haddadpour et al. (2019)). Gradient (model) compression (Alistarh et al. (2018); Albasyoni et al. (2020)), sparsification (Wangni et al. (2017); Wang et al. (2018)), low-rank approximation (Vogels et al. (2020); Wang et al. (2021)), and quantization (Alistarh et al. (2017); Wen et al. (2017); Reisizadeh et al. (2020)) techniques directly reduce the local data size. Employing heterogeneous model architectures across clients is also a communication-efficient approach (Diao et al. (2020)). While all these works effectively tackle the expensive communication issue from different angles, they commonly assume the underlying periodic full model aggregation.
To break such a convention of periodic full model aggregation, we propose FedLAMA, a novel layerwise adaptive model aggregation scheme for scalable and accurate Federated Learning. FedLAMA first prioritizes all the layers based on their contributions to the total model discrepancy. We present a metric for estimating the layer-wise degree of model discrepancy at run-time. The aggregation intervals are adjusted based on the layer-wise model discrepancy such that the layers with a smaller degree of model discrepancy is assigned with a longer aggregation interval than the other layers. The above steps are repeatedly performed once the entire model is synchronized once.
Our focus is on how to relax the model aggregation frequency at each layer, jointly considering the communication efficiency and the impact on the convergence properties of federated optimization. By adjusting the aggregation interval based on the layer-wise model discrepancy, the local models can be effectively synchronized while reducing the number of communications at each layer. The model accuracy is marginally affected since the intervals are increased only at the layers that have the smallest contribution to the total model discrepancy. Our empirical study demonstrates that FedLAMA automatically finds the interval settings that make a practical trade-off between the communication cost and the model accuracy. We also provide a theoretical convergence analysis of FedLAMA for smooth and non-convex problems under non-IID data settings.
We evaluate the performance of FedLAMA across three representative image classification benchmark datasets: CIFAR-10 (Krizhevsky et al. (2009)), CIFAR-100, and Federated Extended MNIST (Cohen et al. (2017)). Our experimental results deliver novel insights on how to aggregate the local models efficiently consuming the network bandwidth. Given a fixed number of training iterations, as the aggregation interval increases, FedLAMA reduces the communication cost by up to 60% under IID data settings and 70% under non-IID data settings, while having only a marginal accuracy drop.

heading: 2 RELATED WORKS
text: Compression Methods – The communication-efficient global model update methods can be categorized into two groups: structured update and sketched update (Konečnỳ et al. (2016)). The structured update indicates the methods that enforce a pre-defined fixed structure of the local updates, such as low-rank approximation and random mask methods. The sketched update indicates the methods that post-process the local updates via compression, sparsification, or quantization. Both directions are well studied and have shown successful results (Alistarh et al. (2018); Albasyoni et al. (2020); Wangni et al. (2017); Wang et al. (2018); Vogels et al. (2020); Wang et al. (2021); Alistarh et al. (2017); Wen et al. (2017); Reisizadeh et al. (2020)). The common principle behind these methods is that the local updates can be replaced with a different data representation with a smaller size.
These compression methods can be independently applied to our layer-wise aggregation scheme such that the each layer’s local update is compressed before being aggregated. Since our focus is on adjusting the aggregation frequency rather than changing the data representation, we do not directly compare the performance between these two approaches. We leave harmonizing the layer-wise aggregation scheme and a variety of compression methods as a promising future work.
Similarity Scores – Canonical Correlation Analysis (CCA) methods are proposed to estimate the representational similarity across different models (Raghu et al. (2017); Morcos et al. (2018)). Centered Kernel Alignment (CKA) is an improved extension of CCA (Kornblith et al. (2019)). While these methods effectively quantify the degree of similarity, they commonly require expensive computations. For example, SVCCA performs singular vector decomposition of the model and CKA computes Hilbert-Schmidt Independence Criterion multiple times (Gretton et al. (2005)). In addition, the representational similarity does not deliver any information regarding the gradient difference that is strongly related to the convergence property. We will propose a practical metric for estimating the layer-wise model discrepancy, which is cheap enough to be used at run-time.
Layer-wise Model Freezing – Layer freezing (dropping) is the representative layer-wise technique for neural network training (Brock et al. (2017); Kumar et al. (2019); Zhang & He (2020); Goutam et al. (2020)). All these methods commonly stop updating the parameters of the layers in a bottomup direction. These empirical techniques are supported by the analysis presented in (Raghu et al. (2017)). Since the layers converge from the input-side sequentially, the layer-wise freezing can reduce the training time without strongly affecting the accuracy. These previous works clearly demonstrate the advantages of processing individual layers separately.

heading: 3 BACKGROUND
text: Federated Optimization – We consider federated optimization problems as follows.
min x∈Rd
[ F (x) :=
m∑ i=1 piFi(x)
] , (1)
Algorithm 1: FedLAMA: Federated Layer-wise Adaptive Model Aggregation. Input: τ ′: base aggregation interval, φ: interval increasing factor, pi, i ∈ {1, · · · ,m}.
1 τl ← τ ′,∀l ∈ {1, · · · , L}; 2 for k = 1 to K do 3 SGD step: xik = x i k−1 − η∇f(wik−1, ξk); 4 for l = 1 to L do 5 if k mod τl is 0 then 6 Synchronize layer l: u(l,k) ← ∑m i=1 pix i (l,k);
7 dl ← ∑m i=1 ( pi‖u(l,k) − xi(l,k)‖ 2 ) /(τl(dim(u(l,k))) ;
8 if k mod φτ ′ is 0 then 9 Adjust aggregation interval at all L layers (Algorithm 2).;
10 Output: uK ;
where pi = ni/n is the ratio of local data to the total dataset, and Fi(x) = 1ni ∑ ξ∈D fi(x, ξ) is the local objective function of client i. n is the global dataset size and ni is the local dataset size.
FedAvg is a basic algorithm that solves the above minimization problem. As the degree of data heterogeneity increases, FedAvg converges more slowly. Several variants of FedAvg, such as FedProx, FedNova, and SCAFFOLD, tackle the data heterogeneity issue. All these algorithms commonly aggregate the local solutions using the periodic full aggregation scheme.
Model Discrepancy – All local SGD-based algorithms allow the clients to independently train their local models within each communication round. The variance of stochastic gradients and heterogeneous data distribution can lead the local models to different directions on parameter space during the local update steps. We formally define such a discrepancy among the models as follows.
m∑ i=1 pi‖u− xi‖2,
where m is the number of clients, u is the synchronized model, and xi is client i’s local model. This quantity bounds the difference between the local gradients and the global gradients under a smoothness assumption on objective functions.

heading: 4 LAYER-WISE ADAPTIVE MODEL AGGREGATION
text: Layer Prioritization – In theoretical analysis, it is common to assume the smoothness of objective functions such that the difference between local gradients and global gradients is bounded by a scaled difference of the corresponding sets of parameters. Motivated by this convention, we define ‘layer-wise unit model discrepancy’, a useful metric for prioritizing the layers as follows.
dl =
∑m i=1 pi‖ul − xil‖2
τl(dim(ul)) , l ∈ {1, · · · , L} (2)
where L is the number of layers, l is the layer index, u is the global parameters, xi is the client i’s local parameters, τ is the aggregation interval, and dim(·) is the number of parameters. This metric quantifies how much each parameter contributes to the model discrepancy at each iteration. The communication cost is proportional to the number of parameters. Thus, ∑m i=1 pi‖ul − xil‖2/dim(ul) shows how much model discrepancy can be eliminated by synchronizing the layer at a unit communication cost. This metric allows prioritizing the layers such that the layers with a smaller dl value has a lower priority than the others.
Adaptive Model Aggregation Algorithm – We propose FedLAMA, a layer-wise adaptive model aggregation scheme. Algorithm 1 shows FedLAMA algorithm. There are two input parameters: τ ′ is the base aggregation interval and φ is the interval increase factor. First, the parameters at layer l are synchronized across the clients after every τl iterations (line 6). Then, the proposed metric
Algorithm 2: Layer-wise Adaptive Interval Adjustment. Input: d: the observed model discrepancy at all L layers, τ ′: the base aggregation interval, φ:
the interval increasing factor. 1 Sorted model discrepancy: d̂← sort (d); 2 Sorted index of the layers: î← argsort (d); 3 Total model size: λ← ∑L l=1 dim(ul);
4 Total model discrepancy: δ ← ∑L l=1 dl ∗ dim(ul); 5 for l = 1 to L do 6 δl ← ( ∑l i=1 d̂i ∗ dim(ui))/δ;
7 λl ← ( ∑l i=1 dim(ui))/λ; 8 Find the layer index: i← îl ; 9 if δl < λl then
10 τi ← φτ ′; 11 else 12 τi ← τ ′;
13 Output: τ : the adjusted aggregation intervals at all L layers.;
dl is calculated using the synchronized parameters ul (line 7). At the end of every φτ ′ iterations, FedLAMA adjusts the model aggregation interval at all the L layers. (line 9).
Algorithm 2 finds the layers that can be less frequently aggregated making a minimal impact on the total model discrepancy. First, the layer-wise degree of model discrepancy is estimated as follows.
δl = ∑l i=1 d̂i ∗ dim(ui)∑L i=1 d̂i ∗ dim(ui) , (3)
where d̂i is the ith smallest element in the sorted list of the proposed metric d. Given l layers with the smallest dl values, δl quantifies their contribution to the total model discrepancy. Second, the communication cost impact is estimated as follows.
λl = ∑l i=1 dim(ui)∑L i=1 dim(ui)
(4)
λl is the ratio of the parameters at the l layers with the smallest dl values. Thus, 1 − λl estimates the number of parameters that will be more frequently synchronized than the others. As l increases, δl increases while 1− λl decreases monotonically. Algorithm 2 loops over the L layers finding the l value that makes δl and 1 − λl similar. In this way, it finds the aggregation interval setting that slightly sacrifices the model discrepancy while remarkably reducing the communication cost.
Figure 1 shows the δl and 1 − λl curves collected from a) CIFAR-10 (ResNet20) training and b) CIFAR-100 (Wide-ResNet28-10) training. The x-axis is the number of layers to increase the aggregation interval and the y-axis is the δl and 1− λl values. The cross point of the two curves is much lower than 0.5 on y-axis in both charts. For instance, in Figure 1.a), the two curves are crossed when x value is 9, and the corresponding y value is near 0.2. That is, when the aggregation interval is increased at those 9 layers, 20% of the total model discrepancy will increase by a factor of φ while 80% of the total communication cost will decrease by the same factor. Note that the cross points are below 0.5 since the δl and 1− λl are calculated using the dl values sorted in an increasing order. It is worth noting that FedLAMA can be easily extended to improve the convergence rate at the cost of having minor extra communications. In this work, we do not consider finding such interval settings because it can increase the latency cost, which is not desired in Federated Learning. However, in the environments where the latency cost can be ignored, such as high-performance computing platforms, FedLAMA can accelerate the convergence by adjusting the intervals based on the cross point of 1− δl and λl calculated using the list of dl values sorted in a decreasing order. Impact of Aggregation Interval Increasing Factor φ – In Federated Learning, the communication latency cost is usually not negligible, and the total number of communications strongly affects the
scalability. When increasing the aggregation interval, Algorithm 2 multiplies a pre-defined small constant φ to the fixed base interval τ ′ (line 10). This approach ensures that the communication latency cost is not increased while the network bandwidth consumption is reduced by a factor of φ.
FedAvg can be considered as a special case of FedLAMA where φ is set to 1. When φ > 1, FedLAMA less frequently synchronize a subset of layers, and it results in reducing their communication costs. When increasing the aggregation interval, FedLAMA multiplies φ to the base interval τ ′. So, it is guaranteed that the whole model parameters are fully synchronized after φτ ′ iterations. Because of the layers with the base aggregation interval τ ′, the total model discrepancy of FedLAMA after φτ ′ iterations is always smaller than that of FedAvg with an aggregation interval of φτ ′.

heading: 5 CONVERGENCE ANALYSIS

heading: 5.1 PRELIMINARIES
text: Notations – All vectors in this paper are column vectors. x ∈ Rd denotes the parameters of one local model and m is the number of clients. The stochastic gradient computed from a single training data point ξ is denoted by g(x, ξ). For convenience, we use g(x) instead. The full batch gradient is denoted by∇F (x). We use ‖·‖ and ‖·‖op to denote l2 norm and matrix operator norm, respectively. Assumptions – We analyze the convergence rate of FedLAMA under the following assumptions.
1. (Smoothness). Each local objective function is L-smooth, that is, ‖∇Fi(x) −∇Fi(y)‖ ≤ L‖x− y‖,∀i ∈ {1, · · · ,m}.
2. (Unbiased Gradient). The stochastic gradient at each client is an unbiased estimator of the local full-batch gradient: Eξ[gi(x, ξ)] = ∇Fi(x).
3. (Bounded Variance). The stochastic gradient at each client has bounded variance: Eξ[‖gi(x, ξ)−∇Fi(x)‖2 ≤ σ2],∀i ∈ {1, · · · ,m}, σ2 ≥ 0.
4. (Bounded Dissimilarity). For any sets of weights {pi ≥ 0}mi=1, ∑m i=1 pi = 1, there exist
constants β2 ≥ 1 and κ2 ≥ 0 such that ∑m i=1 pi‖∇Fi(x)‖2 ≤ β2‖ ∑m i=1 pi∇Fi(x)‖2+κ2. If local objective functions are identical to each other, β2 = 1 and κ2 = 0.

heading: 5.2 ANALYSIS
text: We begin with showing two key lemmas. All the proofs can be found in Appendix. Lemma 5.1. (Framework) Under Assumption 1 ∼ 3, if the learning rate satisfies η ≤ 12L , FedLAMA ensures
1
K K∑ k=1 E [ ‖∇F (uk)‖2 ] ≤ 2 ηK E [F (u1)− F (u∗)] + 2ηLσ2 m∑ i=1 (pi) 2
+ L2
K K∑ k=1 m∑ i=1 pi E [∥∥uk − xik∥∥2] .
(5)
Lemma 5.2. (Model Discrepancy) Under Assumption 1 ∼ 4, if the learning rate satisfies η < 1 2(τmax−1)L , FedLAMA ensures
1
K K∑ k=1 m∑ i=1 pi E [∥∥uk − xik∥∥2] ≤ 2η2(τmax − 1)σ21−A + Aκ2L2(1−A)
+ Aβ2
KL2(1−A) K∑ k=1 E [ ‖∇F (uk)‖2 ] ,
(6)
where A = 4η2(τmax − 1)2L2 and τmax is the largest averaging interval across all the layers.
Based on Lemma 5.1 and 5.2, we analyze the convergence rate of FedLAMA as follows.
Theorem 5.3. Suppose all m local models are initialized to the same point u1. Under Assumption 1 ∼ 4, if FedLAMA runs for K iterations and the learning rate satisfies η ≤
min
{ 1
2(τmax−1)L , 1 L √ 2τmax(τmax−1)(2β2+1)
} , FedLAMA ensures
E
[ 1
K K∑ i=1 ‖∇F (uk)‖2 ] ≤ 4 ηK (E [F (u1)− F (u∗)]) + 4η m∑ i=1 p2iLσ 2 (7)
+ 3η2(τmax − 1)L2σ2 + 6η2τmax(τmax − 1)L2κ2,
where u∗ indicates a local minimum and τmax is the largest averaging interval across all the layers.
Remark 1. (Linear Speedup) With a sufficiently small diminishing learning rate and a large number of training iterations, FedLAMA achieves linear speedup. If the learning rate is η =
√ m√ K
and pi = 1 m ,∀i ∈ {1, · · · ,m}, we have
E
[ 1
K K∑ i=1
‖∇F (uk)‖2 ] ≤ O ( 1√ mK ) +O (m K ) (8)
If K > m3, the first term on the right-hand side becomes dominant and it achieves linear speedup.
Remark 2. (Impact of Interval Increase Factor φ) The worst-case model discrepancy depends on the largest averaging interval across all the layers, τmax = φτ ′. The larger the interval increase factor φ, the larger the model discrepancy terms in (7). In the meantime, as φ increases, the communication frequency at the selected layers is proportionally reduced. So, φ should be appropriately tuned to effectively reduce the communication cost while not much increasing the model discrepancy.

heading: 6 EXPERIMENTS
text: Experimental Settings – We evaluate FedLAMA using three representative benchmark datasets: CIFAR-10 (ResNet20 (He et al. (2016))), CIFAR-100 (WideResNet28-10 (Zagoruyko & Komodakis (2016))), and Federated Extended MNIST (CNN (Caldas et al. (2018))). We use TensorFlow 2.4.3 for local training and MPI for model aggregation. All our experiments are conducted on 4 compute nodes each of which has 2 NVIDIA v100 GPUs.
Due to the limited compute resources, we simulate Federated Learning such that each process sequentially trains multiple models and then the models are aggregated across all the processes at once. While it provides the same classification results as the actual Federated Learning, the training time is serialized within each process. Thus, instead of wall-clock time, we consider the total communication cost calculated as follows.
C = L∑ l=1 Cl = L∑ l=1 dim(ul) ∗ κl, (9)
where κl is the total number of communications at layer l during the training.
Hyper-Parameter Settings – We use 128 clients in our experiments. The local batch size is set to 32 and the learning rate is tuned based on a grid search. For CIFAR-10 and CIFAR-100, we artificially generate heterogeneous data distributions using Dirichlet’s distribution. When using Non-IID data, we also consider partial device participation such that randomly chosen 25% of the clients participate in training at every φτ ′ iterations. We report the average accuracy across at least three separate runs.

heading: 6.1 CLASSIFICATION PERFORMANCE ANALYSIS
text: To evaluate the proposed model aggregation scheme, we keep all the other factors the same, such as optimizer, the number of clients, the degree of heterogeneity, and compare the performance across different model aggregation schemes. We compare the performance across three different model aggregation settings as follows.
• Periodic full aggregation with an interval of τ ′
• Periodic full aggregation with an interval of φτ ′
• Layer-wise adaptive aggregation with intervals of τ ′ and φ
The first setting provides the baseline communication cost, and we compare it to the other settings’ communication costs. The third setting is FedLAMA with the base aggregation interval τ ′ and the interval increase factor φ. Due to the limited space, we present a part of experimental results that deliver the key insights. More results can be found in Appendix.
Experimental Results with IID Data – We first present CIFAR-10 and CIFAR-100 classification results under IID data settings. Table 1 and 2 show the CIFAR-10 and CIFAR-100 results, respectively. Note that the learning rate is individually tuned for each setting using a grid search, and we report the best settings. In both tables, the first row shows the performance of FedAvg with a short interval τ ′ = 6. As the interval increases, FedAvg significantly loses the accuracy while the communication cost is proportionally reduced. FedLAMA achieves a comparable accuracy to FedAvg with τ ′ = 6 while its communication cost is similar to that of FedAvg with φτ ′. These results demonstrate that Algorithm 2 effectively finds the layer-wise interval settings that maximize the communication cost reduction while minimizing the model discrepancy increase.
Experimental Results with Non-IID Data – We now evaluate the performance of FedLAMA using non-IID data. FEMNIST is inherently heterogeneous such that it contains the hand-written digit pictures collected from 3, 550 different writers. We use random 10% of the writers’ training samples in our experiments. Table 3 shows the FEMNIST classification results. The base interval τ ′ is set to 10. FedAvg (φ = 1) significantly loses the accuracy as the aggregation interval increases. For example, when the interval increases from 10 to 40, the accuracy is dropped by 2.1% ∼ 2.7%. In contrast, FedLAMA maintains the accuracy when φ increases, while the communication cost is remarkably reduced. This result demonstrates that FedLAMA effectively finds the best interval setting that reduces the communication cost while maintaining the accuracy.
Table 4 and 5 show the non-IID CIFAR-10 and CIFAR-100 experimental results. We use Dirichlet’s distribution to generate heterogeneous data across all the clients. The detailed settings regarding Dirichlet’s distribution can be found in Appendix. The base aggregation interval τ ′ is set to 6. The interval increase factor φ is set to 2 for FedLAMA. Likely to the IID data experiments, we observe that the periodic full averaging significantly loses the accuracy as the model aggregation interval increases, while it has a proportionally reduced communication cost. For both datasets, FedLAMA achieves a comparable accuracy to the periodic full averaging with the interval of τ ′ while having the communication cost that is close to the periodic full averaging with the increased interval of φτ ′. Especially, FedLAMA works effectively even when the Dirichlet’s coefficient is set to 0.1. The coefficient of 0.1 represents an extremely high degree of data heterogeneity in terms of not only the number of samples per client but also the balance of the classes assigned to each client. These results imply that FedLAMA is a practical algorithm for Federated Learning applications with highly heterogeneous data distributions.

heading: 6.2 COMMUNICATION EFFICIENCY ANALYSIS
text: We analyze the total number of communications and the accumulated data size to evaluate the communication efficiency of FedLAMA. Figure 2 shows the total number of communications at the individual layers. The τ ′ is set to 6 and φ is 2 for FedLAMA. The key insight is that FedLAMA increases the aggregation interval mostly at the output-side large layers. This means the dl value shown in Equation (2) at the these layers are smaller than the others. Since these large layers take up most of the total model parameters, the communication cost is remarkably reduced when their aggregation intervals are increased. Figure 3 shows the layer-wise local data size shown in Equation 9. FedLAMA shows the significantly smaller total data size than FedAvg. The extra computational cost of FedLAMA is almost negligible since it calculates dl after each communication round only. Therefore, given the virtually same computational cost, FedLAMA aggregates the local models at a cheaper communication cost, and thus it improves the scalablity of Federated Learning.
We found that the amount of the reduced communication cost was not strongly affected by the degree of data heterogeneity. As shown in Table 4 and 5, the reduced communication cost is similar across different Dirichlet’s coefficients and device participation ratios. That is, FedLAMA can be considered as an effective model aggregation scheme regardless of the degree of data heterogeneity.

heading: 7 CONCLUSION
text: We proposed a layer-wise model aggregation scheme that adaptively adjusts the model aggregation interval at run-time. Breaking the convention of aggregating the whole model parameters at once, this novel model aggregation scheme introduces a flexible communication strategy for scalable Federated Learning. Furthermore, we provide a solid convergence guarantee of FedLAMA under the assumptions on the non-convex objective functions and the non-IID data distribution. Our empirical study also demonstrates the efficacy of FedLAMA for scalable and accurate Federated Learning.
Harmonizing FedLAMA with other advanced optimizers, gradient compression, and low-rank approximation methods is a promising future work.

heading: 8 CODE OF ETHICS
text: Our work does not deliver potentially harmful insights or conflicts of interests. We also do not find any potential inappropriate application or privacy/security issues. The datasets we used in our study are all public benchmark datasets, and our source code will be opened once the paper is accepted.

heading: 9 REPRODUCIBILITY STATEMENT
text: The software versions, implementation details, hyper-parameter settings can be found in the first two paragraphs of Section 6. The entire source code used in our experiments will be published as an open source once the paper is accepted. We believe one can exactly reproduce our experimental results following the provided descriptions.

heading: A APPENDIX

heading: A.1 CONVERGENCE ANALYSIS
text: Herein, we provide the proofs of the lemmas and theorem shown in Section 5.

heading: A.1.1 PRELIMINARIES
text: FedLAMA periodically chooses a few layers that will be less frequently synchronized. We call these layers Least Critical Layers (LCL) for short.
Notations – All vectors in this paper are column vectors. x ∈ Rd denotes the parameters of one local model and m is the number of workers. The stochastic gradient computed from a single training data point ξ is denoted by g(x, ξ). For convenience, we use g(x) instead. The full batch gradient is denoted by∇F (x). We use ‖ · ‖ and ‖ · ‖op to denote l2 norm and matrix operator norm, respectively.
Objective Function – In this paper, we consider federated optimization problems as follows.
min x∈Rd
[ F (x) :=
m∑ i=1 piFi(x)
] , (10)
where pi = ni/n is the ratio of local data to the total dataset, and Fi(x) = 1ni ∑ ξ∈D fi(x, ξ) is the local objective function of client i. n is the global dataset size and ni is the local dataset size. Note that, by definition, ∑m i=1 pi = 1.
Averaging Matrix – We define a time-varying averaging matrix Wk ∈ Rdm×dm as follows.
Wk =  P, if k mod τmin is 0 J, if k mod τmax is 0 I, otherwise
(11)
I is an identity matrix, P is also a time-varying averaging matrix, and J is a full averaging matrix. First, P1i is a d× d diagonal matrix that has 1 for the diagonal elements that correspond to the LCL parameters and pi for all the other diagonal elements. Likewise, P0i is another d×d diagonal matrix that has 0 for the diagonal elements that correspond to the LCL parameters and pi for all the other diagonal elements. Then, P is defined as follows.
P = { P1, for m diagonal blocks P0, for all the other blocks
(12)
The ith block column of P consists of P1i and P 0 i following the above definition.
Here we present an example of P wherem = 2 and d = 2. In this example, p0 = 1/3 and p1 = 2/3. Saying the LCL is the second parameter, P is defined as follows.
P10 = [ 1 3 0 0 1 ] ,P00 = [ 1 3 0 0 0 ] ,P11 = [ 2 3 0 0 1 ] ,P01 = [ 2 3 0 0 0 ] (13)
P =
[ P10 P 0 1
P00 P 1 1
] =  1 3 0 2 3 0
0 1 0 0 1 3 0 2 3 0 0 0 0 1  . (14) The full-averaging matrix J is defined as follows. First, Ji is a d× d diagonal matrix that has pi for the diagonal elements. Then, J consists of m×m blocks of Ji such that each column block is m of Ji blocks. Here we present an example of J where m = 2 and d = 2 as follows.
J0 = [ 1 3 0 0 13 ] ,J1 = [ 2 3 0 0 23 ] (15)
J = [ J0 J1 J0 J1 ] =  1 3 0 2 3 0 0 13 0 2 3
1 3 0 2 3 0 0 13 0 2 3  . (16) The averaging matrix P and J have the following properties:
1. P1dm = 1dm, J1dm = 1dm. 2. The product of any two averaging matrices consists only of diagonal block matrices because
all the blocks in P and J are diagonal. 3. PJ = JP = J regardless of which layers are chosen as the LCL. 4. PP = P regardless of which layers are chosen as the LCL.
Vectorization – We define a vectorized form of m local model parameters xk ∈ Rdm, its stochastic gradients gk ∈ Rdm, and the full gradients fk ∈ Rdm as follows
xk = vec { x1k,x 2 k, · · · ,xmk } gk = vec { g1(x 1 k), g2(x 2 k), · · · , gm(xmk )
} fk = vec { ∇F1(x1k),∇F2(x2k), · · · ,∇Fm(xmk ) } .
(17)
The full model aggregation can be written using the vectorized form of local models xk and the averaging matrix J as follows.
Jxk =  1 3 0 2 3 0 0 13 0 2 3
1 3 0 2 3 0 0 13 0 2 3
  x (1,1) k x (1,2) k
x (2,1) k x (2,2) k
 =  (x (1,1) k + 2x (2,1) k )/3 (x (1,2) k + 2x (2,2) k )/3
(x (1,1) k + 2x (2,1) k )/3 (x (1,2) k + 2x (2,2) k )/3  (18) where x(i,j)k is the j th model parameter of local model i at iteration k.
We also define the following additional vectorized forms of the weighted model parameters and gradients for convenience.
x̂k = vec {√ p1x 1 k, √ p2x 2 k, · · · , √ pmx m k } ĝk = vec {√ p1g1(x 1 k), √ p2g2(x 2 k), · · · , √ pmgm(x m k ) }
f̂k = vec {√ p1∇F1(x1k), √ p2∇F2(x2k), · · · , √ pm∇Fm(xmk ) } (19) Assumptions – We analyze the convergence rate of FedLAMA under the following assumptions.
1. (Smoothness). Each local objective function is L-smooth, that is, ‖∇Fi(x) −∇Fi(y)‖ ≤ L‖x− y‖,∀i ∈ {1, · · · ,m}.
2. (Unbiased Gradient). The stochastic gradient at each client is an unbiased estimator of the local full-batch gradient: Eξ[gi(x, ξ)] = ∇Fi(x).
3. (Bounded Variance). The stochastic gradient at each client has bounded variance: Eξ[‖gi(x, ξ)−∇Fi(x)‖2 ≤ σ2],∀i ∈ {1, · · · ,m}, σ2 ≥ 0.
4. (Bounded Dissimilarity). For any sets of weights {pi ≥ 0}mi=1, ∑m i=1 pi = 1, there exist
constants β2 ≥ 1 and κ2 ≥ 0 such that ∑m i=1 pi‖∇Fi(x)‖2 ≤ β2‖ ∑m i=1 pi∇Fi(x)‖2+κ2. If local objective functions are identical to each other, β2 = 1 and κ2 = 0.

heading: A.1.2 PROOFS
text: Theorem 5.1. Suppose all m local models are initialized to the same point u1. Under Assumption 1 ∼ 4, if FedLAMA runs for K iterations and the learning rate satisfies η ≤
min
{ 1
2(τmax−1)L , 1 L √ 2τmax(τmax−1)(2β2+1)
} , FedLAMA ensures
E
[ 1
K K∑ i=1 ‖∇F (uk)‖2 ] ≤ 4 ηK (E [F (u1)− F (u∗)]) + 4η m∑ i=1 p2iLσ 2
+ 3η2(τmax − 1)L2σ2 + 6η2τmax(τmax − 1)L2κ2, (20) where u∗ indicates a local minimum.
Proof. Based on Lemma 5.1 and 5.2, we have
1
K K∑ k=1 E [ ‖∇F (uk)‖2 ] ≤ 2 ηK (E [F (u1)− F (u∗)]) + 2η m∑ i=1 p2iLσ 2
+ L2
( η2(τmax − 1)σ2j
1−A +
Aβ2
KL2(1−A) K∑ k=1 E [ ‖∇F (uk)‖2 ] +
Aκ2
L2(1−A)
) .
After re-writing the left-hand side and a minor rearrangement, we have
1
K K∑ k=1 E [ ‖∇F (uk)‖2 ] ≤ 2 ηK (E [F (u1)− F (u∗)]) + 2η m∑ i=1 p2iLσ 2
+ 1
K K∑ k=1 Aβ2 1−A E [ ‖∇F (uk)‖2 ] + L2 ( η2(τmax − 1)σ2
1−A +
Aκ2
L2(1−A)
) .
By moving the third term on the right-hand side to the left-hand side, we have
1
K K∑ k=1 ( 1− Aβ 2 1−A ) E [ ‖∇jF (uk)‖2 ] ≤ 2 ηK (E [F (u1)− F (u∗)]) + 2η m∑ i=1 p2iLσ 2
+ L2 ( η2(τmax − 1)σ2
1−A +
Aκ2
L2(1−A)
) . (21)
If A ≤ 12β2+1 , then Aβ2 1−A ≤ 1 2 . Therefore, (21) can be simplified as follows.
1
K K∑ k=1 E [ ‖∇F (uk)‖2 ] ≤ 4 ηK (E [F (u1)− F (u∗)]) + 4η m∑ i=1 p2iLσ 2 (22)
+ 2L2 ( η2(τmax − 1)σ2
1−A
) + 2 Aκ2
1−A .
The learning rate condition A ≤ 12β2+1 also ensures that 1 1−A ≤ 1 + 1 2β2 . Based on Assumption 4, 1 2β2 ≤ 2 3 , and thus 1 1−A ≤ 2 3 . Therefore, we have
1
K K∑ k=1 E [ ‖∇F (uk)‖2 ] ≤ 4 ηK (E [F (u1)− F (u∗)]) + 4η m∑ i=1 p2iLσ 2
+ 3η2(τmax − 1)L2σ2 + 6η2τmax(τmax − 1)L2κ2. We complete the proof.
Learning Rate Constraints – In Theorem 5.3, we have two learning rate constraints, one from (22) and the other from (51) as follows.
A < 1
2β2 + 1 from (22)
A < 1 from (51) After a minor rearrangement, we have a unified learning rate constraint as follows.
η ≤ min
{ 1
2(τmax − 1)L ,
1 L √ 2τmax(τmax − 1)(2β2 + 1)
}
Lemma 5.1. (Framework) Under Assumption 1 ∼ 3, if the learning rate satisfies η ≤ 12L , FedLAMA ensures
1
K K∑ k=1 E [ ‖∇F (uk)‖2 ] ≤ 2 ηK E [F (u1)− F (u∗)] + 2ηLσ2 m∑ i=1 (pi) 2
+ L2
K K∑ k=1 m∑ i=1 pi E [∥∥uk − xik∥∥2] .
(23)
Proof. Based on Assumption 1, we have
E [F (uk+1)− F (uk)] ≤ −ηE [ 〈∇F (uk),
m∑ i=1 pigi(x i k)〉 ] ︸ ︷︷ ︸
T1
+ η2L
2 E ∥∥∥∥∥ m∑ i=1 pigi(x i k) ∥∥∥∥∥ 2 
︸ ︷︷ ︸ T2
(24)
First, T1 can be rewritten as follows.
T1 = E [ 〈∇F (uk),
m∑ i=1 pi ( gi(x i k)−∇Fi(xik) ) 〉
] + E [ 〈∇F (uk),
m∑ i=1 pi∇Fi(xik)〉
]
= E [ 〈∇F (uk),
m∑ i=1 pi∇Fi(xik)〉
]
= 1
2 ‖∇F (uk)‖2 +
1 2 E ∥∥∥∥∥ m∑ i=1 pi∇Fi(xik) ∥∥∥∥∥ 2 − 1 2 E ∥∥∥∥∥∇F (uk)− m∑ i=1 pi∇Fi(xik) ∥∥∥∥∥ 2  ,
(25) where the last equality holds based on a basic equality: 2a>b = ‖a‖2 + ‖b‖2 − ‖a− b‖2 . Then, T2 can be bounded as follows.
T2 = E ∥∥∥∥∥ m∑ i=1 pi ( gi(x i k)− E [ gi(x i k) ]) + m∑ i=1 pi E [ gi(x i k) ]∥∥∥∥∥ 2 
= E ∥∥∥∥∥ m∑ i=1 pi ( gi(x i k)−∇Fi(xik) ) + m∑ i=1 pi∇Fi(xik) ∥∥∥∥∥ 2 
≤ 2E ∥∥∥∥∥ m∑ i=1 pi ( gi(x i k)−∇Fi(xik) )∥∥∥∥∥ 2 + 2E ∥∥∥∥∥ m∑ i=1 pi∇Fi(xik) ∥∥∥∥∥ 2 
= 2 m∑ i=1 p2i E [∥∥gi(xik)−∇Fi(xik)∥∥2]+ 2E ∥∥∥∥∥ m∑ i=1 pi∇Fi(xik) ∥∥∥∥∥ 2  ≤ 2σ2 m∑ i=1 p2i + 2E ∥∥∥∥∥ m∑ i=1 pi∇Fi(xik) ∥∥∥∥∥ 2  , (26)
where the last equality holds because gi(xik) − ∇Fi(xik) has 0 mean and is independent across i, and the last inequality follows Assumption 3.
By plugging in (25) and (26) into (24), we have the following.
E [F (uk+1)− F (uk)] ≤ − η
2 ‖∇F (uk)‖2 −
η 2 E ∥∥∥∥∥ m∑ i=1 pi∇Fi(xik) ∥∥∥∥∥ 2 
+ η
2 E ∥∥∥∥∥∇F (uk)− m∑ i=1 pi∇Fi(xik) ∥∥∥∥∥ 2 + η2Lσ2 m∑ i=1 p2i
+ η2LE ∥∥∥∥∥ m∑ i=1 pi∇Fi(xik) ∥∥∥∥∥ 2 
= −η 2 ‖∇F (uk)‖2 − η 2 (1− 2ηL)E ∥∥∥∥∥ m∑ i=1 pi∇Fi(xik) ∥∥∥∥∥ 2 
+ η
2 E ∥∥∥∥∥∇F (uk)− m∑ i=1 pi∇Fi(xik) ∥∥∥∥∥ 2 + η2Lσ2 m∑ i=1 p2i
If η ≤ 12L , it follows
E [F (uk+1)− F (uk)] η ≤ −1 2 ‖∇F (uk)‖2 + ηLσ2 m∑ i=1 p2i
+ 1
2 E ∥∥∥∥∥∇F (uk)− m∑ i=1 pi∇Fi(xik) ∥∥∥∥∥ 2 
≤ −1 2 ‖∇F (uk)‖2 + ηLσ2 m∑ i=1 p2i (27)
+ 1
2 m∑ i=1 pi E [∥∥∇Fi(uk)−∇Fi(xik)∥∥2]
≤ −1 2 ‖∇F (uk)‖2 + ηLσ2 m∑ i=1 p2i + L2 2 m∑ i=1 pi E [∥∥uk − xik∥∥2] ,
where (27) holds based on the convexity of `2 norm and Jensen’s inequality.
By taking expectation and averaging across K iterations, we have.
1
K K∑ k=1 E [F (uk+1)− F (uk)] η ≤ − 1 2K K∑ k=1 ‖∇F (uk)‖2 + ηLσ2 m∑ i=1 p2i
+ L2
2K K∑ k−1 m∑ i=1 pi E [∥∥uk − xik∥∥2] .
After a minor rearrangement, we have a telescoping sum as follows.
1
K K∑ k=1 E [ ‖∇F (uk)‖2 ] ≤ 2 ηK E [F (u1)− F (uk+1)] + 2ηLσ2 m∑ i=1 p2i
+ L2
K K∑ k=1 m∑ i=1 pi E [∥∥uk − xik∥∥2]
≤ 2 ηK E [F (u1)− F (u∗)] + 2ηLσ2 m∑ i=1 p2i
+ L2
K K∑ k=1 m∑ i=1 pi E [∥∥uk − xik∥∥2] ,
where u∗ indicates the local minimum. Here, we complete the proof.
Lemma 5.2. (Model Discrepancy) Under Assumption 1 ∼ 4, if the learning rate satisfies η < 1 2(τmax−1)L , FedLAMA ensures
1
K K∑ k=1 m∑ i=1 pi E [∥∥uk − xik∥∥2] ≤ 2η2(τmax − 1)σ21−A + Aκ2L2(1−A)
+ Aβ2
KL2(1−A) K∑ k=1 E [ ‖∇F (uk)‖2 ] ,
(28)
where A = 4η2(τmax − 1)2L2 and τmax is the largest averaging interval across all the layers.
Proof. We begin with rewriting the weighted average of the squared distance using the vectorized form of the local models as follows.
m∑ i=1 pi ∥∥uk − xik∥∥2 = m∑ i=1 ∥∥√pi (uk − xik)∥∥2 = ‖Jx̂k − x̂k‖2 (29) = ‖(J− I)x̂k‖2 ,
where (29) holds by the commutative property of multiplication.
Then, according to the parameter update rule, we have
(J− I)x̂k = (J− I)Wk−1(x̂k−1 − ηĝk−1) = (J− I)Wk−1x̂k−1 − (J−Wk−1)ηĝk−1, (30)
where (30) holds because JW = J based on the averaging matrix property 3, and IW = W.
Then, expanding the expression of xk−1, we have
(J− I)x̂k = (J− I)Wk−1(Wk−2(x̂k−2 − ηĝk−2))− (J−Wk−1)ηĝk−1 = (J− I)Wk−1Wk−2x̂k−2 − (J−Wk−1Wk−2)ηĝk−2 − (J−Wk−1)ηĝk−1.
Repeating the same procedure for x̂k−2, x̂k−3, · · · , x̂2, we have
(J− I)x̂k = (J− I) k−1∏ s=1 Wsx̂1 − η k−1∑ s=1 (J− k−1∏ l=s Wl)ĝs
= −η k−1∑ s=1 (J− k−1∏ l=s Wl)ĝs, (31)
where (31) holds because xi1 is the same across all the workers and thus (J− I)x̂1 = 0.
Based on (31), we have
1
K K∑ k=1 m∑ i=1 pi E [∥∥uk − xik∥∥2]
= 1
K K∑ k=1 ( E [ ‖(J− I)x̂k‖2 ])
= 1
K K∑ k=1 η2 E ∥∥∥∥∥ k−1∑ s=1 (J− k−1∏ l=s Wl)ĝs ∥∥∥∥∥ 2 
= 1
K K∑ k=1 η2 E ∥∥∥∥∥ k−1∑ s=1 (J− k−1∏ l=s Wl)(ĝs − f̂s) + k−1∑ s=1 (J− k−1∏ l=s Wl)f̂s ∥∥∥∥∥ 2 
≤ 2η 2
K  K∑ k=1 E ∥∥∥∥∥ k−1∑ s=1 (J− k−1∏ l=s Wl)(ĝs − f̂s) ∥∥∥∥∥ 2 
︸ ︷︷ ︸ T3
+ K∑ k=1 E ∥∥∥∥∥ k−1∑ s=1 (J− k−1∏ l=s Wl)f̂s ∥∥∥∥∥ 2 
︸ ︷︷ ︸ T4  (32)
where (32) holds based on the convexity of `2 norm and Jensen’s inequality. Now, we focus on bounding T3 and T4, separately.

heading: Bounding T3
text: K∑ k=1 E ∥∥∥∥∥ k−1∑ s=1 (J− k−1∏ l=s Wl)(ĝs − f̂s) ∥∥∥∥∥ 2 
= K∑ k=1 k−1∑ s=1 E ∥∥∥∥∥(J− k−1∏ l=s Wl)(ĝs − f̂s) ∥∥∥∥∥ 2  (33) ≤ K∑ k=1 k−1∑ s=1 E ∥∥∥(ĝs − f̂s)∥∥∥2 ∥∥∥∥∥(J− k−1∏ l=s Wl) ∥∥∥∥∥ 2
op
 , (34)
where (33) holds because ĝs − f̂s has 0 mean and independent across s, and (34) holds based on Lemma A.1.
Without loss of generality, we replace k with aτmax+ b, where a is the communication round index and b is the iteration index within each communication round. Then, we have
K/τmax−1∑ a=0 τmax∑ b=1 aτmax+b−1∑ s=1 E ∥∥∥(ĝs − f̂s)∥∥∥2 ∥∥∥∥∥(J− k−1∏ l=s Wl) ∥∥∥∥∥ 2
op  =
K/τmax−1∑ a=0 τmax∑ b=1 aτ∑ s=1 E ∥∥∥(ĝs − f̂s)∥∥∥2 ∥∥∥∥∥(J− aτmax+b−1∏ l=s Wl) ∥∥∥∥∥ 2
op  +
K/τmax−1∑ a=0 τmax∑ b=1 aτmax+b−1∑ s=aτmax+1 E ∥∥∥(ĝs − f̂s)∥∥∥2 ∥∥∥∥∥(J− aτmax+b−1∏ l=s Wl) ∥∥∥∥∥ 2
op  =
K/τmax−1∑ a=0 τmax∑ b=1 aτmax+b−1∑ s=aτmax+1 E ∥∥∥(ĝs − f̂s)∥∥∥2 ∥∥∥∥∥(J− aτmax+b−1∏ l=s Wl) ∥∥∥∥∥ 2
op  (35) =
K/τmax−1∑ a=0 τmax∑ b=1 aτ+b−1∑ s=aτmax+1 E [∥∥∥(ĝs − f̂s)∥∥∥2] (36)
= K/τmax−1∑ a=0 τmax∑ b=1 aτmax+b−1∑ s=aτmax+1 m∑ i=1 pi E [∥∥(gi(xis)−∇Fi(xis))∥∥2] ≤ K/τmax−1∑
a=0 τmax∑ b=1 aτmax+b−1∑ s=aτmax+1 m∑ i=1 piσ 2 (37)
= K/τmax−1∑ a=0 τmax∑ b=1 (b− 1)σ2 = K/τmax−1∑ a=0 τmax(τmax − 1) 2 σ2
≤ K (τmax − 1) 2 σ2. (38)
Remember FedLAMA synchronizes the whole parameters at least once after every τmax iterations. Thus, (35) holds because ∏aτmax+b−1 l=s Wl is J when s ≤ aτmax, and thus J − ∏aτmax+b−1 l=s Wl becomes 0. (36) holds based on Lemma A.2. (37) holds based on Assumption 3.
Bounding T4
K−τmax∑ k=1 E ∥∥∥∥∥ k−1∑ s=1 (J− k−1∏ l=s Wl)f̂s ∥∥∥∥∥ 2 
= K/τmax−1∑ a=0 τmax∑ b=1 E ∥∥∥∥∥ aτ+b−1∑ s=1 (J− aτmax+b−1∏ l=s Wl)f̂s ∥∥∥∥∥ 2 
= K/τmax−1∑ a=0 τmax∑ b=1 E ∥∥∥∥∥ aτmax+b−1∑ s=aτmax+1 (J− aτmax+b−1∏ l=s Pl)f̂s ∥∥∥∥∥ 2  (39) ≤ K/τmax−1∑
a=0 τmax∑ b=1 (b− 1) aτmax+b−1∑ s=aτmax+1 E ∥∥∥∥∥(J− aτmax+b−1∏ l=s Pl)f̂s ∥∥∥∥∥ 2  (40)
≤ K/τmax−1∑
a=0 τmax∑ b=1 (b− 1) aτmax+b−1∑ s=aτmax+1 E ∥∥∥f̂s∥∥∥2 ∥∥∥∥∥(J− aτmax+b−1∏ l=s Pl) ∥∥∥∥∥ 2
op  (41) ≤ K/τmax−1∑
a=0 τmax∑ b=1
( (b− 1)
aτmax+b−1∑ s=aτmax+1
E [∥∥∥f̂s∥∥∥2]) (42)
≤ τmax(τmax − 1) 2 K/τmax−1∑ a=0 ( aτmax+τmax−1∑ s=aτmax+1 E [∥∥∥f̂s∥∥∥2])
≤ τmax(τmax − 1) 2 K∑ k=1 E [∥∥∥f̂k∥∥∥2]
= τmax(τmax − 1)
2
K∑ k=1 m∑ i=1 pi E [∥∥∇Fi(xik)∥∥2] , (43)
where (39) holds because J − ∏aτmax+b−1 l=s Pl becomes 0 when s ≤ aτ . (40) holds based on the convexity of `2 norm and Jensen’s inequality. (41) holds based on Lemma A.1. (42) holds based on Lemma A.2.

heading: Final Result
text: By plugging in (38) and (43) into (32), we have
1
K K∑ k=1 m∑ i=1 pi E [∥∥uk − xik∥∥2]
≤ 2η 2
K
( K
(τmax − 1) 2 σ2 + τmax(τmax − 1) 2 ( K∑ k=1 m∑ i=1 pi E [∥∥∇Fi(xik)∥∥2] ))
= η2(τmax − 1)σ2 + η2τmax(τmax − 1)
K
( K∑ k=1 m∑ i=1 pi E [∥∥∇Fi(xik)∥∥2] ) (44)
The local gradient term on the right-hand side in (44) can be rewritten using the following inequality. E [∥∥∇Fi(xik)∥∥2] = E [∥∥∇Fi(xik)−∇Fi(uk) +∇Fi(uk)∥∥2]
≤ 2E [∥∥∇Fi(xik)−∇Fi(uk)∥∥2]+ 2E [‖∇Fi(uk)‖2] (45)
≤ 2L2 E [∥∥uk − xik∥∥2]+ 2E [‖∇Fi(uk)‖2] , (46)
where (45) holds based on the convexity of `2 norm and Jensen’s inequality.
Plugging in (46) into (44), we have
1
K K∑ k=1 m∑ i=1 pi E [∥∥uk − xik∥∥2]
≤ η2(τmax − 1)σ2 + 2η2τmax(τmax − 1)L2
K
K∑ k=1 m∑ i=1 pi E [∥∥uk − xik∥∥2]
+ 2η2τmax(τmax − 1)
K
K∑ k=1 m∑ i=1 pi E [ ‖∇Fi(uk)‖2 ] (47)
After a minor rearranging, we have
1
K K∑ k=1 m∑ i=1 pi E [∥∥uk − xik∥∥2] ≤ η2(τmax − 1)σ21− 2η2τmax(τmax − 1)L2
+ 2η2τmax(τmax − 1)
K(1− 2η2τmax(τmax − 1)L2) K∑ k=1 m∑ i=1 pi E [ ‖∇Fi(uk)‖2 ] (48)
Let us define A = 2η2τmax(τmax − 1)L2. Then (48) is simplified as follows.
1
K K∑ k=1 m∑ i=1 pi E [∥∥uk − xik∥∥2]
≤ η 2(τmax − 1)σ2
1−A +
A
KL2(1−A) K∑ k=1 m∑ i=1 pi E [ ‖∇Fi(uk)‖2 ] Based on Assumption 4, we have
1
K K∑ k=1 m∑ i=1 pi E [∥∥uk − xik∥∥2]
≤ η 2(τmax − 1)σ2
1−A +
Aβ2
KL2(1−A) K∑ k=1 E ∥∥∥∥∥ m∑ i=1 pi∇Fi(uk) ∥∥∥∥∥ 2 + Aκ2 L2(1−A) (49)
= η2(τmax − 1)σ2
1−A +
Aβ2
KL2(1−A) K∑ k=1 E [ ‖∇F (uk)‖2 ] +
Aκ2
L2(1−A) , (50)
where (50) holds based on the definition of the objective function (10).
Note that (49) is true only when 1− A > 0. Thus, after a minor rearrangement, we have a learning rate constraint as follows.
η < 1
2(τmax − 1)L (51)
Here, we complete the proof.

heading: A.1.3 PROOF OF OTHER LEMMAS
text: Lemma A.1. Consider a real matrix A ∈ Rmdj×mdj and a real vector b ∈ Rmdj . If b 6= 0mdj , we have
‖Ab‖ ≤ ‖A‖op‖b‖ (52)
Proof.
‖Ab‖2 = ‖Ab‖ 2
‖b‖2 ‖b‖2
≤ ‖A‖2op‖b‖2 (53)
where (53) holds based on the definition of operator norm.
Lemma A.2. Suppose an md×md averaging matrix P and the full-averaging matrix J, then
‖J−P‖2op = 1. (54)
regardless of which layers are chosen as the LCL.
Proof. First, by the definition of averaging matrix P, all the columns that do not correspond to the LCL are zeroed out in J−P. Then, based on the averaging matrix property 1 and 2, the remaining columns in P has 1 at all different rows. By the definition of J, all the non-zero elements in ith column are the same pi, i ∈ {1, · · · ,m}. Consequently, the remaining columns in J − P are always orthogonal regardless of which layers are chosen as the LCL, and thus the eigenvalues of J − P are either 1 or −1. Finally, by the definition of the matrix operator norm, ‖J − P‖2op = max{|λ(J−P)|} = 1, where λ(·) indicates the eigenvalues of the input matrix.

heading: A.2 ADDITIONAL EXPERIMENTAL RESULTS
text: In this section, we provide extra experimental results with extensive hyper-parameter settings. We commonly use 128 clients and a local batch size of 32 in all the experiments. The gradual learning rate warmup (Goyal et al. (2017)) is also applied to the first 10 epochs in all the experiments. Overall, the learning curve charts and the validation accuracy tables deliver the key insight that FedLAMA achieves a comparable convergence speed to the periodic full aggregation with the base interval
(τ ’) while having the communication cost that is similar to the periodic full aggregation with the increased interval (φτ ′).
Artificial Data Heterogeneity – For CIFAR-10 and CIFAR-100, we artificially generate the heterogeneous data distribution using Dirichlet’s distribution. The concentration coefficient α is set to 0.1, 0.5, and 1.0 to evaluate the performance of FedLAMA across a variety of degree of data heterogeneity. Note that the small concentration coefficient represents the highly heterogeneous numbers of local samples across clients as well as the balance of the samples across the labels. We used the data distribution source code provided by FedML (He et al. (2020)).
CIFAR-10 – Figure 4 shows the full learning curves for IID and non-IID CIFAR-10 datasets. The hyper-parameter settings correspond to Table 4 and 1. First, as the aggregation interval increases from 6 to 24, FedAvg suffers from the slower convergence, and it results in achieving a lower validation accuracy, regardless of the data distribution. In contrast, FedLAMA learning curves are marginally affected by the increased aggregation interval. Table 6 and 7 show the CIFAR-10 classification performance of FedLAMA across different φ settings. As expected, the accuracy is reduced as φ increases. The IID and non-IID data settings show the common trend. Depending on the system network bandwidth, φ can be tuned to be an appropriate value. When φ = 2, the accuracy is almost the same as or even slightly higher than FedAvg accuracy. If the network bandwidth is limited, one can increase φ and slightly increase the epoch budget to achieve a good accuracy. Table 8 shows the CIFAR-10 accuracy across different τ ′ settings. We see that the accuracy is significantly dropped as τ ′ increases.
CIFAR-100 – Figure 5 shows the learning curves for IID and non-IID CIFAR-100 datasets. Likely to CIFAR-10 results, FedAvg learning curves are strongly affected as the aggregation interval increases from 6 to 24 while FedLAMA learning curves are not strongly affected. Table 9 and 10 show the CIFAR-100 classification performance of FedLAMA across different φ settings. FedLAMA achieves a comparable accuracy to FedAvg with a short aggregation interval, even when the degree of data heterogeneity is extreamly high (25% device sampling and Direchlet’s coefficient of 0.1). Table 11 shows the FedAvg accuracy with different τ ′ settings. Under the strongly heterogeneous data distributions, FedAvg with a large aggregation interval (τ ≥ 12) do not achieve a reasonable accuracy.
FEMNIST – Figure 6 shows the learning curves of CNN training. Likely to the previous two datasets, the periodic full aggregation suffers from the slower convergence as the aggregation interval increases. FedLAMA learning curves are not much affected by the increased aggregation interval, and it results in achieving a higher validation accuracy after the same number of iterations. Table 12 shows the FEMNIST classification performance of FedLAMA across different φ settings. FedLAMA achieves a similar accuracy to the baseline (FedAvg with τ ′ = 10) even when using a large interval increase factor φ ≥ 4. These results demonstrate the effectiveness of the proposed layer-wise adaptive model aggregation method on the problems with heterogeneous data distributions.

Review summaries:
This paper developed an adaptive aggregation method for federated learning. The theoretical analysis shows how the interval affects the convergence rate. The experiments show that it can reduce the communication cost. 

In this paper, the authors propose a layer-wise model aggregation scheme in federated learning cases to reduce the communication cost. Specifically, they quantified the model discrepancy between local models and global models and adaptively adjusted the aggregation interval in a layer-wise manner. By increasing the aggregation intervals and relaxing the aggregation frequency, the method can reduce the communication cost in federated learning cases. The experimental results show it can reduce the communication cost for IID data and non-IID data compared to FedAvg.