Global Momentum Compression for Sparse Communication in Distributed SGD

1 INTRODUCTION . Many machine learning models can be formulated as the following empirical risk minimization problem : min w∈Rd F ( w ) : = 1 n n∑ i=1 f ( w ; ξi ) , ( 1 ) where w denotes the model parameter , ξi denotes the ith training data , n is the number of training data , and d is the size of the model . For example , let ξi = ( ai , yi ) , where ai denotes the feature of the ith training data and yi denotes the label . Then in logistic regression f ( w ; ξi ) = log ( 1 + e−yia T i w ) + λ2 ‖w‖ 2 , and in SVM f ( w ; ξi ) = max ( 0 , 1 − yiaTi w ) + λ2 ‖w‖ 2 . Many deep learning models can also be formulated as ( 1 ) . One of the efficient ways to solve ( 1 ) is stochastic gradient descent ( SGD ) ( Robbins & Monro , 1951 ) . In each iteration , SGD calculates one stochastic gradient ∇f ( w ; ξi ) and updates w by w← w−η∇f ( w ; ξi ) , or updates w with a mini-batch of stochastic gradients . Inspired by momentum and nesterov ’ s accelerated gradient descent , momentum SGD ( MSGD ) ( Polyak , 1964 ; Tseng , 1998 ; Lan , 2012 ; Kingma & Ba , 2015 ) has been proposed and widely used in machine learning . In practice , MSGD can achieve better performance than SGD ( Krizhevsky et al. , 2012 ; Sutskever et al. , 2013 ) . Many machine learning platforms like TensorFlow , PyTorch , and MXNet adopt MSGD as one of the optimization methods . With the rapid growth of data , distributed SGD ( DSGD ) ( Dekel et al. , 2012 ; Li et al. , 2014b ) has attracted much attention since it can parallelly calculate a batch of stochastic gradients . DSGD can be formulated as follows : wt+1 = wt − ηt p p∑ k=1 gt , k , ( 2 ) where p is the number of workers , gt , k is the stochastic gradient or a mini-batch of stochastic gradients calculated by the kth worker at the tth iteration . DSGD can be implemented on distributed frameworks like parameter server and all-reduce framework . Each worker calculates gt , k and sends it to the server or other workers for updating w. Recently , more and more large models , such as deep learning models , are used in machine learning to improve the generalization ability . In large models , gt , k is a high dimensional vector . Due to the latency and limited bandwidth of network , communication cost has become the bottleneck of traditional DSGD or distributed MSGD ( DMSGD ) . For example , when we implement DSGD on parameter server , the server needs to receive p high dimension vectors from workers , which will lead to communication traffic jam and make the convergence of DSGD slow . Hence , we need to compress gt , k to reduce the communication cost . Recently , researchers have proposed two main categories of communication compression techniques for reducing communication cost in DSGD and DMSGD . The first category is quantization ( Wen et al. , 2017 ; Alistarh et al. , 2017 ; Jiang & Agrawal , 2018 ) . In machine learning problems , 32-bit float number is typically adopted for representation . Quantization methods quantize the value ( gradient or parameter ) representation from 32 bits to some low bit-width like 8 bits or 4 bits . Since the quantized gradients in most methods are an unbiased estimation of the original ones , the convergence rate of these methods has the same order of magnitude as that of DSGD , but slower due to the extra quantization variance . It is easy to find that the communication cost can be reduced by 31 fold in the ideal case . In practice , at least 4 bits should be adopted for representation in most cases to keep original accuracy . In these cases , the communication cost is reduced by 7 fold . The other category is based on sparsified gradient , which is called sparse communication . In sparse communication , after calculating the update vector gt , k at each iteration , each worker only sends a subset of coordinates in gt , k , denoted as S ( gt , k ) , to the server or other workers . Here , S ( gt , k ) is a sparse vector and hence it can reduce the communication cost . In recent works ( Aji & Heafield , 2017 ; Lin et al. , 2018 ) , each worker will typically remember those values which are not sent to the server , i.e. , gt , k−S ( gt , k ) , and store it in the memory rather than dropping them . The gt , k−S ( gt , k ) is called memory gradient and will be used to calculate the next update vector gt+1 , k . This is intuitively necessary because a subset of coordinates in the stochastic gradient can not reflect the real descent direction and can make errors with higher probability than the original stochastic gradient . This memory gradient based sparse communication strategy has been widely adopted by recent communication compression methods and has achieved better performance than quantization methods and other sparse communication methods without memory gradient . In these memory gradient based sparse communication methods , some are for vanilla SGD ( Aji & Heafield , 2017 ; Alistarh et al. , 2018 ; Stich et al. , 2018 ; Karimireddy et al. , 2019 ; Tang et al. , 2019 ) . The convergence rate of vanilla SGD with sparse communication has been proved in ( Alistarh et al. , 2018 ; Stich et al. , 2018 ; Karimireddy et al. , 2019 ; Tang et al. , 2019 ) . Very recently , there has appeared one sparse communication method for distributed MSGD ( DMSGD ) , called deep gradient compression ( DGC ) ( Lin et al. , 2018 ) , which has achieved better performance than vanilla DSGD with sparse communication in practice . However , the theory about the convergence of DGC is still lack . Furthermore , although DGC uses momentum SGD , the momentum in DGC is calculated by each worker . Hence it is a local momentum without global information . In this paper , we propose a novel method , called global momentum compression ( GMC ) , for sparse communication in DMSGD which includes DSGD as a special case . The main contributions of this paper are summarized as follows : • GMC combines memory gradient and momentum SGD to achieve sparse communication in DMSGD ( DSGD ) . But different from DGC which adopts local momentum , GMC adopts global momentum . • We theoretically prove the convergence rate of GMC for both convex and non-convex problems . To the best of our knowledge , this is the first work that proves the convergence of DMSGD with sparse communication and memory gradient . • Empirical results show that , compared with the DMSGD counterpart without sparse communication , GMC can reduce the communication cost by approximately 100 fold with no loss of generalization accuracy . • GMC can also achieve comparable ( sometimes better ) performance with comparable com- munication compression ratio , with an extra theoretical guarantee . 2 PRELIMINARY . In this paper , we use ‖ · ‖ to denote L2 norm , use w∗ to denote the optimal solution of ( 1 ) , use ∇f ( w ; I ) , 1|I| ∑ ξi∈I ∇f ( w ; ξi ) to denote one stochastic gradient with respect to a mini-batch of samples I such that EI [ ∇f ( w ; I ) |w ] = ∇F ( w ) , use to denote element-wise product , use 1 to denote the vector ( 1 , 1 , . . . , 1 ) T ∈ Rd , use ej to denote the vector ( 0 , . . . , 1 , . . . 0 ) T ∈ Rd with only the jth coordinate being 1 , use I to denote an identity matrix . For a vector a , we use a ( j ) to denote its jth coordinate value . ‖a‖0 denotes the number of non-zero values in a . Definition 1 ( smooth function ) Function h ( · ) is L-smooth ( L > 0 ) if ∀w , w′ , |h ( w ) − h ( w′ ) −∇h ( w′ ) T ( w −w′ ) | ≤ L 2 ‖w −w′‖2 . Definition 2 ( convex function ) Function h ( · ) is convex if ∀w , w′ , |h ( w ) − h ( w′ ) −∇h ( w′ ) T ( w −w′ ) | ≥ µ 2 ‖w −w′‖2 , where µ ≥ 0 . If µ > 0 , h ( · ) is called a µ-strongly convex function . 2.1 DISTRIBUTED MOMENTUM SGD . The widely used momentum SGD ( MSGD ) ( Polyak , 1964 ) for solving ( 1 ) can be written as gt =βgt−1 + ηt∇f ( wt ; It ) , wt+1 =wt − gt , where β ∈ [ 0 , 1 ) . The gt is the Polyak ’ s momentum and ∇f ( wt ; It ) is an unbiased estimation of ∇F ( wt ) . Since gt = ( wt −wt+1 ) , MSGD can also be written as wt+1 =wt − ηt ( ∇f ( wt ; It ) + β ηt ( wt−1 −wt ) ) . ( 3 ) Please note that if β = 0 , MSGD degenerates to SGD . One simple way to implement distributed MSGD ( DMSGD ) is that each worker parallelly calculates some stochastic gradients and then the stochastic gradients of all workers are aggregated to get ∇f ( wt ; It ) + β ( wt−1 − wt ) /ηt . The update process of w in this way is equivalent to the serial MSGD . We call ( wt−1 −wt ) /ηt the global momentum , because it captures the global information from all workers . Another way to implement DMSGD is using local momentum : gt , k =βgt−1 , k + ηtvt , k , k = 1 , 2 , . . . , p , wt+1 =wt − p∑ k=1 gt , k , where vt , k is the stochastic gradient calculated by the kth worker and ∑p k=1 vt , k = ∇f ( wt ; It ) . gt−1 , k is the local momentum . Since ∑p k=1 gt−1 , k = ( wt−1 − wt ) , this DMSGD with local momentum can also be written as the formulation in ( 3 ) . Hence , the global momentum contains all information of the local momentums . We will find that DGC ( Lin et al. , 2018 ) is mainly based on the local momentum while GMC is mainly based on the global momentum . Hence , each worker in DGC can not capture the global information from its own local momentum , while that in GMC can capture the global information from the global momentum even if sparse communication is adopted . In the later section , we will see that global momentum is better than local momentum when using memory gradient for sparse communication . Algorithm 1 Global Momentum Compression ( GMC ) on Parameter Server 1 : Initialization : p workers , w−1 = w0 , β ∈ [ 0 , 1 ) , batch size b ; 2 : Set g0 , k = u0 , k = 0 , k = 1 , . . . , p , 3 : for t = 0 , 1 , 2 , ... T − 1 do 4 : Workers : 5 : for k = 1 , 2 . . . , p , each worker parallelly do 6 : if t > 0 then 7 : Receive wt −wt−1 from the server ; 8 : Get wt by wt = wt−1 + ( wt −wt−1 ) ; 9 : end if 10 : Randomly pick a mini-batch of training data It , k ⊆ Dk with |It , k| = b ; 11 : gt , k = 1 pb ∑ ξi∈It , k ∇f ( wt ; ξi ) − β pηt ( wt −wt−1 ) ; 12 : Generate a sparse vector mt , k ∈ { 0 , 1 } d ; 13 : Send mt , k ( gt , k + ut , k ) to the server ; 14 : ut+1 , k = ( 1−mt , k ) ( gt , k + ut , k ) ; 15 : end for 16 : Server : 17 : wt+1 = wt − ηt ∑p k=1 mt , k ( gt , k + ut , k ) ; 18 : Send wt+1 −wt to workers ; 19 : end for

Summary: Gradient sparsification is an important technique to reduce the communication overhead in distributed training. In this paper, the authors proposed a training method called global momentum compression (GMC) for distributed momentum SGD with sparse gradient. Following existing gradient sparsification techniques such DGC, GMC is also built up on the memory gradient approach; the major distinction between GMC and existing techniques is that GMC keeps track of global gradient to maintain the memory gradient, while the existing technique keeps track of worker-local gradients for memory gradient. The primary contributions in the paper are as the following: