Layer-wise Adaptive Model Aggregation for Scalable Federated Learning

1 INTRODUCTION . 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 . 2 RELATED WORKS . 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 . 3 BACKGROUND . 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 . 4 LAYER-WISE ADAPTIVE MODEL AGGREGATION . 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 φτ ′ .

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