Title: GLOBAL MOMENTUM COMPRESSION FOR SPARSE COMMUNICATION IN DISTRIBUTED SGD

Abstract: With the rapid growth of data, distributed stochastic gradient descent (DSGD) has been widely used for solving large-scale machine learning problems. Due to the latency and limited bandwidth of network, communication has become the bottleneck of DSGD when we need to train large scale models, like deep neural networks. Communication compression with sparsified gradient, abbreviated as sparse communication, has been widely used for reducing communication cost in DSGD. Recently, there has appeared one method, called deep gradient compression (DGC), to combine memory gradient and momentum SGD for sparse communication. DGC has achieved promising performance in practice. However, the theory about the convergence of DGC is lack. In this paper, we propose a novel method, called global momentum compression (GMC), for sparse communication in DSGD. GMC also combines memory gradient and momentum SGD. 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 distributed momentum SGD (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 compared with DGC, with an extra theoretical guarantee.

1. INTRODUCTION:

Many machine learning models can be formulated as the following empirical risk minimization problem:

EQUATION

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 = (a i , y i ), where a i denotes the feature of the ith training data and y i 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 -y i a T i 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:

EQUATION

where p is the number of workers, g t,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 g t,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, g t,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 g t,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 g t,k at each iteration, each worker only sends a subset of coordinates in g t,k , denoted as S(g t,k ), to the server or other workers. Here, S(g t,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., g t,k -S(g t,k ), and store it in the memory rather than dropping them. The g t,k -S(g t,k ) is called memory gradient and will be used to calculate the next update vector g t+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 communication compression ratio, with an extra theoretical guarantee.

2. PRELIMINARY:

In this paper, we use • to denote L 2 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 E I [∇f (w; I)|w] = ∇F (w), use to denote element-wise product, use 1 to denote the vector (1, 1, . . . , 1) T ∈ R d , use e j to denote the vector (0, . . . , 1, . . . 0) T ∈ R d 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 g t =βg t-1 + η t ∇f (w t ; I t ), w t+1 =w t -g t , where β ∈ [0, 1). The g t is the Polyak's momentum and ∇f (w t ; I t ) is an unbiased estimation of ∇F (w t ). Since g t = (w t -w t+1 ), MSGD can also be written as

EQUATION

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 (w t ; I t ) + β(w t-1 -w t )/η t . The update process of w in this way is equivalent to the serial MSGD. We call (w t-1 -w t )/η t the global momentum, because it captures the global information from all workers.

Another way to implement DMSGD is using local momentum:

g t,k =βg t-1,k + η t v t,k , k = 1, 2, . . . , p, w t+1 =w t - p k=1 g t,k ,

where v t,k is the stochastic gradient calculated by the kth worker and p k=1 v t,k = ∇f (w t ; I t ). g t-1,k is the local momentum. Since p k=1 g t-1,k = (w t-1 -w t ), 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 cannot 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 = w 0 , β ∈ [0, 1), batch size b; 2: Set g 0,k = u 0,k = 0, k = 1, . . . , p, 3: for t = 0, 1, 2, .. • Each worker generates a sparse vector m t,k and sends m t,k (g t,k + u t,k ) to the server or other workers;

• Each worker updates u t+1,k = (1 -m t,k ) (g t,k + u t,k );

• Update parameter w t+1 = w t -η t p k=1 m t,k (g t,k + u t,k );

Below, we introduce the framework and two essential elements of GMC: memory gradient u t,k and global momentum (w t-1 -w t )/η t .

3.1. FRAMEWORK OF GMC:

GMC can be easily implemented on the all-reduce distributed framework in which each worker sends the sparse vector m t,k (g t,k + u t,k ) to all the other workers, then each worker updates w t+1 after receiving the sparse vectors from other workers.

Recently, parameter server (Li et al., 2014a) has been one of the most popular distributed frameworks in machine learning. GMC can also be implemented on parameter server. Although the theories in this paper can be adopted for the all-reduce framework, in this paper we adopt parameter server for illustration. The details of GMC implemented on parameter server are shown in Algorithm 1. The difference between GMC and traditional DSGD on parameter server is that in GMC after updating w t+1 , server will send w t+1 -w t , rather than w t+1 , to workers. Since m t,k is sparse, w t+1 -w t is sparse as well. Hence, sending w t+1 -w t can reduce the communication cost compared with sending w t . In our experiments, we find that GMC can make w t+1 -w t 0 ≤ 0.01d with no loss of accuracy when training large scale models. Workers can get w t+1 by w t+1 = w t + (w t+1 -w t ).

Remark 1 We can also use the sparse communication technique on the server (denoted as GMC * ) to make (w t+1 -w t ) sparser than that in GMC so that when the server sends w t+1 -w t to workers, the communication cost can be further reduced, compared with GMC. The convergence of GMC * can also be theoretically proved. But in practice, GMC * will be slightly worse than GMC in terms of accuracy. Due to the space limitation, we move both theory and experiments about GMC * to Appendix A.

3.2. NECESSITY OF MEMORY GRADIENT:

In GMC, after sending a sparse vector m t,k (g t,k + u t,k ) to the server, each worker will remember the coordinates which are not sent and store them in u t+1,k :

EQUATION

Hence we call u t,k the memory gradient, which is important for the convergence guarantee of GMC.

Here, we give an intuitive explanation about why GMC needs to remember the coordinate values which are not sent. We consider the simple case that β = 0, which means g t,k is a stochastic gradient of F (w) and GMC degenerates to (Aji & Heafield, 2017) . Since m t,k is a sparse vector, GMC can be seen as a method achieving sparse communication by combining stochastic coordinate descent (SCD) (Nesterov, 2012) and DSGD. In SCD, each -∇F (w) (j) e j denotes a true descent direction. When we use a stochastic gradient ∇f (w; I) to replace ∇F (w), -∇f (w; I) (j) e j will make errors with high probability, especially when m t,k adopts the strategy that choosing s coordinates with the largest absolute values (denoted as top-s strategy).

For further explaining the importance of memory gradient, we consider the following simple example:

let n = p = 2, f (w; ξ 1 ) = (-α, )w, f (w; ξ 2 ) = (α + , γ)w, where w ∈ [-1, 0] × [-1, 0], 0 < < α < γ < α + . f (w; ξ 1

) is on the first worker and f (w; ξ 2 ) is on the second worker. Then we run GMC with β = 0 to solve min

F (w) = 1 2 (f (w; ξ 1 ) + f (w; ξ 2 )).

The optimal solution is w * = (-1, -1) T . We adopt the top-1 strategy for m t,k in this example.

If we do not use the memory gradient, which means each worker directly sends m t,k g t,k to the server, then the two workers will send (-α/2, 0) T and ((α + )/2, 0) T respectively to the server. The parameter is updated by w ← w -η t ( /2, 0). We observe that w (2) will never be updated. This is due to the pseudo large gradient value which cheats m t,k . Since ∇F (w) = ( /2, (γ + )/2) T , we can see that the second coordinate has the true large gradient value and we should have mainly focused on updating w (2) . However, in the two stochastic gradients ∇f (w; ξ 1 ) and ∇f (w; ξ 2 ), the first coordinate has larger absolute value. Hence, ∇f (w; ξ k ) cheats m t,k which leads to the error.

If we use memory gradient, at the beginning, m t,1 = m t,2 = (1, 0) T . After some iterations, m t,1 = (0, 1) T and m t,2 = (0, 1) T due to the memory gradient. Specifically, let t 1 , t 2 be two integers satisfying α/ ≤ t 1 < α/ + 1, (α + )/γ ≤ t 2 < (α + )/γ + 1, then it is easy to verify that m ct1,1 = (0, 1) T , m ct2,2 = (0, 1) T , ∀c ≥ 1. It implies that if we use the memory gradient, both w (1) and w (2) will be updated. Hence, GMC can make w converge to the optimum (-1, -1) T .

Hence, the memory gradient is necessary for sparse communication. It can overcome the disadvantage of combining DSGD and SCD.

3.3. BENEFIT OF GLOBAL MOMENTUM:

In GMC, each worker calculates g t,k as

EQUATION

When β = 0, it degenerates to DSGD with sparse communication, which is also called gradient dropping in (Aji & Heafield, 2017) . We use g

(GD)

t,k to denote this degenerated version with β = 0. We can see that GMC uses the global momentum (w t-1 -w t )/η t . While in DGC (Lin et al., 2018) , the g

EQUATION

which uses the local momentum g (DGC)

t-1,k . We use u

(DGC) t,k

to denote the memory gradient in DGC.

For sparse m t,k , we can find that g

(DGC) t,k and u (DGC) t,k

only contain the information based on the local data, while g t,k and u t,k can contain the global information due to the global momentum. Hence, global momentum is better than local momentum when using memory gradient for sparse communication. Assume that E[F (w t )] converges to F (w * ), then w t -w t-1 denotes the descent direction with high probability. Since I t,k ⊆ D k , which only contains partial information of the whole training data, the global momentum (w t-1 -w t )/η t can make compensation for the error between the stochastic gradient and the full gradient. Specifically, if (w t -w t-1 ) T (-∇F (w t )) ≥ 0, then we get that

g T t,k ∇F (w t ) =( 1 pb ξi∈I t,k ∇f (w t ; ξ i ) - β pη t (w t -w t-1 )) T ∇F (w t ) ≥ 1 pb ξi∈I t,k ∇f (w t ; ξ i ) T ∇F (w t ) =(g (GD) t,k ) T ∇F (w t ).

It implies that g t,k is a better estimation of ∇F (w t ) than g

(GD)

t,k . The experiments in (Lin et al., 2018) also show that (Aji & Heafield, 2017) has some loss of accuracy with m t,k 0 ≈ 0.001d, while we can set such a sparsity for GMC with no loss of accuracy in practice. Hence, GMC is better than (Aji & Heafield, 2017) .

Remark 2 There are some other ways to combine global momentum and memory gradient. For example, we can put the global momentum on the server. However, these ways lead to some loss of performance. More discussions are put in Appendix B.

4. CONVERGENCE OF GMC:

In this section, we prove the convergence rate of GMC for both convex and non-convex problems. The detailed proofs are put in Appendix C. We define a diagonal matrix M t,k ∈ R d×d such that diag(M t,k ) = m t,k to replace the symbol . Then we obtain

EQUATION

EQUATION

For convenience, we denote ∇f (w;

I t ) = 1 pb p k=1 ξi∈I t,k ∇f (w t ; ξ i ), ũt = p k=1 u t,k , gt = p k=1 g t,k

. By eliminating M t,k , we obtain w t+1 =w t -η t ∇f (w t ; I t ) + β(w t -w t-1 ) -η t ũt + η t ũt+1 .

We can find that if we do not use the sparse communication technique, then ũt = 0 and ( 9) is the same as momentum SGD in (3).

First, we give the following lemma:

Lemma 1 Let x t = w t + β 1-β (w t -w t-1 ) -ηt 1-β ũt .

Then we have:

EQUATION

Equation ( 10) is similar to the update equation in SGD except the additional term ũt+1 . Inspired by this, we only need to prove the convergence of x t and x t -w t . To get the convergence results, we further make the following three assumptions:

Assumption 1 (bounded gradient) E ∇f (w t ; ξ i ) 2 ≤ G 2 , ∀t. Assumption 2 (bounded memory) E ũt 2 ≤ U 2 , ∀t. Assumption 3 (bounded parameter) E x t -w * 2 ≤ D 2 , ∀t.

Remark 3 Assumption 1 is common in stochastic optimization. Assumption 2 is easy to be guaranteed. We discuss it in Appendix D. Assumption 3 is only for convenience in the analysis for strongly convex and convex cases. We can add one projection operation on Equation ( 7) to guarantee Assumption 3 (if both w t and ũt are bounded, x t is bounded as well), which can be written as

w t+1 = Π Ω (w t -η t p k=1 M t,k (g t,k + u t,k )).

The convergence can still be guaranteed if we use the projection operation. More details are put in Appendix E.

For the new variable x t in Lemma 1, the gap between x t and w t has the following property:

Lemma 2 With Assumption 1 and Assumption 2, we get that

• If η t = η, then E x t -w t 2 ≤ [ 2β 2 (2G 2 + 2U 2 ) (1 -β) 4 + 2U 2 (1 -β) 2 ]η 2 .

(11)

EQUATION

)

where

A = max t0 t=0 { η 2 t +βη 2 t-1 +...+β t η 2 0 η 2 t+1 }, t 0 satisfies ( t0+1+q t0+2+q ) 2ρ ≥ 1+β 2 .

Lemma 2 implies that E x t -w t 2 ≤ O(η 2 t ), i.e. the distance between x t and w t is O(η t ). For convenience, below we use the constant C r,q,ρ,β,η to denote the upper bound of E x t -w t 2 /η 2 t . If r = 0, C r,q,ρ,β,η denotes the upper bound in (11) and if η = 0, C r,q,ρ,β,η denotes the upper bound in ( 12). Then we have the following convergence results:

Theorem 1 (strongly convex case) Let F (•) be L-smooth and µ-strongly convex. With Assumption 1, Assumption 2, Assumption 3, and η t = 1-β µ(t+1) , m = T /2 , we obtain

1 m T -1 t=T -m E(F (w t ) -F (w * )) ≤ 3A + 2G C (1-β)/µ,1,1,β,0 (1 -β) µT .

where

A = max{µ 2 D 2 , 2(1 -β) C (1-β)/µ,1,1,β,0 LD + µU D + 2G 2 + 2U 2 }.

Theorem 2 (convex case) Let F (•) be a convex function. With Assumption 1, Assumption 2, Assumption 3, and

η t = 1-β √ t+1 , we obtain T -1 t=0 2 √ t + 1 E(F (w t ) -F (w * )) ≤ w 0 -w * 2 + A log(T ),

where

A = 2(1 -β) C 1-β,1,0.5,β,0 G + 2U D + 2G 2 + 2U 2 . Since T -1 t=0 1 √ t+1 = O( √ T ), Theorem 2 implies that if the objective function is convex, GMC has a convergence rate of O(log(T )/ √ T ).

Theorem 3 (non-convex case) Let F (•) be L-smooth. With Assumption 1, Assumption 2, and η t = η being a constant, we obtain

1 (1 -β)T T -1 t=0 E ∇F (w t ) 2 ≤ F (w 0 ) -F (w * ) T η + Aη, where A = LG √ C 0,0,1,β,η 1-β + LG 2 2(1-β) 2 .

By taking η = O(1/ √ T ), it is easy to find that GMC has a convergence rate of O(1/ √ T ), if the objective function is non-convex.

5. EXPERIMENTS:

We conduct experiments on a PyTorch based parameter server framework with one server and eight workers. Each worker has access to one TITAN Xp GPU. We compare GMC with distributed momentum SGD (DMSGD) and DGC. In the experiments of (Lin et al., 2018) , DGC gets far better performance on both accuracy and communication compression ratio than quantization methods. Hence, we do not compare with quantization methods in this paper. We use warm-up strategy (run DMSGD 5 epochs) for both GMC and DGC. The momentum factor masking trick is used in DGC. The β is set as 0.9. In our experiments, we consider the communication cost on the server which is the busiest node. It includes receiving vectors from the p workers and sending one vector to the p workers. So the cost of DMSGD is 2pd. In GMC and DGC, since m t,k is sparse, workers send the vectors using the structure of (key, value). The cost of each (key, value) is 2. Server sends w t+1 -w t using this structure as well. Hence, the cost of GMC and DGC is 2( p k=1 m t,k 0 + p w t+1 -w t 0 ). We define the communication compression ratio (CR) of GMC and DGC as:

CR = 1 T T +4 t=5 1 pd (p w t+1 -w t 0 + p k=1 m t,k 0 ).

The CR of DMSGD is 100% (no compression). Here, all numbers have the same unit (float value).

5.1. CONVEX MODEL:

We use the dataset MNIST and the model logistic regression (LR) to evaluate GMC on convex problems. Since the sizes of dataset and model are small, m t,k directly adopts top-s strategy with m t,k 0 = s where s = 0.01d or 0.001d. We use 8 workers for this experiment. The weight decay is 0.0001 and total batch size is 128. The result is shown in Figure 1 . We can see that GMC gets similar performance as DGC and DMSGD. Figure 1 : Comparison on convex model. The CRs of DGC with s = 0.01d and 0.001d are 7.87% and 0.797% respectively. The CRs of GMC with s = 0.01d and 0.001d are 8.0% and 0.797% respectively.

5.2. NON-CONVEX MODEL:

For large models, we use the approximate top-s strategy for m t,k : given a vector a, we first randomly choose a subset S ⊂ {1, . . . , d} and |S| is 0.001d ∼ 0.01d. We get the threshold θ such that |{j||a (j) | ≥ θ, j ∈ S}| = 0.001|S|. Then we set m t,k by choosing the indexes in {j||a (j) | ≥ θ, j = 1, 2, . . . , d}. It implies that m t,k 0 is approximately 0.001d.

First, we use the dataset CIFAR-10 and three popular deep models (AlexNet, ResNet20, ResNet56) to evaluate GMC on non-convex problems. We use both 4 and 8 workers with the total batch size 128. The weight decay is set as 0.0001. Figure 2 shows the learning process about training loss and test accuracy of the three methods. We can find that GMC and DGC have the same performance as that of DMSGD on ResNet20 and ResNet56. Compared to the two residual networks, AlexNet has more parameters. We can see that GMC also gets the performance as that of DMSGD and is slightly better than DGC. Table 1 shows the finial training results of the three methods. Compared with DMSGD, GMC can reduce the communication cost by more than 100 fold with no loss of accuracy. This is far better than those quantization methods. Furthermore, GMC achieves slightly better accuracy with comparable communication compression ratio, compared with DGC.

We also evaluate GMC on larger dataset ImageNet and larger model ResNet18. We use 8 workers with total batch size 256. The weight decay is set as 0.0001. In Figure 3 , GMC gets the same 

6. CONCLUSION:

In this paper, we propose a novel method, called global momentum compression (GMC), for sparse communication in DMSGD (DSGD). 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 GMC can achieve state-of-the-art performance.

A USING SPARSE COMMUNICATION ON SERVER:

If we use the sparse communication technique on the server, the update for w t can be written as:

EQUATION

EQUATION

EQUATION

where g t,k is the same as that in GMC. We call this method GMC * . The m t,k and a t are generated by the worker and server respectively. The u t,k and ∆ t are the memory gradients on the worker and server respectively. We can call u t,k the local memory gradient and ∆ t the global memory gradient. Details are presented in Algorithm 2.

Algorithm 2 GMC * 1: Initialization: p workers,

w -1 = w 0 , β ∈ [0, 1), batch size b; 2: Set g 0,k = u 0,k = ∆ 0 = 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 w t -w t-1 from the server; 8:

Get w t by w t = w t-1 + (w t -w t-1 );

9: end if 10:

Randomly pick a mini-batch of training data

I t,k ⊆ D k with |I t,k | = b; 11: g t,k = 1 pb ξi∈I t,k ∇f (w t ; ξ i ) -β pηt (w t -w t-1 ); 12:

Generate a sparse vector m t,k ∈ {0, 1} d ; 13:

Send m t,k (g t,k + u t,k ) to the server; 14:

u t+1,k = (1 -m t,k ) (g t,k + u t,k ); 15:

end for 16:

Server:

17:

Generate a sparse vector a t ∈ {0, 1} d ; 18:

w t+1 = w t -η t a t [ p k=1 m t,k (g t,k + u t,k ) + ∆ t ];

19:

∆ t+1 = (1 -a t ) [ p k=1 m t,k (g t,k + u t,k ) + ∆ t ] 20:

Send w t+1 -w t to workers; 21: end for We eliminate m t,k , a t and obtain

w t+1 -η t ∆ t+1 =w t -η t [ p k=1 m t,k (g t,k + u t,k ) + ∆ t ] =w t -η t [g t + ũt -ũt+1 + ∆ t ] By denoting h t = ũt + ∆ t , we obtain w t+1 -η t h t+1 = w t -η t h t -η t ∇f (w t ; I t ) + β(w t -w t-1 )

The equation is the same as (9). Hence our convergence analysis is also suitable for GMC * . In Figure 4 , we set a t using top-s strategy with a t 0 = s = 0.001d, which implies that w t+1 -w t is nearly as sparse as m t,k . We can see that GMC * achieves the same loss as DMSGD. But it is slightly worse than GMC and DMSGD in terms of accuracy.

B OTHER WAYS FOR SPARSE COMMUNICATION:

It is easy to get one way to combine memory gradient and global momentum that we can put the global momentum w t -w t-1 on the server. We briefly present it in Algorithm 3. Algorithm 3 for t = 0, 1, 2, ...T -1 do Workers: for k = 1, 2 . . . , p, each worker parallelly do

g t,k = 1 pb ξi∈I t,k ∇f (w t ; ξ i ); Generate a sparse vector m t,k ∈ {0, 1} d ; Send m t,k (g t,k + u t,k ) to the server; u t+1,k = (1 -m t,k ) (g t,k + u t,k ); end for Server: w t+1 = w t -η t p k=1 m t,k (g t,k + u t,k ) + β(w t -w t-1 );

Send w t+1 -w t to workers; end for Let {w t }, {g t,k }, {u t,k } be the sequences produced by Algorithm 3, then we can get that

w t+1 =w t -η t p k=1 M t,k (g t,k + u t,k ) + β(w t -w t-1 ), u t+1,k =(I -M t,k )(g t,k + u t,k ), k = 1, 2, . . . , p.

By eliminating M t,k , we obtain the same equation as that in (9). Hence our convergence analysis is also suitable for Algorithm 3. The difference from GMC is that the memory gradient u t,k in Algorithm 3 only remembers the dropped stochastic gradient. After receiving sparse update vectors from workers, server updates parameter using the global momentum. Compared with GMC, the disadvantage is that its memory gradient does not contain the momentum information which can play the role of correcting the update direction. Hence, GMC gets better performance than Algorithm 3. (Lin et al., 2018) also points out that Algorithm 3 has some loss of convergence performance.

Another way is that we put η t in the memory gradient. We briefly present it in Algorithm 4.

Let {w t }, {g t,k }, {u t,k } be the sequences produced by Algorithm 4, then we can get that

w t+1 =w t - p k=1 M t,k (g t,k + u t,k ) u t+1,k =(I -M t,k )(g t,k + u t,k ), k = 1, 2, . . . , p.

By eliminating M t,k , we obtain the same equation as that in (9). Hence our convergence analysis is also suitable for Algorithm 3. However, the norm of memory gradient is larger than that of GMC. In particular, let y t,k = u t,k /η t , then we have 

y t+1,k = η t η t+1 (I -M t,k )(

C CONVERGENCE PROOFS:

First, we introduce the Young's Inequality as follow: ∀a, b

(a + b) 2 ≤ (1 + γ)a 2 + (1 + 1 γ )b 2 , ∀γ > 0,

which is used frequently in our proofs.

C.1 PROOF OF LEMMA 1

x t+1 =w t+1 + β 1 -β (w t+1 -w t ) - η t+1 1 -β ũt+1 = 1 1 -β (w t+1 -βw t ) - η t+1 1 -β ũt+1 = 1 1 -β (w t -η t ũt -η t ∇f (w t ; I t ) + β(w t -w t-1 ) -βw t + η t ũt+1 ) - η t+1 1 -β ũt+1 = 1 1 -β (w t -βw t-1 -η t ũt -η t ∇f (w t ; I t ) + η t ũt+1 ) - η t+1 1 -β ũt+1 = 1 1 -β (w t -βw t-1 -η t ũt -η t ∇f (w t ; I t )) + η t -η t+1 1 -β ũt+1 = 1 1 -β (w t -βw t-1 ) - η t 1 -β ũt - η t 1 -β ∇f (w t ; I t ) + η t -η t+1 1 -β ũt+1 =x t - η t 1 -β ∇f (w t ; I t ) + η t -η t+1 1 -β ũt+1 . C.2 PROOF OF LEMMA 2 Since w t+1 -w t = η t p k=1 M t,k (g t,k + u t,k ), we have ∀γ > 0, E w t+1 -w t 2 ≤η 2 t E gt + ũt 2 =η 2 t E β(w t-1 -w t )/η t + ∇f (w t ; I t ) + ũt 2 ≤(1 + γ)β 2 E w t -w t-1 2 + (1 + 1 γ )(2G 2 + 2U 2 )η 2 t .

Let γ = 1 β -1, we get that

E w t+1 -w t 2 ≤βE w t -w t-1 2 + η 2 t 1 -β (2G 2 + 2U 2 ) ≤ 2G 2 + 2U 2 1 -β (η 2 t + βη 2 t-1 + . . . + β t η 2 0 ). If η t = η, we get that E w t -w t-1 2 ≤ (2G 2 +2U 2 )η 2 (1-β) 2

. Then we obtain

E x t -w t 2 ≤ 2β 2 (1 -β) 2 E w t -w t-1 2 + 2η 2 (1 -β) 2 U 2 ≤[ 2β 2 (2G 2 + 2U 2 ) (1 -β) 4 + 2U 2 (1 -β) 2 ]η 2 . If η t = r (t+q) ρ , ρ ∈ [0.5, 1], r > 0, q ≥ 0, let t 0 satisfy ( t0+1+q t0+2+q ) 2ρ ≥ 1+β 2 , A = max t0 t=1 { η 2 t +βη 2 t-1 +...+β t η 2 0 η 2 t+1 }. Then by induction, we get that E w t -w t-1 2 ≤ (2G 2 +2U 2 ) max{(1-β)A,2}η 2 t (1-β) 2

. Then we obtain

E x t -w t 2 ≤ 2β 2 (1 -β) 2 E w t -w t-1 2 + 2η 2 t (1 -β) 2 U 2 ≤[ 2β 2 (2G 2 + 2U 2 ) max{(1 -β)A, 2} (1 -β) 4 + 2U 2 (1 -β) 2 ]η 2 t .

C.3 PROOF OF THEOREM 1:

According to Lemma 1, we get that

EQUATION

)

Using the strongly convex property, we get that

E x t+1 -w * 2 ≤(1 - 2 t )E x t -w * 2 + 2η 2 t 1 -β C (1-β)/µ,0,1,β,0 LD + 2(η t -η t+1 ) 1 -β U D + 2G 2 η 2 t (1 -β) 2 + 2(η t -η t+1 ) 2 (1 -β) 2 U 2 ≤(1 - 2 t )E x t -w * 2 + [2(1 -β) C (1-β)/µ,0,1,β,0 LD + 2µU D + 2G 2 + 2U 2 ] 1 µ 2 t 2 . Since E x 1 -w * 2 ≤ D 2 , we get that E x t -w * 2 ≤ A µ 2 t , ∀t ≥ 1, where A = max{µ 2 D 2 , 2(1 -β) C (1-β)/µ,0,1,β,0 LD + 2µU D + 2G 2 + 2U 2 }. Using (16) again, we get that 2E(F (x t ) -F (w * )) ≤ µt(E x t -w * 2 -E x t+1 -w * 2 ) + A µt .

By summing up the equation from t

= T -m to T -1, we get that 2 T -1 t=T -m E(F (x t ) -F (w * )) ≤ T -1 t=T -m [µt(E x t -w * 2 -E x t+1 -w * 2 ) + A µt ] ≤µ(T -m)E x T -m -w * 2 + T -1 t=T -m (µE x t -w * 2 + A µt ) ≤µ(T -m)E x T -m -w * 2 + T -1 t=T -m 2A µt ≤ A µ + T -1 t=T -m 2A µt .

By setting m = T /2 , we get that

T -1 t=T -m 1 t ≤ 1. Then 1 m T -1 t=T -m E(F (x t ) -F (w * )) ≤ 3A µT . Since E[F (x t ) -F (w t )] ≥ E∇F (w t ) T (x t -w t ) ≥ -G C (1-β)/µ,0,1,β,0 1-β µt , we get the result 1 m T -1 t=T -m E(F (w t ) -F (w * )) ≤ 3A + 2G C (1-β)/µ,0,1,β,0 (1 -β) µT .

C.4 PROOF OF THEOREM 2:

Similar to (16), we get that

EQUATION

) Since η t = 1-β √ t+1 , we get that T -1 t=0 2 √ t + 1 E(F (w t ) -F (w * )) ≤ w 0 -w * 2 + T -1 t=0 A t + 1 ≤ w 0 -w * 2 + A log(T ),

