Contractive error feedback for gradient compression

1 INTRODUCTION . Many modern machine learning tasks can be cast as an optimization problem min x∈Rd f ( x ) : = 1 N N∑ i=1 Eξ∼D [ fi ( x , ξ ) ] . ( 1 ) Typically , the possibly nonconvex loss function fi ( x , ξ ) is continuously differentiable with respect to ( w.r.t . ) x . We assume that f∗ > −∞ is a global minimum of ( 1 ) . In this work , we will focus on training neural networks ( NN ) by solving ( 1 ) in two important scenarios : i ) data parallelism in a multi-GPU setting , and ii ) federated learning where clients could be edge devices such as smart phones . In both cases , N denotes the number of workers ( GPUs or devices ) and f∗ is no smaller than 0 given common choices of loss functions such as cross entropy and mean square . While training NNs is a resource-hungry task , bandwidth and memory prevent us from fully harvesting the merits of parallelism . Limited bandwidth slows down the speed of information ( e.g. , gradients ) exchange among workers . In multi-GPU settings , it is observed that communication of gradients becomes a bottleneck that significantly drags down the training speed . While for federated learning , the system capacity is usually confined by communication latency between the server and workers . Memory constraints are less studied compared with bandwidth ( Sohoni et al. , 2019 ) . However , given the trend of an increasing model size in both computer vision and natural language processing , e.g. , ViT ( 0.6B ) ( Dosovitskiy et al. , 2020 ) , Megatron-LM ( 8.3B ) ( Shoeybi et al. , 2019 ) , T5 ( 11B ) ( Raffel et al. , 2019 ) , it evidences that memory is becoming an unavoidable constraint for training large models in current GPU generation with 16/32GB memory . For federated learning , it is also necessary to reduce the memory footprint on workers especially for those IoT applications . In this work , we will focus on the parsimonious setting , jointly dealing with limited bandwidth and memory . The priority is still communication efficiency for speeding up training time , otherwise one can simply rely on SGD with almost no additional memory consumption . Many existing works are unable to cope with such a challenging problem because memory is not carefully taken care of . Communication efficiency . A well-documented approach to alleviate communication bottleneck is to compress the ( stochastic ) gradients such that training can be made faster with reduced overhead ( Seide et al. , 2014 ; Alistarh et al. , 2017 ; Karimireddy et al. , 2019 ) . Depending on whether the gradient compressor is biased , these schemes are categorized as follows . Unbiased compressors maintain the unbiasedness of stochastic gradients at the price of enlarged variance . For example , uniform quantization of stochastic gradients is studied in QSGD ( Alistarh et al. , 2017 ) . A concurrent work ( Wen et al. , 2017 ) focuses on a special case of QSGD by quantizing each entry of the gradient into { ±1 , 0 } . Other variants include natural quantization ( Horvath et al. , 2019 ) , non-uniform quantization ( Ramezani-Kebrya et al. , 2021 ) and adaptive quantization ( Faghri et al. , 2020 ) . Another route to obtain an unbiased compressor is through ( scaled ) gradient sparsification ( Wangni et al. , 2018 ) or its generalized form , atomic decomposition ( Wang et al. , 2018 ) . However , such methods may be practically slow with performance degradation ( Vogels et al. , 2019 ) . Biased gradient compressors , on the other hand , have been more successful in practice since not only do they support a more impressive compression ratio , but also a test accuracy comparable with SGD can be obtained in many scenarios . Examples of such compressors contain top-k or ( unscaled ) random-k sparsification ( Lin et al. , 2018 ; Alistarh et al. , 2019 ) , signSGD ( Bernstein et al. , 2018a ; b ) , and PowerSGD ( Vogels et al. , 2019 ) . Due to the bias , such compressors typically rely on error feedback ( EF ) schemes to ensure convergence ( Stich et al. , 2018 ; Karimireddy et al. , 2019 ) . Memory concerns in communication efficient methods . Memory concerns are highly entangled with communication efficiency in distributed training . On the one hand , memory footprint can be alleviated through a smaller batchsize ( Krause et al. , 2016 ; Spring et al. , 2019 ) , leading to increased iterations as well as communication rounds per epoch . Hence , communication efficient methods are particularly useful for accelerating training time in such cases . On the other hand , existing communication efficient methods are challenged by limited memory due to several reasons . First , a smaller batch size , which enlarges the variance , degrades the applicability of unbiased gradient compressors . Confirmed in our experiments , the exploded variance typically calls for a small step size which decelerates convergence . Unbiased gradient compressors such as quantization , are also hindered by the need of AllGather ( Xu et al. , 2020 ) , since peak memory proportional to the number of workers is required . This renders them inefficient for the parsimonious setting , especially when the number of workers is large . Memory is also not handled well in methods with biased gradient compressors due to the need of error feedback , where additional memory the same as model size has to be consumed to keep track of accumulated compression error . Note that although no memory footprint is introduced in Local SGD ( Stich , 2019 ) , it suffers from variance blown problem as well . More importantly , local SGD can not be integrated with ZeRO 3 ( Rajbhandari et al. , 2020 ) , state-ofthe-art method for coping with limited memory , due to partitioned optimizer states among workers . In this work , we will focus on EF-based methods because of the capability of adopting a biased gradient compressor , which is typically contractive thus more robust to the gradient variance in a small batchsize setting . The additional memory consumption of EFSGD over SGD is the local error vector , which is delicate and critical for ensuring convergence . Meanwhile , the designed manner for memory efficiency can only come with lightweight and negligible runtime overhead , otherwise it violates the ultimate goal of communication efficient methods . Contractive error feedback is thus introduced to address these challenges . In a nutshell , our contributions can be summarized as : • To the best of our knowledge , this is the first work that systematically studies memory concerns in communication efficient methods . Contractive error feedback ( ConEF ) , is introduced to effectively manage memory given the preference of an allreducable and biased gradient compressor . • Convergence of ConEF is established . Our theoretical results suggest a tradeoff between memory efficiency and a faster convergence , and ConEF finds the sweet spot . • ConEF is capable of saving 80 % – 90 % additional memory of EFSGD on various learning problems such as image classification , language modeling , and machine translation . With almost the same runtime of EFSGD , ConEF achieves 1.3x – 5x speedup over SGD . Though test performance slightly drops on smaller models , the proposed method is more useful and reliable for larger networks where improved test accuracy over EFSGD is observed . Notation . Bold lowercase ( uppercase ) letters denote vectors ( matrices ) ; ‖x‖ stands for ` 2 norm of x ; and 〈x , y〉 denotes the inner product of x and y . In addition , we use [ x ] i ( [ X ] i , j ) to denote the i-th entry of vector x ( i , j-th entry of matrix X ) . 2 PRELIMINARIES . This section briefly recaps EFSGD ( Stich et al. , 2018 ; Karimireddy et al. , 2019 ) listed under Alg . 1 . We use git to denote the stochastic gradient at xt on worker i , and assume that the stochastic gradients are mutually independent among different workers per iteration . In line 5 , the scaled Algorithm 1 EFSGD 1 : Initialize : x0 ∈ Rd , ei0 = 0 ∈ Rd , ∀i , η 2 : for t = 0 , 1 , . . . , T − 1 do 3 : assert xt = xit for every worker i 4 : for worker i = 1 , . . . , N in parallel do 5 : pit = ηg i t + e i t 6 : ∆it = Q ( pit ) 7 : ∆t = Aggregate ( ∆it , ∀i ) 8 : xt+1 = xt −∆t 9 : eit+1 = p i t −∆it 10 : end for 11 : end for Algorithm 2 ConEF 1 : Initialize : x0 ∈ Rd , ei0 = 0 ∈ Rd , ∀i , η 2 : for t = 0 , 1 , . . . , T − 1 do 3 : assert xt = xit for every worker i 4 : for worker i = 1 , . . . , N in parallel do 5 : pit = ηg i t + e i t 6 : ∆it = Q ( pit ) 7 : ∆t = Aggregate ( ∆it , ∀i ) 8 : xt+1 = xt −∆t 9 : eit+1 = C ( pit −∆it ) 10 : end for 11 : end for stochastic gradient is augmented via accumulated compression error eit , then compressed and communicated . There is no restriction to the biasedness of gradient compressor Q . Examples of Q include ( scaled ) sign , random/top-k , and powerSGD . The “ aggregate ” in line 7 refers to ∆t = 1 N ∑ i ∆ i t , and the compression error e i t+1 is recomputed in line 9 after updating model xt . More implementation details on communication of gradients , and reasons for the preference of AllReduce , are deferred to Apdx . A.1 to save space . Besides the widely appreciated allreducable gradient compressors , those biased ones are more tailored for settings where the inclination is to rely on a smaller batchsize for memory saving . This is because biased gradient compressors are usually contractive ; see more details shortly in Assumption 4 . However , due to the need of error feedback to ensure convergence , the additional memory consumption of eit can confine the applicability of biased compressors as well as EFSGD in parsimonious setups , and this issue is somehow omitted by existing works . Other error feedback variants , e.g. , ( Wu et al. , 2018 ; Basu et al. , 2019 ; Xu et al. , 2021 ; Richtárik et al. , 2021 ) , also rely on additional vectors for the compression error , hence may benefit from the proposed technique as well . 3 MEMORY SAVING VIA CONTRACTIVE ERROR FEEDBACK . To endow EFSGD with memory efficiency , contractive error feedback is introduced in this section . The key idea is to apply another compressor C on the error vector eit such that the memory footprint can be mitigated . The proposed method is summarized in Alg . 2 , where the name contractive error feedback ( ConEF ) originates from the fact that the error vector is represented in its compressed format . Note that the encoding and decoding of compression in Alg . 2 are omitted for notational convenience . Q and C are adopted to denote the gradient and error compressors , respectively , to highlight their different roles . The gradient compressor Q is general and a biased one is more recommended for small batchsizes that can squeeze more memory out . An unbiased error compressor C is preferable since it is helpful for generalization as we shall discuss later in Section 4 . Moreover , the implementation of C has to be lightweight to avoid slowing down runtime . We will focus on theoretical properties first , and practical guidances for the choice of C are deferred to Section 4 .

Summary: The paper introduces a two-stage error correction procedure for approximated (doubly compressed) gradients. Multiple theorems are derived to prove convergence and scaling properties. The algorithm can be combined with any compression algorithm for the gradient. The algorithm is applied to CIFAR-10 with results showing favorable results.