AQUILA: Communication Efficient Federated Learning with Adaptive Quantization of Lazily-Aggregated Gradients

1 INTRODUCTION . With the deployment of ubiquitous sensing and computing devices , the Internet of things ( IoT ) as well as many other distributed systems have gradually grown from concept to reality , bringing dramatic convenience to people ’ s daily life ( Du et al. , 2020 ; Liu et al. , 2020a ; Hard et al. , 2018 ) . To fully utilize such distributed computing resources , distributed learning provides a promising framework that can achieve comparable performance with the traditional centralized learning scheme . However , the privacy and security of sensitive data during the updating and transmitting processes in distributed learning have been a growing concern . In this context , federated learning ( FL ) ( McMahan et al. , 2017 ; Yang et al. , 2019 ) has been developed , enabling distributed devices to collaboratively learn a global model without privacy leakage by keeping private data sets isolated and masking transmitted information with secure approaches like differential privacy ( Abadi et al. , 2016 ) , secret sharing techniques ( Bonawitz et al. , 2017 ) and homomorphic encryption ( Liu et al. , 2020b ) . Due to its privacy and security preserving property and great potentials in some distributed but privacysensitive fields like finance and health , FL has attracted tremendous attentions from both academia and industry in recent years . Unfortunately , in many FL applications like image classification and objective recognition , the models to be trained tend to be high-dimensional , which lead to heavy communication overheads , for example , a Resnet-152 network has over 58 million parameters ( He et al. , 2016 ) . Hence , communication efficiency has become one of the key bottlenecks of FL . To this end , recent researches have tried to reduce the communication frequency , for example , Sun et al . ( 2020 ) proposes the lazily-aggregated quantization ( LAQ ) method to reduce communication rounds by skipping some unnecessary parameter uploads . To further reduce transmitted bits per communication , LAQ can be used jointly with gradient compression techniques , e.g . quantization and sparsification ( Strom , 2015 ; Wangni et al. , 2018 ; Lin et al. , 2018 ; Han et al. , 2020 ) . Moreover , Mao et al . ( 2021 ) develops the Adaptive Quantized Gradient ( AQG ) for LAQ to adjust the quantization bit among multiple given levels during training . However , the AQG is not sufficiently adaptive , for example , in the two-level AQG with 4 bit and 2 bit , the situation of 3 bit and 1 bit is not covered at all . In a separate line of work , Jhunjhunwala et al . ( 2021 ) develops an adaptive quantization rule ( AdaQuantFL ) for FL which can search in a given range for an optimal quantization level and achieve a better error-communication tradeoff . Previous work has investigated how to optimize communication frequency or adjust quantization levels in a highly adaptive fashion , but not both . Intuitively , we ask the question , can we adaptively adjust the quantization level in LAQ to further reduce communication rounds and transmitted bits simultaneously ? A straightforward approach is to train LAQ jointly with state-of-the-art adaptive quantization methods like AdaQuantFL . However , quantization level and communication frequency are mutually dependent in establishing the model convergence and must be adjusted together . For example , lazy aggregation leads to the skip of some gradients that no longer require quantization , while the choice of quantization levels directly affects the quality of a gradient update and thus whether it is selected for transmission and aggregation . Therefore , the key question of our work is how to jointly leverage these two complementary yet mutually-dependent degrees of freedom for further optimizing communication efficiency in FL . The key idea of this paper is to select the optimal quantization bits for each communication round in LAQ by optimizing the gradient loss caused by skipping quantized updates , which gives a novel quantization criterion that can cooperate with LAQ strategy to further reduce overall transmitted bits while maintaining the desired convergence properties of LAQ . The contributions of this paper are trifold and summarized as follows . 1 ) We propose a FL framework with adaptive quantization of lazily-aggregated gradients termed AQUILA , which simultaneously adjusts the communication frequency and quantization level in a synergistic fashion . 2 ) We formulate AQUILA as an optimization problem and develop an upperbound for the gradient loss caused by communication skipping , which gives a novel adaptive quantization criterion which is theoretically proven to be more efficient compared to AdaQuantFL while maintaining the same convergence properties . 3 ) We experimentally evaluate the performance of AQUILA in a number of non-homogeneous FL settings , including Non-IID data distribution and various heterogeneous model architecture . The experimental results show that AQUILA can significantly mitigate the communication overhead over a number of baselines including fixed-bit LAQ and the naive combination of LAQ and AdaQuantFL . Our approach is highly adaptive and compatible to existing FL settings . 2 BACKGROUND AND RELATED WORK . Consider a FL system with one central server and a clients set M of M distributed clients to collaboratively train a global model parameterized by θ∗ . At iteration k , each client m ∈ M first trains the global model θk on its local data Dkm , and sends the local gradient g k m = ∇fm ( Dkm ; θ k ) to the central server . Then the server aggregates the parameters and updates the global parameter by : θk+1 = θk − α M ∑ m∈M gkm . ( 1 ) To reduce communication overheads with gradient quantization , the stochastic uniform quantizer ( Alistarh et al. , 2017 ) is usually adopted . For any local gradient g ∈ Rd , the quantized value of its i-th dimension with quantization level b is defined as : Qb ( gi ) = ∥g∥2 · sign ( gi ) · ξi ( g , b ) , ( 2 ) where ξi ( g , b ) is a random variable defined as follows . Let l ∈ { 0 , 1 , 2 , ... , b−1 } be an integer satisfying |gi| /∥g∥2 ∈ [ l/b , ( l + 1 ) /b ) , then : ξi ( g , b ) = { ( ℓ+ 1 ) /b with probability ( b · ∥gi∥ ) /∥g∥2 − l ℓ/b otherwise . ( 3 ) It is clear that with quantization level b , the number of bits for transmitting a quantized gradient from a client to the central server is Cb = d ⌈log2 ( b+ 1 ) ⌉+ d+ 32 , with 32 bits for ∥g∥2 , 1 bit for each sign ( gi ) , and log2 ( b+ 1 ) bits for each ξi ( g , b ) . For communication rounds reduction , the LAQ proposes to let the client m ∈ M upload its newlyquantized local gradient Qb ( gkm ) at iteration k only when the change in local gradient is sufficiently large , i.e. , ∥∥∥Qb ( gkm ) −Qb ( ĝk−1m ) ∥∥∥2 2 ≥ ∑D d=1 ξd ∥∥∥θk+1−d − θk−d∥∥∥2 2 α2M2 + 3 ( ∥εb ( ĝk−1m ) ∥22 + ∥∥εb ( gkm ) ∥∥22 ) , ( 4 ) where Qb ( ĝ k−1 m ) is the last quantized upload from client m , εb ( ĝ k−1 m ) and εb ( g k m ) denote quantization errors , and { ξd } Dd=1 are some predetermined constant weights . Notice here a fixed quantization level b is used . In LAQ , if the difference between client m ’ s newly-quantized local gradient Qb ( gkm ) and the last upload is smaller than a threshold involving quantization errors and global model ’ s innovation , client m will skip the upload of Qb ( gkm ) at iteration k and the central server will reuses Qb ( ĝ k−1 m ) for such lazy aggregation : θ̂ k+1 = θk − α M ∑ m∈M Qb ( ĝ k m ) = θ k − α M ∑ m∈M\Mk0 Qb ( g k m ) − α M ∑ m∈Mk0 Qb ( ĝ k−1 m ) , ( 5 ) where Mk0 denotes the subset of clients that skip the new gradient update and reuse the old quantized gradient at iteration k. Besides , ĝkm represents the actual gradient for aggregation from client m , which is gkm for m ∈ M \Mk0 , while ĝ k−1 m for m ∈ Mk0 . Recently , AdaQuantFL is proposed to achieve a better error-communication tradeoff by adaptively adjusting the quantization levels during FL training ( Jhunjhunwala et al. , 2021 ) . Specifically , AdaQuantFL computes iteration k ’ s optimal quantization level b∗k based on the following criterion involving training loss and initial quantization level b0 : b∗k = √ f ( θ0 ) /f ( θk ) · b0 , ( 6 ) where f ( θ0 ) and f ( θk ) are the global training loss at iteration 0 and k , respectively . However , AdaQuantFL transmits quantized gradients at every iteration . In order to skip unnecessary communication rounds and adaptively adjust quantization level for each communication jointly , an naive approach is to quantize lazily aggregated gradients with AdaQuantFL . Nevertheless , it fails to achieve efficient communication due to a number of reasons . Firstly , given the descending trend of training loss , AdaQuantFL ’ s criterion ( 6 ) may lead to high quantization bit number even exceeding 32 late in the training process , which is too much for cases where the global convergence is already approaching . Secondly , higher quantization level results in smaller quantization error , which will lead to lower communication threshold in LAQ ’ s criterion ( 4 ) and thus higher frequency of transmission . Therefore it is desirable to develop more efficient adaptive quantization method in the lazily-aggregated setting to systematically improve communication efficiency in FL . 3 METHOD . 3.1 ADAPTIVE QUANTIZATION OF LAZILY-AGGREGATED GRADIENTS ( AQUILA ) . Given the above limitations of the naive joint use of existing adaptive quantization criterion and lazy aggregation strategy , this paper aims to design a unifying framework for communication efficiency optimization where the quantization level and communication frequency are adjusted in a synergistic and interactive way . Based on a careful investigation on LAQ , we design a novel quantization criterion where the optimal quantization level is selected by minimizing the expected gradient loss caused by skipping quantized updates . The rationale behind such strategy is , by formulating the adaptive quantization problem of lazily-aggregated gradients as optimizing the expected gradient loss with respect to the number of quantization bits , we can get an adaptive quantization criterion based on local gradient updates while maintaining or even improving the convergence property of LAQ . To get the optimization target , we first derive an upperbound for the expected gradient loss in terms of quantization bits as elaborated in Section 3.2 , which gives a novel adaptive quantization criterion ( 7 ) that selects the quantization level for client m at iteration k based on initial quantization level b0 and the change between client m ’ s newly computed gradient gkm and the last uploaded gradient ĝk−1m : ( bkm ) ∗ = b0 · √∥∥g1m − ĝ0m∥∥22/∥∥gkm − ĝk−1m ∥∥22 . ( 7 ) The superiority of ( 7 ) comes from the following two aspects . Firstly , the gradient updates tend to fluctuate along with the training process instead of keeping descending like the loss value , and thus prevent the quantization level from increasing tremendously compared with the initial level . Secondly , with lazy aggregation criterion based on gradient updates like ( 4 ) , the transmitted bits in AQUILA is further controlled , since the gradient update for actual transmission in ( 7 ) is lower bounded by the lazy aggregation criterion , and therefore high-bit transmission for small update is more likely to be skipped . To better fit the larger quantization errors induced by fewer quantization bits in ( 7 ) , we modify the communication criterion as follows to avoid the potential expansion of clients group to be skipped : ∥∥∥Qbkm ( gkm ) −Qb̂k−1m ( ĝk−1m ) ∥∥∥22≥ ∑D d=1ξd ∥∥∥θk+1−d−θk−d∥∥∥2 2 α2M2 +2 ∥∥∥εbkm ( gkm ) −εb̂k−1m ( ĝk−1m ) ∥∥∥22 , ( 8 ) where all the notations are the same as in ( 4 ) except the heterogeneous quantization level bkm and b̂k−1m for each client . For detailed development of ( 8 ) , please refer to the Appendix . The cooperation of the novel adaptive quantization criterion ( 7 ) and the modified lazy aggregation strategy ( 8 ) is illustrated in Fig . 1a . Compared to the naive combination of AdaQuantFL and LAQ where the mutual influence between adaptive quantization and lazy aggregation has not been considered as shown in Fig . 1b , our AQUILA framework adaptively optimizes the allocation of quantization bits throughout training to promote the convergence of lazy aggregation , and at the same time utilizes the lazy aggregation strategy to improve the efficiency of adaptive quantization by skipping high-bit transmission . The proposed AQUILA ’ s effect of suppressing the transmission of high quantization bits has been verified in our experiments , as shown in Fig . 20 in the Appendix . Besides , with the adjusted lazy aggregation strategy ( 8 ) , AQUILA well addresses the problem of high communication frequency in the late training process of naive combination of AdaQuantFL and LAQ , as indicated by Fig . 2 . The proposed AQUILA is summarized as follows in Algorithm 1 . At iteration k = 0 , all clients are forced to transmit local gradients quantized with the initial level b0 . At iteration k ∈ { 1 , 2 , ... , K } , the server first broadcasts the global model θk to all clients . Each client m computes gkm with local training data , and then uses it to select an optimal quantization level by ( 7 ) . Then , each client computes its gradient update after quantization and determines whether to upload the update or not based on the communication criterion ( 8 ) . Finally , the server updates the new global model θk+1 with up-to-date quantized gradients Qbkm ( g k m ) for those clients who send the uploads at iteration k , while reusing the old quantized gradients Qb̂k−1m ( ĝ k−1 m ) for those who skip the uploads . Algorithm 1 Communication Efficient FL with AQUILA Input : the number of communication rounds K , the learning rate α , the maximum communication level bmax Initialize : the initial global model parameter θ0 and the initial quantization level b0 . 1 : Server broadcasts θ0 to all clients . 2 : for each client m ∈ M in parallel do 3 : Calculates local gradient g0m and sends the quantized gradient Qb0 ( g 0 m ) . 4 : Set ĝ0m = g 0 m and b̂ 0 m = b0 on both sides . 5 : end for 6 : for k = 1 , 2 , ... , K do 7 : Server broadcasts θk to all clients . 8 : for each client m ∈ M in parallel do 9 : Calculates local gradient gkm . 10 : Computes the optimal local quantization level bkm by ( 7 ) . 11 : if bkm ≥ bmax then 12 : bkm = bmax . 13 : end if 14 : if ( 8 ) holds for client m then 15 : Client m computes and sends the quantized gradient Qbkm ( g k m ) . 16 : Set ĝkm = g k m and b̂ k m = b k m on both sides . 17 : else 18 : Client m sends nothing . 19 : Set ĝkm = ĝ k−1 m and b̂ k m = b̂ k−1 m on both sides . 20 : end if 21 : end for 22 : Server updates θk+1 by θk − α ∑M m=1 Qb̂km ( ĝkm ) . 23 : end for

Summary: The paper proposes combining two orthogonal algorithms -- the lazily aggregated gradient(LAQ)method and adaptive quantization (AdaQuantFL)-- to reduce communication complexity in federated learning. In particular, while LAQ focuses on the frequency of gradient updates in each iteration, AdaQuantFL focuses on the adaptive allocation of bits across iterations. This paper proposes to combine both of these approaches for better communication complexity in federated learning.   At the heart of the combining strategy is Eq. 7, which allocates the bits used in each iteration. This equation is intern based on minimizing the upper bound on the expected deviation in the path of the iterate due to gradient skipping. The authors theoretically show that such a bit allocation strategy leads to reduced communication compared to the bit allocation strategy of AdaQuantFl.