where A = 2(1 -β) C 1-β,1,0.5,β,0 G + 2U D + 2G LG θ.

2 η 2 2(1 -β) 2 ≤EF (x t ) - η 1 -β E[ ∇F (w t ) 2 -LG x t -w t ] + LG 2 η 2 2(1 -β) 2 ≤EF (x t ) - η 1 -β E ∇F (w t ) 2 + LG C 0,0,1,β,η η 2 1 -β + LG 2 η 2 2(1 -β) 2

E MORE DISCUSSION ABOUT ASSUMPTION 3:

If we do not use Assumption 3 in the proofs of Theorem 1 and Theorem 2, we can add one projection operation on equation ( 7), which can be written as

w t+1 = Π Ω (w t -η t p k=1 M t,k (g t,k + u t,k )),

where Ω = {w| w ≤ W } and w * ≤ W 2 . Then w t is bounded. Hence, we reasonably assume that ∀t, i, ũt ≤ U , ∇f (w t ; ξ i ) ≤ G and we set η t ≤ η 0 ≤ (1-β)W

2U:

. Please note that the U can be irrelevant to Ω according to Appendix D. We still define x t as

x t = w t + β 1 -β (w t -w t-1 ) - η t 1 -β ũt .

Then x t is bounded as well.

To get the convergence results for strongly convex and convex functions, we need to prove two essential inequalities: E x t -w t 2 ≤ O(η 2 t ) and x t+1 -w * 2 ≤ x t -w * -ηt 1-β ∇f (w t ; I t ) + By setting γ = 1-β β , we obtain w t+1 -w t 2 ≤ β w t -w t-1 2 + 2G 2 +2U 2 1-β η 2 t . Then using the proof in Lemma 2, we obtain the first inequality that E x t -w t 2 ≤ O(η 2 t ). Next, we get the relation between x t+1 and x t . According to the definition of x t , we have

