L AYER - WISE A DAPTIVE M ODEL A GGREGATION FOR S CALABLE F EDERATED L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
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 inefﬁcient 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 ﬁnely control the aggregation interval to relax the aggregation frequency without a signiﬁcant 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.
I NTRODUCTION
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.
, FedProx (Li et al.
, 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 signiﬁcantly different across the layers of neural networks (You et al.
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 inefﬁcient network bandwidth consumption can signiﬁcantly 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.
Gradient (model) compression (Alistarh et al.
; Albasyoni et al.
(2020)), sparsiﬁcation (Wangni et al.
; Wang et al.
(2018)), low-rank approximation (Vogels et al.
; Wang et al.
(2021)), and quantization (Alistarh et al.
(2017); Wen et al.
; Reisizadeh et al.
techniques directly reduce the local data size.
Employing heterogeneous model architectures across clients is also a communication-efﬁcient approach (Diao et al.
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 ﬁrst 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 efﬁciency 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 ﬁnds 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 classiﬁcation benchmark datasets: CIFAR-10 (Krizhevsky et al.
(2009)), CIFAR-100, and Federated Extended MNIST (Cohen et al.
Our experimental results deliver novel insights on how to aggregate the local models efﬁciently consuming the network bandwidth.
Given a ﬁxed 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.
R ELATED W ORKS
Compression Methods – The communication-efﬁcient global model update methods can be categorized into two groups: structured update and sketched update (Kone ˇ cn ` y et al.
The structured update indicates the methods that enforce a pre-deﬁned ﬁxed 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, sparsiﬁcation, or quantization.
Both directions are well studied and have shown successful results (Alistarh et al.
; Albasyoni et al.
Wangni et al.
; Wang et al.
; Vogels et al.
; Wang et al.
Alistarh et al.
(2017); Wen et al.
; Reisizadeh et al.
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.
; Morcos et al.
Centered Kernel Alignment (CKA) is an improved extension of CCA (Kornblith et al.
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.
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.
; Kumar et al.
(2019); Zhang & He (2020)
; Goutam et al.
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.
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.
B ACKGROUND
Federated Optimization – We consider federated optimization problems as follows.
[Equation]
where p i = n i /n is the ratio of local data to the total dataset, and F i ( x ) = 1 n i (cid:80) ξ ∈D f i ( x , ξ ) is the local objective function of client i .
n is the global dataset size and n i 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 deﬁne such a discrepancy among the models as follows.
[Equation]
where m is the number of clients, u is the synchronized model, and x i 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.
L AYER - WISE A DAPTIVE M ODEL A GGREGATION
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 deﬁne ‘layer-wise unit model discrepancy’, a useful metric for prioritizing the layers as follows.
[Equation]
where L is the number of layers, l is the layer index, u is the global parameters, x i is the client
i ’s local parameters, τ is the aggregation interval, and dim ( · ) is the number of parameters.
This metric quantiﬁes how much each parameter contributes to the model discrepancy at each iteration.
The communication cost is proportional to the number of parameters.
Thus, (cid:80) mi =1 p i (cid:107) u l − x il (cid:107) 2 / dim ( u l ) 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 d l 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: τ (cid:48) 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 d is using the synchronized
At the of
FedLAMA adjusts the model aggregation interval at all the L layers
Algorithm 2 ﬁnds 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.
[Equation]
where ˆ d i is the i th smallest element in the sorted list of the proposed metric d .
Given l layers with the smallest d l values, δ l quantiﬁes their contribution to the total model discrepancy.
Second, the communication cost impact is estimated as follows.
[Equation]
λ l is the ratio of the parameters at the l layers with the smallest d l 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 ﬁnding the l value that makes δ l and 1 − λ l similar.
In this way, it ﬁnds the aggregation interval setting that slightly sacriﬁces 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 .
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 d l 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 ﬁnding 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 d l 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 a) CIFAR-10 (ResNet20) b) CIFAR-100 (WRN28-10)     25 0 .0 .0 .0 .0 .1 .no
r m a li
z
ed m agn i t ude # o f l a y e r s w i
t h
i n c r ea s ed i n t e r v a l de l t a l a m bda    20 0 .0 .0 .0 .0 .1 .no
r m a li
z
ed m agn i t ude # o f l a y e r s w i
t h
i n c r ea s ed i n t e r v a l de l t a l a m bda scalability.
When increasing the aggregation interval, Algorithm 2 multiplies a pre-deﬁned small constant φ to the ﬁxed base interval τ (cid:48) (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 τ (cid:48) .
So, it is guaranteed that the whole model parameters are fully synchronized after φτ (cid:48) iterations.
Because of the layers with the base aggregation interval τ (cid:48) , the total model discrepancy of FedLAMA after φτ (cid:48) iterations is always smaller than that of FedAvg with an aggregation interval of φτ (cid:48) .
C ONVERGENCE A NALYSIS 5. P RELIMINARIES
Notations – All vectors in this paper are column vectors.
x ∈ R d 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 (cid:107)·(cid:107) and (cid:107)·(cid:107) op to denote l 2 norm and matrix operator norm, respectively.
Assumptions – We analyze the convergence rate of FedLAMA under the following assumptions.
(Smoothness).
Each local objective function is L-smooth, that is, (cid:107)∇ F i ( x ) − ∇ F i ( y ) (cid:107) ≤ L (cid:107) x − y (cid:107) , ∀ i ∈ { 1 , · · · , m } .
(Unbiased Gradient).
The stochastic gradient at each client is an unbiased estimator of the local full-batch gradient: E ξ [ g i ( x , ξ )]
= ∇ F i ( x ) .
(Bounded Variance).
The stochastic gradient at each client has bounded variance:
E ξ [ (cid:107) g i ( x , ξ ) − ∇ F i ( x ) (cid:107) 2 ≤ σ 2 ] , ∀ i ∈ { 1 , · · · , m } , σ 2 ≥ 0 .
(Bounded Dissimilarity).
For any sets of weights { p i ≥ 0 } mi =1 , (cid:80) mi =1 p i = 1 , there exist constants β 2 ≥ 1 and κ 2 ≥ 0 such that (cid:80) mi
=1 p
i (cid:107)∇ F i ( x ) (cid:107) 2 ≤ β 2 (cid:107) (cid:80) mi =1 p i ∇ F i ( x ) (cid:107) 2 + κ 2 .
If local objective functions are identical to each other, β 2 = 1 and κ 2 = 0 .
5. A NALYSIS
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 satisﬁes η ≤ 12 L , FedLAMA ensures
[Equation]
Lemma 5.2.
(Model Discrepancy)
Under Assumption 1 ∼ 4 , if the learning rate satisﬁes η <2( τ max − 1) L , FedLAMA ensures
[Equation]
where A = 4 η 2 ( τ max − 1) 2 L 2 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 u 1 .
Under Assumption 1 ∼ 4 , if FedLAMA runs for K iterations and the learning rate satisﬁes η ≤ min
[Equation]
, the average-squared gradient norm of u k is bounded as follows
[Equation]
where u ∗ indicates a local minimum and τ max is the largest averaging interval across all the layers.
Remark 1.
(Linear Speedup)
With a sufﬁciently small diminishing learning rate and a large number of training iterations, FedLAMA achieves linear speedup.
If the learning rate is η = √ m √ K , we have
[Equation]
If K > m 3 , the ﬁrst 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 = φτ (cid:48) .
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.
E XPERIMENTS
Experimental Settings – We evaluate FedLAMA using three representative benchmark datasets:
CIFAR-10 (ResNet20 (He et al.
, CIFAR-100 (WideResNet28-10 (Zagoruyko & Komodakis (2016))), and Federated Extended MNIST (CNN (Caldas et al.
We use TensorFlow 2.4.
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 classiﬁcation 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.
[Equation]
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 toand the learning rate is tuned based on a grid search.
For CIFAR-10 and CIFAR-100, we artiﬁcially 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 φτ (cid:48) iterations.
We report the average accuracy across at least three separate runs.
[Table] [Table]
6. C LASSIFICATION P ERFORMANCE A NALYSIS
To evaluate the proposed model aggregation scheme
[Table]
we keep all the other factors the same, such as optimizer, the number of clients, the degree of heterogeneity
[Table]
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 τ (cid:48) • Periodic full aggregation with an interval of φτ (cid:48)
• Layer-wise adaptive aggregation with intervals of τ (cid:48) and φ
The ﬁrst 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 τ (cid:48) 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 ﬁrst present CIFAR-10 and CIFAR-100 classiﬁcation 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 ﬁrst row shows the performance of FedAvg with a short interval τ (cid:48)
= 6 .
As the interval increases, FedAvg signiﬁcantly loses the accuracy while the communication cost is proportionally reduced.
FedLAMA achieves a comparable accuracy to
FedAvg with τ (cid:48)
= 6 while its communication cost is similar to that of FedAvg with φτ (cid:48) .
These results demonstrate that Algorithm 2 effectively ﬁnds 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 classiﬁcation results.
The base interval τ (cid:48) is set to 10 .
FedAvg ( φ = 1 ) signiﬁcantly 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 .
In contrast, FedLAMA maintains the accuracy when φ increases, while the communication cost is remarkably reduced.
This result demonstrates that FedLAMA effectively ﬁnds 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 τ (cid:48) 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 signiﬁcantly loses the accuracy as the model aggregation interval
[Table]
[Table] [Table]
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 τ (cid:48) while having the communication cost that is close to the periodic full averaging with the increased interval of φτ (cid:48) .
Especially, FedLAMA works effectively even when the Dirichlet’s coefﬁcient is set to 0 .
The coefﬁcient 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.
6. C OMMUNICATION E FFICIENCY A NALYSIS
We analyze the total number of communications and the accumulated data size to evaluate the communication efﬁciency of FedLAMA.
Figure 2 shows the total number of communications at the
[Figure] [Figure]
individual layers.
The τ (cid:48) 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 d l 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 signiﬁcantly smaller total data size than FedAvg.
The extra computational cost of FedLAMA is almost negligible since it calculates d l 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 coefﬁcients and device participation ratios.
That is, FedLAMA can be considered as an effective model aggregation scheme regardless of the degree of data heterogeneity.
C ONCLUSION
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 ﬂexible 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 efﬁcacy 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.
C ODE OF E THICS
Our work does not deliver potentially harmful insights or conﬂicts of interests.
We also do not ﬁnd 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.
R EPRODUCIBILITY S TATEMENT
The software versions, implementation details, hyper-parameter settings can be found in the ﬁrst 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.
A A PPENDIX A. C ONVERGENCE A NALYSIS
Herein, we provide the proofs of the lemmas and theorem shown in Section 5.
A.1. P RELIMINARIES
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 ∈ R d 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 (cid:107) · (cid:107) and (cid:107) · (cid:107) op to denote l 2 norm and matrix operator norm, respectively.
Objective Function – In this paper, we consider federated optimization problems as follows.
[Equation]
where p i = n i /n is the ratio of local data to the total dataset, and F i ( x ) = 1 n i (cid:80) ξ ∈D f i ( x , ξ ) is the local objective function of client i .
n is the global dataset size and n i is the local dataset size.
Note that, by deﬁnition, (cid:80) mi =1 p i = 1 .
Averaging Matrix – We deﬁne a time-varying averaging matrix W k ∈ R dm × dm as follows.
[Equation]
I is an identity matrix, P is also a time-varying averaging matrix, and J is a full averaging matrix.
First, P 1 i is a d × d diagonal matrix that has 1 for the diagonal elements that correspond to the LCL parameters and p i for all the other diagonal elements.
Likewise, P
i is another d × d diagonal matrix that has 0 for the diagonal elements that correspond to the LCL parameters and p i for all the other diagonal elements.
Then, P is deﬁned as follows.
[Equation]
The i th block column of P consists of P 1 i and P
i following the above deﬁnition.
Here we present an example of P where m = 2 and d = 2 .
In this example, p 0 = 1 / 3 and p 1 = 2 / 3 .
Saying the LCL is the second parameter, P is deﬁned as follows.
[Equation]
The full-averaging matrix J is deﬁned as follows.
First, J i is a d × d diagonal matrix that has p i for the diagonal elements.
Then, J consists of m × m blocks of J i such that each column block is m of
J i blocks.
Here we present an example of J where m = 2 and d = 2 as follows.
[Equation]
J = (cid:20) J
J 1 J
J(cid:21)
=                    .
The averaging matrix P and J have the following properties:
P1 dm = 1 dm , J1 dm = 1 dm .
The product of any two averaging matrices consists only of diagonal block matrices because all the blocks in P and J are diagonal.
PJ =
JP =
J regardless of which layers are chosen as the LCL.
PP = P regardless of which layers are chosen as the LCL.
Vectorization – We deﬁne a vectorized form of m local model parameters x k ∈ R dm , its stochastic gradients g k ∈ R dm , and the full gradients f k ∈ R dm as follows
[Equation]
The full model aggregation can be written using the vectorized form of local models x k and the averaging matrix J as follows.
[Equation]
where x ( i,j ) k is the j th model parameter of local model i at iteration k .
We also deﬁne the following additional vectorized forms of the weighted model parameters and gradients for convenience.
[Equation]
Assumptions – We analyze the convergence rate of FedLAMA under the following assumptions.
(Smoothness).
Each local objective function is L-smooth, that is, (cid:107)∇ F i ( x ) − ∇ F i ( y ) (cid:107) ≤ L (cid:107) x − y (cid:107) , ∀ i ∈ { 1 , · · · , m } .
(Unbiased Gradient).
The stochastic gradient at each client is an unbiased estimator of the local full-batch gradient: E ξ [ g i ( x , ξ )]
= ∇ F i ( x ) .
(Bounded Variance).
The stochastic gradient at
each
client has bounded variance:
E ξ [ (cid:107) g i ( x , ξ ) − ∇ F i ( x ) (cid:107) 2 ≤ σ 2 ] , ∀ i ∈ { 1 , · · · , m } , σ 2 ≥ 0 .
(Bounded Dissimilarity).
For any sets of weights { p i ≥ 0 } mi =1 , (cid:80) mi =1 p
i = 1 , there exist constants β 2 ≥ 1 and κ 2 ≥ 0 such that (cid:80) mi
[Equation]
p i (cid:107)∇
[Equation]
If local objective functions are identical to each other, β 2 = 1 and κ 2 = 0 .
A.1. P ROOFS
Theorem 5.1.
Suppose all m local models are initialized to the same point u 1 .
Under Assumption 1 ∼ 4 , if FedLAMA runs for K iterations and the learning rate satisﬁes η ≤ min
[Equation]
, the average-squared norm of ∇ F ( u k ) is bounded as follows.
[Equation]
where u ∗ indicates a local minimum.
Proof.
Based on Lemma 5.1 and 5.2, we have
[Equation]
After re-writing the left-hand side and a minor rearrangement, we have
[Equation]
By moving the third term on the right-hand side to the left-hand side
[Equation]
we have
[Equation]
If A ≤ 1 2 β 2 +
[Equation]
then Aβ 2 1 − A ≤ 12 .
Therefore
[Equation]
(21) can be simpliﬁed as follows.
[Equation]
The learning rate condition A ≤ 1 2 β 2 +1 also ensures that 1 1 − A ≤ 1 + 12 β 2 .
Based on Assumption 4, 12 β 2 ≤ 23 , and thus 1 1 − A ≤ 23 .
Therefore, we have
[Equation]
We complete the proof.
Lemma 5.1.
(Framework)
Under Assumption 1 ∼ 3 , if the learning rate satisﬁes η ≤ 12 L , FedLAMA ensures
[Equation]
Proof.
Based on Assumption
[Equation]
we have
[Equation]
First, T 1 can be rewritten as follows.
[Equation]
where the last equality holds based on a basic equality:
[Equation]
Then, T 2 can be bounded as follows.
[Equation]
2   =
E   (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) m (cid:88)
i = p
i (cid:0) g i ( x ik ) − ∇ F i ( x ik ) (cid:1) + m (cid:88) i
= p
i ∇ F i ( x ik ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 2   ≤ 2 E   (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) m (cid:88) i
= p
i (cid:0) g i ( x ik ) − ∇ F i ( x ik ) (cid:1)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 2   + 2 E   (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) m (cid:88) i
= p
i ∇ F i ( x ik ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 2   =m (cid:88) i
= p 2 i E (cid:104)(cid:13)(cid:13) g i ( x ik ) − ∇ F i ( x ik ) (cid:13)(cid:13) 2 (cid:105) + 2 E   (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) m (cid:88) i
= p
i ∇ F i ( x ik ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
  ≤ 2 σm (cid:88)
i = p 2 i + 2 E   (cid:13) (cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
where the last equality holds because g i ( x ik ) − ∇ F i ( x ik ) has 0 mean and is independent across i , and the last inequality follows Assumption 3.
By plugging in (24) and (25) into (23), we have the following.
[Equation]
If η ≤ 12 L
[Equation]
it follows
[Equation]
where (26) holds based on the convexity of (cid:96) 2 norm and Jensen’s inequality.
By taking expectation and averaging across K iterations
[Equation]
we have.
[Equation]
F
[Equation]
] η
[Equation]
After a minor rearrangement, we have a telescoping sum as follows.
[Equation]
where u ∗ indicates the local minimum.
Here, we complete the proof.
Lemma 5.2.
(Model Discrepancy)
Under Assumption 1 ∼ 4 , if the learning rate satisﬁes η <2( τ max − 1) L , FedLAMA ensures
[Equation]
where A = 4 η 2 ( τ max − 1) 2 L 2 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.
[Equation]
where (28) holds by the commutative property of multiplication.
Then, according to the parameter update rule
[Equation]
we have
[Equation]
where (29) holds because
JW = J based on the averaging matrix property [Equation] and IW = W .
Then, expanding the expression of x k − [Equation] we have
[Equation]
Repeating the same procedure for ˆx k − 2 , ˆx k − 3 , · · · , ˆx 2 , we have
[Equation]
where (30) holds because x i 1 is the same across all the workers and thus ( J − I ) ˆx 1 = 0 .
Based on (30), we have
[Equation]
[Table]
[Equation]
[Table]
 
[Equation]
[Table]
[Equation]
[Table]
(cid:13)(cid:13)
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table]
[Equation]
[Table] [Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]

[Equation]
31) where (31) holds based on the convexity of (cid:96) 2 norm and Jensen’s inequality.
Now, we focus on bounding T 3 and T 4 , separately.
Bounding T
[Equation]
(32) ≤
[Equation]
k = k −
[Equation]
where (32) holds because ˆ g s − ˆ f s has 0 mean and independent across s , and (33) 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
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ( J − aτ
[Equation]
) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) op   =
[Equation]
aτ
[Equation]
E
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) (
[Equation]
aτ
[Equation]
(cid:89)
[Equation]
s
[Equation]
l ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) op 
 (34) =
[Equation]
(cid:20)(cid:13)(cid:13)(cid:13) (
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
(cid:21) (35) =
[Equation]
≤
[Equation]
a
[Equation] [Equation]
b
[Equation]
s = aτ
[Equation]
(cid:88)
[Equation]
Remember FedLAMA synchronizes the whole parameters at least once after every τ max iterations.
Thus, (34) holds because (cid:81)
aτ max + b − 1 l = s
W l is J when s ≤ aτ max , and thus J − (cid:81)
aτ max + b − 1 l = s
W l becomes 0 .
(35) holds based on Lemma A.2.
(36) holds based on Assumption 3.
Bounding T
[Equation]
max (cid:88) k
=
E   (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) k − 1 (cid:88) s =
( J − k − 1 (cid:89) l = s
W l ) ˆf s (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 2   =
K/τ max − 1 (cid:88) a = τ max (cid:88) b =
E   (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) aτ
+ b − 1 (cid:88) s =
( J − aτ max + b − 1 (cid:89) l
= s
W l ) ˆf s (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 2   =
K/τ max − 1 (cid:88) a = τ max (cid:88) b =
E   (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) aτ max + b − 1 (cid:88) s
= aτ max
+ ( J − aτ max + b − 1 (cid:89) l
= s
P l ) ˆf s (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 2   (38) ≤
K/τ max − 1 (cid:88) a = τ max (cid:88) b =   ( b − 1)
aτ max + b − 1 (cid:88) s
= aτ max +
E   (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ( J − aτ max + b − 1 (cid:89) l
= s
P l ) ˆf s (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 2   
 (39) ≤
K/τ max − 1 (cid:88) a = τ max (cid:88) b =   ( b − 1)
aτ max + b − 1 (cid:88) s
= aτ max +
E   (cid:13)(cid:13) (cid:13) ˆ f s (cid:13)(cid:13) (cid:13) 2 (cid:13)(cid:13)(cid:13) (cid:13)(cid:13)
( J − aτ max + b − 1 (cid:89) l
= s
P l ) (cid:13)(cid:13)(cid:13) (cid:13)(cid:13) op 
 
 (40) ≤
K/τ max − 1 (cid:88) a = τ max (cid:88) b = (cid:32) ( b − 1)
aτ max + b − 1 (cid:88) s
= aτ max +
E (cid:20)(cid:13)(cid:13)(cid:13) ˆf s (cid:13)(cid:13)(cid:13) 2 (cid:21)(cid:33) (41) ≤ τ max ( τ max − 1)
K/τ max − 1 (cid:88) a = (cid:32) aτ max + τ max − 1 (cid:88) s
= aτ max +
E (cid:20)(cid:13)(cid:13)(cid:13) ˆf s (cid:13)(cid:13)(cid:13) 2 (cid:21)(cid:33) ≤ τ max ( τ max − 1)
K (cid:88) k
=
E (cid:20)(cid:13)(cid:13)(cid:13) ˆf k (cid:13)(cid:13)(cid:13)
[Equation]
where (38) holds because J − (cid:81)
aτ max + b − 1 l = s
P l becomes 0 when s ≤ aτ .
(39) holds based on the convexity of (cid:96) 2 norm and Jensen’s inequality.
(40) holds based on Lemma A.1.
(41) holds based on
Lemma A.2.
Final Result
By plugging in (37) and (42) into (31)
[Equation]
we have
[Equation]
The local gradient term on the right-hand side in (43) can be rewritten using the following inequality.
[Equation]
where (44) holds based on the convexity of (cid:96) 2 norm and Jensen’s inequality.
Plugging in (45) into (43), we have
[Equation]
After a minor rearranging
[Equation]
we have
[Equation]
Let us deﬁne A = 2 η 2 τ max ( τ max − 1) L 2 .
Then (47) is simpliﬁed as follows.
[Equation]
Based on Assumption [Equation] we have
[Equation]
where (49) holds based on the deﬁnition of the objective function (10).
Note that (48) is true only when 1 − A > 0 .
Thus, after a minor rearrangement, we have a learning rate constraint as follows.
[Equation]
Here, we complete the proof.
A.1. P ROOF OF O THER L EMMAS
Lemma A.1.
Consider a real matrix A ∈ R md j × md j and a real vector b ∈ R md j .
If b (cid:54) =
md j , we have
[Equation]
[Equation]
Proof.
[Equation]
where (51) holds based on the deﬁnition of operator norm.
Lemma A.2.
Suppose an md × md averaging matrix P and the full-averaging matrix J , then [Equation] P
[Equation]
regardless of which layers are chosen as the LCL.
Proof.
First, by the deﬁnition 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 deﬁnition of J , all the non-zero elements in i th column are the same p
i , 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 deﬁnition of the matrix operator norm, (cid:107) J − P (cid:107) 2 op = max {| λ ( J − P ) |} = 1 , where λ ( · ) indicates the eigenvalues of the input matrix.
[Figure]
T r a i n i
[Figure]
ft
[Figure] [Figure]
A. A DDITIONAL E XPERIMENTAL R ESULTS
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 ﬁrst 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
[Figure]
( τ ’) while cost that is similar to aggregation
Artiﬁcial Data Heterogeneity – For CIFAR-10 and CIFAR-100, we artiﬁcially generate the heterogeneous data distribution using Dirichlet’s distribution.
The concentration coefﬁcient α is set to 0 .
5 , and 1 .
to evaluate the performance of FedLAMA across a variety of degree of data heterogeneity.
Note that the small concentration coefﬁcient 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.
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 classiﬁcation 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 τ (cid:48) settings.
We see that the accuracy is signiﬁcantly dropped as τ (cid:48) increases.
CIFAR-
– 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 andshow the CIFAR-100 classiﬁcation 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 coefﬁcient of 0 . 1 ).
Table 11 shows the FedAvg accuracy with different τ (cid:48) 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 classiﬁcation performance of FedLAMA across different φ settings.
FedLAMA achieves a similar accuracy to the baseline (FedAvg with τ (cid:48)
= 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.
[Table]
[Table] [Table]
[Table] [Table]
results of FedAvg with different τ (cid:48) settings.
[Table] [Table]