On Distributed Adaptive Optimization with Gradient Compression

1 INTRODUCTION . Deep neural network has achieved the state-of-the-art learning performance on numerous AI applications , e.g. , computer vision ( Goodfellow et al. , 2014 ; He et al. , 2016 ; Voulodimos et al. , 2018 ) , Natural Language Processing ( Graves et al. , 2013 ; Young et al. , 2018 ; Zhang et al. , 2018 ) , Reinforcement Learning ( Mnih et al. , 2013 ; Silver et al. , 2017 ) and recommendation systems ( Covington et al. , 2016 ; Wei et al. , 2017 ) . With the increasing sheet size of data observations and growing complexity of deep neural networks , standard single-machine training procedures encounter at least two major challenges : • Due to the limited computing power of a single-machine , processing the massive number of data samples takes a long time — training is too slow . Many real-world applications can not even afford spending that much time on training . • In many scenarios , data are stored in multiple servers , possibly at different locations , due to the storage constraints ( massive user behavior data , Internet images , etc . ) or privacy reasons ( Chang et al. , 2018 ) . Hence , transmitting data among servers might be costly . Distributed learning framework ( Dean et al. , 2012 ) has been a common training strategy to tackle the above two issues . Consider the distributed optimization task where n workers jointly solve a large finite-sum optimization problem in the form of min θ f ( θ ) : = min θ 1 n n∑ i=1 fi ( θ ) = 1 n n∑ i=1 Ex∼Xi [ Fi ( θ ; x ) ] , ( 1 ) where the non-convex function fi represents the average loss over the local data samples for worker i ∈ [ n ] , and θ ∈ Rd the global model parameter . Xi is the data distribution on each local node . In the classical centralized distributed setting , in each iteration the central server uniformly randomly assigns the data to n local workers ( Xi ’ s are the same ) , at which the gradients of the model are computed in parallel . Then the central server aggregates the local gradients , updates the global model ( e.g. , by stochastic gradient descent ( SGD ) ) , and transmits back the updated model to the local nodes for subsequent gradient computation . The scenario where Xi ’ s are different gives rise to the recently proposed Federated Learning ( FL ) ( McMahan et al. , 2017 ) framework , which would not be the major focus of this work . As we can see , distributed training naturally solves aforementioned issues : 1 ) We use n computing nodes to train the model , so the time per training epoch can be largely reduced ; 2 ) There is no need to transmit the local data to central server . Besides , distributed training also provides stronger error tolerance since the training process could continue even one local machine breaks down . As a result of these advantages , there has been a surge of study and applications on distributed systems ( Boyd et al. , 2011 ; Nedic & Ozdaglar , 2009 ; Duchi et al. , 2011 ; Goyal et al. , 2017 ; Hong et al. , 2017 ; Lu et al. , 2019 ; Koloskova et al. , 2019 ) . Gradient compression . Among many optimization strategies , SGD is still the most popular prototype in distributed training for its simplicity and effectiveness ( Chilimbi et al. , 2014 ; Agarwal et al. , 2018 ; Mikami et al. , 2018 ) . Yet , when the deep learning model is very large , the communication between local nodes and central server could be expensive , and the burdensome gradient transmission would slow down the whole training system . Thus , reducing the communication cost in distributed SGD has become an active topic , and an important ingredient of large-scale distributed systems ( e.g. , Seide et al . ( 2014 ) ) . Solutions based on quantization , sparsification and other compression techniques of the local gradients are proposed , e.g. , Alistarh et al . ( 2017 ) ; Wen et al . ( 2017 ) ; Wangni et al . ( 2018 ) ; Stich et al . ( 2018 ) ; Aji & Heafield ( 2017 ) ; Bernstein et al . ( 2018 ) ; De Sa et al . ( 2017 ) ; Yang et al . ( 2019 ) ; Ivkin et al . ( 2019 ) . However , it has been observed both theoretically and empirically ( Stich et al. , 2018 ; Ajalloeian & Stich , 2020 ) , that directly updating with the compressed gradients usually brings non-negligible performance downgrade in terms of convergence speed and accuracy . To tackle this problem , studies ( e.g. , Stich et al . ( 2018 ) ; Karimireddy et al . ( 2019 ) ) show that the technique of error feedback can to a large extent remedy the issue of such gradient compression , achieving same convergence rate as full-gradient SGD . Adaptive optimization . In recent years , adaptive optimization algorithms ( e.g . AdaGrad ( Duchi et al. , 2010 ) , Adam ( Kingma & Ba , 2014 ) and AMSGrad ( Reddi et al. , 2018 ) ) have become popular because of their superior empirical performance . These methods use different implicit learning rates for different coordinates that keep changing adaptively throughout the training process , based on the learning trajectory . In many cases , adaptive methods have been shown to converge faster than SGD , sometimes with better generalization as well . However , the body of literature that extends adaptive methods to distributed training is still fairly limited . In particular , even the simple gradient averaging approach , though appearing standard , has not been analyzed for adaptive optimization algorithms . Given that distributed SGD with compressed gradient averaging can match the performance of standard SGD , one natural question is : is it also true for adaptive methods ? In this work , we try to fill this gap formally , by analyzing COMP-AMS , a distributed adaptive optimization framework using the gradient averaging protocol , with communication-efficient gradient compression . Our contributions . In this work , we study a simple algorithm design leveraging the adaptivity of AMSGrad and the computational virtue of local gradient compression . Our main contributions summarize as follows : • We consider COMP-AMS , a synchronous distributed adaptive optimization framework based on global averaging with gradient compression , which is efficient in both communication and memory as no local moment estimation is needed . We consider the BlockSign and Top-k compressors , coupled with the error-feedback technique to compensate for the bias implied by the compression step for fast convergence . • We provide the convergence analysis of distributed COMP-AMS ( with n workers ) in smooth non-convex optimization . In the special case of n = 1 ( single machine ) , similar to SGD , gradient compression with error feedback in adaptive method achieves same convergence rateO ( 1√ T ) as the standard full-gradient counterpart . Also , we show that with a properly chosen learning rate , COMP-AMS achieves O ( 1√ nT ) convergence , implying a linear speedup in terms of the number of local workers to attain O ( δ ) -stationary point . • Experiments are conducted on various training tasks on image classification and sentiment analysis to validate our theoretical findings on the linear speedup effect . Our results show that COMP-AMS has comparable performance with other distributed adaptive methods , and matches the accuracy of full-precision AMSGrad with substantially reduced communication cost . Thus , it can serve as a convenient distributed training strategy in practice . 2 RELATED WORK . 2.1 DISTRIBUTED SGD WITH COMPRESSED GRADIENTS . Quantization . As we mentioned before , SGD is the most commonly adopted optimization method in distributed training of deep neural nets . To reduce the expensive communication in large-scale distributed systems , extensive works have considered various compression techniques applied to the gradient transaction procedure . The first strategy is quantization . Dettmers ( 2016 ) condenses 32-bit floating numbers into 8-bits when representing the gradients . Seide et al . ( 2014 ) ; Bernstein et al . ( 2018 ) ; Karimireddy et al . ( 2019 ) ; Bernstein et al . ( 2019 ) use the extreme 1-bit information ( sign ) of the gradients , combined with tricks like momentum , majority vote and memory . Other quantization-based methods include QSGD ( Alistarh et al. , 2017 ; Wu et al. , 2018 ; Zhang et al. , 2017 ) and LPC-SVRG ( Yu et al. , 2019b ) , leveraging unbiased stochastic quantization . The saving in communication of quantization methods is moderate : for example , 8-bit quantization reduces the cost to 25 % ( compared with 32-bit full-precision ) . Even in the extreme 1-bit case , the largest compression ratio is around 1/32 ≈ 3.1 % . Sparsification . Gradient sparsification is another popular solution which may provide higher compression rate . Instead of commuting the full gradient , each local worker only passes a few coordinates to the central server and zeros out the others . Thus , we can more freely choose higher compression ratio ( e.g. , 1 % , 0.1 % ) , still achieving impressive performance in many applications ( Lin et al. , 2018 ) . Stochastic sparsification methods , including uniform and magnitude based sampling ( Wangni et al. , 2018 ) , select coordinates based on some sampling probability , yielding unbiased gradient compressors with proper scaling . Deterministic methods are simpler , e.g. , Random-k , Top-k ( Stich et al. , 2018 ; Shi et al. , 2019 ) ( selecting k elements with largest magnitude ) , Deep Gradient Compression ( Lin et al. , 2018 ) , but usually lead to biased gradient estimation . In Ivkin et al . ( 2019 ) , the central server identifies heavy-hitters from the count-sketch ( Charikar et al. , 2002 ) of the local gradients , which can be regarded as a noisy variant of Top-k strategy . More applications and analysis of compressed distributed SGD can be found in Jiang & Agrawal ( 2018 ) ; Shen et al . ( 2018 ) ; Alistarh et al . ( 2018 ) ; Basu et al . ( 2019 ) ; Jiang et al . ( 2018 ) , among others . Error Feedback ( EF ) . Biased gradient estimation , which is a consequence of many aforementioned methods ( e.g. , signSGD , Top-k ) , undermines the model training , both theoretically and empirically , with slower convergence and worse generalization ( Ajalloeian & Stich , 2020 ; Beznosikov et al. , 2020 ) . The technique of error feedback is able to “ correct for the bias ” and fix the convergence issues . In this procedure , the difference between the true stochastic gradient and the compressed one is accumulated locally , which is then added back to the local gradients in later iterations . Stich et al . ( 2018 ) ; Karimireddy et al . ( 2019 ) prove theO ( 1T ) andO ( 1√ T ) convergence rate of EF-SGD in strongly convex and non-convex setting respectively , matching the rates of vanilla SGD ( Nemirovski et al. , 2009 ; Ghadimi & Lan , 2013 ) . More works on the convergence rate of SGD with error feedback include ( Zheng et al. , 2019 ; Stich & Karimireddy , 2019 ) , among other related papers .

Summary: This paper suggests a gradient averaging strategy for distributed adaptive optimization. Gradient compression is used for reducing the communication costs in transmission of gradients, and the tool of error feedback is used for correcting bias injected by the compression step. Convergence analysis results demonstrates that the proposed strategy has a linear speedup as the number of workers increases while the same convergence rate as standard AMSGrad. The experimental results on real-world datasets successfully validate the theoretical results.