x t+1 -w * 2 = w t+1 + β 1 -β (w t+1 -w t ) - η t+1 1 -β ũt+1 -w * 2 = 1 (1 -β) 2 w t+1 -(η t+1 ũt+1 + βw t + (1 -β)w * ) 2 .

Since ũt ≤ U uniformly and w * ≤ W 2 , η t ≤ η 0 ≤ (1-β)W

, we get that

η t+1 ũt+1 + βw t + (1 -β)w * ≤η t+1 U + βW + (1 -β) W 2 ≤ W.

Then we obtain η t+1 ũt+1 + βw t + (1 -β)w * ∈ Ω and

x t+1 -w * 2 ≤ 1 (1 -β) 2 w t -η t p k=1 M t,k (g t,k + u t,k ) -(η t+1 ũt+1 + βw t + (1 -β)w * ) 2 = 1 (1 -β) 2 w t -η t (g t + ũt -ũt+1 ) -(η t+1 ũt+1 + βw t + (1 -β)w * ) 2 = x t -w * -η t 1 -β ∇f (w t ; I t ) + η t -η t+1 1 -β ũt+1 2 .

Then similar to the analysis in Theorem 1 and Theorem 2, we can prove the convergence for strongly convex and convex cases.

:



Figures and Tables:

Figure 2: Learning process of using different non-convex models and workers on CIFAR10.

Figure 3: Learning process.

Figure 4: Comparing GMC * with GMC and DMSGD.

Figure5shows the learning process of ResNet56 on CIFAR10. In the first 80 epochs, GMC and Algorithm 4 are almost the same and both of them are better than Algorithm 3. At the step of scaling step size by factor 0.1, the norm of memory gradient (y t,k ) in Algorithm 4 is 10 fold of that in GMC. Hence, GMC is better than Algorithm 4 after training 80 epochs.

Summing up the above equation fromt = 0 to T -1, we get that 1 ∇F (w t ) 2 ≤ F (w 0 ) -F (w * )

.T -1 do we have p workers. D k denotes the data stored on the kth worker (k = 1, 2, . . . , p) and p k=1 D k = {ξ 1 , ξ 2 , . . . , ξ n }. Our method Global Momentum Compression (GMC) mainly performs the following operations iteratively:

Table content: <table><tr><td>4:</td><td>Workers:</td><td/></tr><tr><td>5:</td><td colspan="3">for k = 1, 2 . . . , p, each worker parallelly do</td></tr><tr><td>6:</td><td colspan="2">if t &gt; 0 then</td></tr><tr><td>7:</td><td colspan="3">Receive w t -w t-1 from the server;</td></tr><tr><td>8:</td><td colspan="3">Get w t by w t = w t-1 + (w t -w t-1 );</td></tr><tr><td>9:</td><td>end if</td><td/></tr><tr><td>10:</td><td colspan="3">Randomly pick a mini-batch of training data I t,k ⊆ D k with |I t,k | = b;</td></tr><tr><td>11:</td><td>g t,k = 1 pb</td><td colspan="2">ξi∈I t,k ∇f (w t ; ξ i ) -β pηt (w t -w t-1 );</td></tr><tr><td>12:</td><td colspan="3">Generate a sparse vector m t,k ∈ {0, 1} d ;</td></tr><tr><td>13:</td><td colspan="3">Send m t,k (g t,k + u t,k ) to the server;</td></tr><tr><td>14:</td><td colspan="3">u t+1,k = (1 -m t,k ) (g t,k + u t,k );</td></tr><tr><td>15:</td><td>end for</td><td/></tr><tr><td>16: 17: 18:</td><td colspan="3">Server: w t+1 = w t -η t Send w t+1 -w t to workers; p k=1 m t,k (g t,k + u t,k );</td></tr><tr><td colspan="2">19: end for</td><td/></tr><tr><td colspan="4">3 GLOBAL MOMENTUM COMPRESSION</td></tr><tr><td/><td colspan="2">• Each worker calculates g t,k = 1 pb</td><td>ξi∈I t,k ∇f (w; ξ i )-β pηt (w t -w t-1 ), where I t,k ⊆ D k</td></tr><tr><td/><td colspan="2">and |I t,k | = b;</td></tr></table>

Experimental results on CIFAR10. test accuracy as that of DMSGD during the learning process. Table2shows that compared with DMSGD, GMC can reduce the communication cost by more than 100 fold with no loss of accuracy and save about 40% of training time. Compared with DGC, GMC achieves slightly better accuracy with comparable communication compress ratio.

Table content: <table><tr><td/><td/><td colspan="2">AlexNet</td><td colspan="2">ResNet20</td><td colspan="2">ResNet56</td></tr><tr><td colspan="3">GPU Method Accuracy</td><td>CR</td><td>Accuracy</td><td>CR</td><td>Accuracy</td><td>CR</td></tr><tr><td/><td>DMSGD</td><td>75.76%</td><td>100%</td><td>91.93%</td><td>100%</td><td>93.23%</td><td>100%</td></tr><tr><td>4</td><td>DGC</td><td>75.67%</td><td>0.48%</td><td>91.35%</td><td>0.37%</td><td>93.08%</td><td>0.47%</td></tr><tr><td/><td>GMC</td><td>76.10%</td><td>0.48%</td><td>91.97%</td><td>0.36%</td><td>93.23%</td><td>0.46%</td></tr><tr><td/><td>DMSGD</td><td>76.08%</td><td>100%</td><td>91.79%</td><td>100%</td><td>93.47%</td><td>100%</td></tr><tr><td>8</td><td>DGC</td><td>75.42%</td><td>0.87%</td><td>91.80%</td><td>0.66%</td><td>93.26%</td><td>0.84%</td></tr><tr><td/><td>GMC</td><td>76.19%</td><td>0.85%</td><td>91.98%</td><td>0.64%</td><td>93.43%</td><td>0.81%</td></tr><tr><td colspan="4">Method Time (second/epoch) Accuracy</td><td>CR</td><td/><td/><td/></tr><tr><td>DMSGD</td><td>2536</td><td/><td>69.79%</td><td>100%</td><td/><td/><td/></tr><tr><td>DGC</td><td>1496</td><td/><td>69.57%</td><td>0.84%</td><td/><td/><td/></tr><tr><td>GMC</td><td>1495</td><td/><td>69.88%</td><td>0.82%</td><td/><td/><td/></tr></table>

Experimental results on ImageNet.

Table content: <table/>

Generate a sparse vector m t,k ∈ {0, 1} d ; Send m t,k (g t,k + u t,k ) to the server;u t+1,k = (1 -m t,k ) (g t,k + u t,k ); (g t,k + u t,k );Send w t+1 -w t to workers; end for If η t is a constant, then Algorithm 4 is the same as GMC. If η t is non-increasing, then η t ≥ η t+1 and

Table content: <table><tr><td>Algorithm 4</td><td/><td/><td/><td/><td/></tr><tr><td colspan="2">for t = 0, 1, 2, ...T -1 do</td><td/><td/><td/><td/></tr><tr><td>Workers:</td><td/><td/><td/><td/><td/></tr><tr><td colspan="7">for k = 1, 2 . . . , p, each worker parallelly do</td></tr><tr><td>g t,k = ηt pb</td><td colspan="6">ξi∈I t,k ∇f (w t ; ξ i ) -β p (w t -w t-1 );</td></tr><tr><td>end for</td><td/><td/><td/><td/><td/></tr><tr><td colspan="2">Server: w t+1 = w t -k=1 m t,k y t+1,k ≥ (I -M t,k )( p</td><td>1 pb</td><td colspan="2">ξi∈I t,k</td><td colspan="2">∇f (w t ; ξ i ) -</td><td>β pη t</td><td>(w t -w t-1 ) + y t,k ) .</td></tr><tr><td/><td/><td colspan="2">1 pb</td><td colspan="2">ξi∈I t,k</td><td>∇f (w t ; ξ i ) -</td><td>β pη t</td><td>(w t -w t-1 ) + y t,k )</td></tr><tr><td/><td colspan="6">the same as the memory gradient of GMC</td></tr></table>

2 + 2U 2 . C.5 PROOF OF THEOREM 3 Since F (•) is L-smooth, we get that EF (x t+1 ) ≤EF (x t ) -η 1 -β E[∇F (x t ) T ∇f (w t ; I t )] + Lη 2 2(1 -β) 2 E ∇f (w t ; I t ) 2 ≤EF (x t ) -η 1 -β E[ ∇F (w t ) 2 + (∇F (x t ) -∇F (w t )) T ∇F (w t )] +

Table content: <table><tr><td/><td/><td>L 2</td><td>E</td><td>η 1 -β</td><td>∇f (w t ; I t ) 2</td></tr><tr><td>=EF (x t ) -</td><td>η 1 -β</td><td>E[∇F (x</td><td/></tr></table>

In the convergence theorems, we need E ũt 2 ≤ U 2 (Assumption 2). Since ũt = p k=1 u t,k , we only need E u t,k 2 to be bounded which is mainly related to the choice of m t,k . Theoretically, we set one large threshold θ > 0 in advance and we denote θ s t,k as the s-largest value of {|g . . . , d}. Then in each iteration, we can choose the values such that |g | ≥ min{ θ, θ s t,k }. Then we obtain u t,k 2 ≤ d θ2 and hence E u t,k 2 ≤ d θ2 . In practice, the large threshold θ is never been used for the choice of m t,k since θ s

Table content: <table><tr><td>+ LG 2 2(1-β) 2 .</td></tr><tr><td>D MORE DISCUSSION ABOUT ASSUMPTION 2</td></tr><tr><td>(j) t,k +u (j) t,k ||j =</td></tr><tr><td>1, (j) t,k + u (j) t,k t,k</td></tr></table>

(g t,k + u t,k ) 2 . According to the definition of M t,k , we have ∀γ > 0, -w t-1 )/η t + ∇f (w t ; I t ) + ũt

Table content: <table><tr><td>ηt-ηt+1 1-β ũt+1</td><td>2 .</td><td/><td/></tr><tr><td colspan="2">First, we get that w t+1 -w t k=1 M t,k w t+1 -w t 2 ≤ η 2 t p 2 ≤η 2 t gt + ũt 2</td><td/><td/></tr><tr><td/><td colspan="4">=η 2 t β(w t 2</td></tr><tr><td/><td>≤(1 + γ)β 2 w t -w t-1</td><td>2 + (1 +</td><td>1 γ</td><td>)(2G 2 + 2U 2 )η 2 t .</td></tr></table>

Peer review (iclr2020):
This paper proposes a scheme for incorporating compressed (sparsified) gradients with momentum in distributed SGD. The approach differs from others in the literature, comes with theoretical guarantees, and improved performance. The results are correct, and the experiments illustrate that the proposed approach can make a difference (albeit, modest) in the quality of the resulting model. The main point I find dissatisfying about the theoretical results of the paper are the use of Assumption 2. The memory vector is a parameter of the algorithm. I realize that one can enforce this with a projection, as argued in the paragraph rationalizing this assumption. However, that specific case isn*t analyzed and it isn*t clear how incorporating that projection would affect the accuracy, since it would essentially be countering the effect of error feedback. I also find Assumption 3 to be strange. In the convex setting, one can typically show that this follows from Assumption 1 alone under the additional assumption of a suitably small step size. In the non-convex setting it isn*t clear what this means, since w^* is not well defined (if there are multiple global minimizers). Assumption 1 is also strong. Typically one assumes that the stochastic gradients are unbiased, and either that the expected gradient is Lipschitz continuous (in the smooth case), or the expected gradient is bounded (in the non-smooth case). Assuming that the stochastic gradients are uniformly bounded essentially implies that the noise vanishes when the gradient gets large. Can you provide examples of functions/problems satisfying these assumptions? Even an example as simple as the case where is a finite sum of quadratic functions, and one randomly samples one of the terms in the finite sum to compute the gradient (i.e., using SGD to solve a large linear least squares problem) doesn*t appear to satisfy Assumption 1. Overall the results are potentially interesting. I would have given a higher rating if the assumptions didn*t appear to be so strong, and if the experimental results demonstrated a more substantial difference with DGC.