A Convergence Analysis of Distributed SGD with Communication-Efficient Gradient Sparsification

1 Introduction

Stochastic gradient descent (SGD) algorithms are commonly used for training many machine-learning models. SGD minimizes the objective function f : Rd → R with stochastic gradients G(xt) using the following update formula:

where xt ∈ Rd is a set of model parameters, and αt ∈ R is the step size at iteration t. With large-scale models (i.e., d is at the order of millions or even billions) and data sets, distributed synchronous SGD (S-SGD) with data-parallelism is the key technique to reduce the overall training time using multiple computational workers [Goyal et al., 2017; Jia et al., 2018]. Given a cluster with P workers, in the tth iteration, the pth worker calculates the gradients Gp (xt) with locally sampled data, and then all workers collaboratively update the model parameters with the aggregated gradients 1P ∑Pp=1 Gp (xt), i.e.,

Ideally, S-SGD with P workers would accelerate the training process by P times. However, the aggregation of gradients requires tremendous data communications among workers, whose time cost becomes significant, especially when the network bandwidth is relatively low [Dean et al., 2012; Shi et al., 2018]. Efficient communication methods have been proposed to alleviate the communication overheads on the system level [Awan et al., 2017; Zhang et al., 2017; Shi et al., 2019a; Chen et al., 2019], while the Top-k sparsification scheme [Chen et al., 2018; Lin et al., 2018] has been proposed to sparsify the gradients to dramatically reduce the communication cost with little impact on the model accuracy on the algorithm level. In Top-k S-SGD, each worker only selects its top-k gradients (in terms of absolute magnitude) to be exchanged with other workers. The update formula becomes:

where ˜Gp (xt) = TopK(Gp (xt)) is the sparsified top-k gradients at the pth worker. Specifically, for a vector x ∈ Rd , TopK(x) ∈ Rd , and the ith (i = 1, 2, ..., d) element of TopK(x) is defined by:

where x(i) denotes the ith element of x and thr is the kth largest value of |x|. In practice, k can be two to three orders of magnitude smaller than d with little impact on the model accuracy [Aji and Heafield, 2017; Lin et al., 2018; Alistarh et al., 2018], which can dramatically reduce the communication overhead. Some work [Wangni et al., 2018; Stich et al., 2018; Alistarh et al., 2018; Jiang and Agrawal, 2018] has provided the theoretical convergence analysis of Top-k S-SGD under different assumptions.

However, it is noted that the indices of non-zero elements of ˜Gpi (xt) for different workers are generally inconsistent.

Therefore, the aggregation of sparsified gradients of Top-k S-SGD from P workers has a communication complexity of O(kP) [Renggli et al., 2018] 1 . Recently, an alternative gradient sparsification scheme named gTop-k [Shi et al., 2019b] has been proposed to reduce the communication complexity from O(kP ) to O(k log P ) by using a tree based approximate reduction algorithm. Their empirical studies show that gTopk S-SGD achieves much better training performance than the Top-k scheme on GPU clusters; yet, there is no theoretical justification on the convergence of gTop-k S-SGD.

In this paper, we provide a detailed theoretical analysis on the convergence performance of gTop-k S-SGD on nonconvex problems. We summarize our main contributions as follows:

• Unlike the existing convergence analysis in [Stich et al., 2018; Alistarh et al., 2018] that uses some assumptions on ∑Pp=1 ˜Gp (xt) which may not hold in practice, we introduce a relatively weak assumption which can be easily verified through real-world experiments.
• We prove that gTop-k S-SGD provides convergence guarantees for non-convex problems under our analytic assumptions. We conduct extensive experiments on representative deep learning models and data sets to verify the soundness of the assumptions and theoretical results.
• We show that gTop-k S-SGD has the same theoretical convergence rate with vanilla mini-batch SGD with properly chosen learning rates. We also discuss the impact of the compression ratio on the convergence rate through experiments.

2 RelatedWork

There are two main types of communication reduction schemes in S-SGD: gradient quantization and sparsification.

In quantization methods, the exchanged gradients at every iteration can be quantified to a small number of bits (e.g., 2 bits) with error compensation [Alistarh et al., 2017; Wen et al., 2017; Bernstein et al., 2018; Jiang and Agrawal, 2018; Wu et al., 2018; Stich et al., 2018; Wangni et al., 2018; Haddadpour et al., 2019] during communication while keeping the model accuracy nearly unchanged. However, even using only one bit for each gradient, the maximum communication compression ratio is 32× compared to the 32-bit counterpart.

In sparsification methods, one can only transmit a small portion of non-zero gradients [Aji and Heafield, 2017; Chen et al., 2018; Lin et al., 2018; Stich et al., 2018; Jiang and Agrawal, 2018; Alistarh et al., 2018; Wangni et al., 2018; Wang et al., 2018] for aggregation so that the communication size can be reduced significantly. Researchers [Aji and Heafield, 2017; Chen et al., 2018; Lin et al., 2018; Renggli et al., 2018] first empirically show the effectiveness of the Top-k or the random-k sparsification in S-SGD with little impact on the model convergence, where k can be only

1 For all-reduce based aggregation, every worker has a communication complexity of O(kP). For parameter-server based aggregation, the parameter server has a communication complexity of O(kP ).

0.1% of the gradient dimension d. Some recent work [Stich et al., 2018; Jiang and Agrawal, 2018; Alistarh et al., 2018; Wangni et al., 2018] provides the convergence analysis on Top-k S-SGD under different assumptions. However, even though the Top-k sparsification scheme can zero-out a large number of gradients, it generates the irregular indices among different workers such that the communication complexity is O(kP) for P workers [Renggli et al., 2018], which limits the system scalability. To further reduce the communication complexity, a communication-efficient sparsification scheme named gTop-k has been recently proposed [Shi et al., 2019b]. gTop-k has a communication complexity of O(k logP) by leveraging a tree structure for gradient communications, and therefore it performs much better than Top-k on large clusters. However, there is no theoretical justification on the convergence of gTop-k S-SGD in [Shi et al., 2019b]. Due to the biased gradients aggregation through the gTop-k sparsification, the theoretical convergence analysis is non-trivial. In this study, we provide the convergence proofs for gTop-k S-SGD, and we conclude that gTop-k S-SGD has the same convergence rate as vanilla S-SGD.

We want to highlight that the existing convergence analysis on Top-k S-SGD [Wangni et al., 2018; Stich et al., 2018; Alistarh et al., 2018; Jiang and Agrawal, 2018] cannot be directly applied to prove the convergence of gTop-k S-SGD. First, gTop-k S-SGD is a biased stochastic compression scheme which is different with the unbiased one in [Wangni et al., 2018]. Second, the analysis in [Stich et al., 2018] is for convex problems, and it requires the sparsification on fully aggregated gradients, while gTop-k S-SGD has no such condition. Third, the analysis in [Jiang and Agrawal, 2018] requires the algorithm to exchange all parameter components in any certain T consecutive iterations, which could also not hold on gTop-k S-SGD since in every iteration only top-k gradients are selected and some very small gradients may not be chosen throughout the training process. Our analysis is closer to the work [Alistarh et al., 2018], but there are three main technical differences. 1) We use a relatively weak analytic assumption on the top-k gradients (and also gTop-k gradients). 2) We eliminate the condition (k > d/2) that is required in [Alistarh et al., 2018] to guarantee the convergence. 3) We prove the convergence of the gTop-k S-SGD algorithm, and derive the convergence rate, and empirically evaluate the impact of compression ratio on the convergence performance.

3 The Algorithm of gTop-k S-SGD

For completeness, in this section we briefly introduce the algorithm of communication-efficient global Top-k (gTop-k) SSGD proposed in [Shi et al., 2019b]. Before describing the algorithm, we define some notations. Let vt and ϵpt denote the local model of each worker and the local gradient residuals of worker p at iteration t, respectively. Note that all workers have the consistent model at any iteration. In gTop-k S-SGD, the model is updated by

where gTopKPp=1 (xp ) = x1⊤x2⊤...⊤xP , and the operator ⊤ is defined as follows. For any two vectors xi ∈ Rd and

where mask = |TopK(xi ) + TopK(xj )| > thr and thr is the kth largest value of |TopK(xi )+TopK(xj )|. Assume that x is the aggregation result by gTopKPp=1 (xp ), it simultaneously generates a vector of gMaskp ∈ Rd which indicates the indices of the selected local values (i.e., TopK(xp )) that contribute to the final x. Specifically, the ith (i = 1, 2, ..., d) element of gMaskp is defined as

The pseudocode of gTop-k S-SGD is shown in Algorithm 1.

Algorithm 1 gTop-k S-SGD at worker p

Input: Stochastic gradients Gp (·) at worker p

Input: Configured value k and the learning rate α

• 1: Initialize v0 = ϵp0 = 0;
• 2: for t = 1 → T do
• 3: accpt = ϵpt-1 +αGpt (vt-1); // Accumulate the residuals
• 4: gt, gMaskp t = gTopKPp=1 (accpt ); // Global top-k and mask
• 5: ϵpt = accpt ⊙¬gMaskp t ; // Store residuals
• 6: vt = vt-1 - 1P gt; // Update the model
• 7: end for

Similar to [Alistarh et al., 2018], we also use xt to denote the auxiliary random variable at iteration t, and

where Gt(vt) = 1P ∑Pp=1 Gpt (vt) and x0 = 0d . The difference between the auxiliary variable xt and the model variable vt can be represented byϵt = vt - xt. (9)

According to Algorithm 1, we have ϵt = 1P ∑Pp=1 ϵpt .

4 Convergence Analysis

4.1 Notations and Assumptions

We mainly discuss the cases that all the computational workers have a full copy of data. We assume that the gTop-k SSGD is applied to solve the non-convex objective function f : Rd →R, which is L-Lipschitz smooth, i.e.,

The sampled stochastic gradients G(·) at every iteration are unbiased, i.e., E[G(vt)] = ∇f(vt). We also assume that the second moment of the stochastic gradients is bounded, i.e.,

where Gp,(i) t (x) are the gradients of the ith sample in a minibatch and || · || is ℓ2-norm. Let b denote the mini-batch size used per worker, and the total mini-batch size with P workers is B = Pb. The mini-batch setting has Gpt (x) =

1b ∑bi=1 Gp,(i) t (x). Thus, the second moment of the average gradients has a smaller bound, i.e., for any t ∈ N,

Assumption 1. The gTopK operator is expected to select k larger values than randomly selecting k values from the accumulated vectors, i.e.,

where randomK(xp ) ∈ Rd is a vector whose k elements are randomly selected from xp following a uniform distribution, and the other d - k elements are zeros.

The assumption will be verified by experiments in Section 5. The key ideas of the proofs are 1) We first bound the difference between the model xt without sparsification and the sparsified model vt. It enables us to bound the expected sumof-squares of gradients of f so that the convergence is guaranteed [Bottou et al., 2018]. 2) Then we bound the expected average-squared gradients of f with some sufficient conditions to derive the convergence rate.

4.2 Main Results

Lemma 1. For any vectors xp ∈ Rd , p = 1, 2, ..., P, and 0 < k ≤ d, it holds that

Proof. In [Stich et al., 2018], the authors have shown that for any vector x ∈ Rd , it holds

Combined with Assumption 1, we easily obtain

Lemma 2. For any iteration t ≥ 1:

where γ = 1 - kd , 0 < k ≤ d and η > 0.

Proof. We derive the difference between vt+1 and xt+1, i.e.,

Proceedings of the Twenty-Eighth International Joint Conference on Artificial Intelligence (IJCAI-19)

Iterating the above inequality from i = 0 → t yields:

Corollary 1.

Proof. Using Lemma 2 and the bound of (12), we have

Theorem 1. Assume that gTop-k S-SGD is applied to minimize the objective function f that satisfies the assumptions in Section 4.1. If one chooses a learning rate schedule such that for any iteration t > 0:

for some constant D > 0, then after running T iterations with Algorithm 1, we have

where x∗ is the optimal solution to the objective function f.

Proof. Under the Assumption of L-smooth of f, we have

Taking the expectation at iteration t, we have

Taking the expectation before t, it yields

Apply (18) to the above inequality, we have

Then we can obtain

Using the L-smooth property of f , we have

Combine with (21), we obtain

Summing up the above inequality for t = 1, 2, ..., T , we have

By dividing the summation of learning rates, we have:

The condition (18) holds if γ(1 + η) < 1 for both fixed learning rates and diminishing learning rates. To derive the bound of η, we have

Therefore, one should choose η < k d-k to satisfy the above inequality. Theorem 1 implies that Algorithm 1 converges to 0 if T is large enough, when αt is set to satisfy the following conditions:

Corollary 2. Under the assumptions in Theorem 1, if τ = γ(1 + η) and αt = θ√B/T, ∀t > 0, where θ > 0 is a constant, we have the convergence rate for Algorithm 1:

Proof. First we prove that αt = θ√B/T, which is a constant step size (or learning rate), satisfies the condition in (18). We set αt = α for simplification (i.e., α = θ√B/T ). We have

Since 0 ≤ τ < 1, we obtain

Therefore, (18) holds by choosing D = ατ 1-τ . Applying Theorem 1, we obtain the inequality of the expected averagesquared gradients of f, i.e.,

which concludes the proof.

From Corollary 2, we can seen that with a properly set learning rate, the gTop-k S-SGD algorithm has a convergence rate of O( 1 √BT ), which is the same as that of mini-batch SGD [Dekel et al., 2012]. It also indicates that k has small impact on the convergence rate if T is large enough.

4.3 Discussion

In Corollary 2, there are two terms to determine the convergence rate of gTop-k S-SGD. The first term indicates that the convergence rate is affected by the constant θ and the minibatch size, and the second term indicates that the convergence rate is also affected by both θ (related to the learning rate) and τ (related to the compression ratio d k ). The second term will be dominated by the first term if T is large enough. However, it is not uncommon that a fixed number of iterations is used for training deep neural networks (DNNs) in practice. As a result, although a larger compression ratio leads to less communications overhead, it would enlarge the bound of the convergence rate.

To understand the details, we expand the second term on the right-hand side of (25). Let c = d/k denote the compression ratio, then γ = 1-1/c and τ = (1-1/c)(1+η). Since η should satisfy the condition of η < k/(d - k), we choose η = k/d = 1/c. Thus,

Therefore, inequality (25) becomes

The above inequality indicates that given a fixed iteration budget (i.e., T ), a higher compression ratio (c) causes a larger bound of the convergence rate. In summary, to achieve a better convergence with a given time budget, one should balance the communication cost and the convergence rate. We will further evaluate the impact of the compression ratio on the convergence performance through experiments in Section 5.

5 Experiments

5.1 Experimental Settings

Our experimental settings cover three deep learning applications. 1) Image classification: Two popular DNNs, VGG-16 [Simonyan and Zisserman, 2014] and ResNet-20 [He et al., 2016], are used for evaluation on the data set of Cifar-10 2 which consists of 50000 training images. 2) Language model: A 2-layer LSTM model (LSTM-PTB) with 1500 hidden units per layer is adopted for evaluation on the data set of PTB [Marcus et al., 1993], which contains 923000 training words. 3) Speech recognition: A 5-layer LSTMmodel (LSTM-AN4) with 800 hidden units per layer is used for evaluation on AN4 [Acero, 1990], which contains 948 training utterances. In all training models, we exploit the warmup strategy in gTop-k S-SGD on the 4-worker distributed environment. The baselines are evaluated using S-SGD without gradient sparsification. The main hyper-parameters adopted in evaluation are shown in Table 1.

5.2 Verification of Assumption and Convergences

We verify Assumption 1 empirically by training DNNs with gTop-k S-SGD. During the training process, we measure

where xp = Gpt (vt) + ϵpt and k = 0.001 × d (i.e., c = 1000). Assumption 1 holds if δ ≤ 1. The measurements of δ corresponded with the training losses on the evaluated DNNs are shown in Fig. 1. It can be seen that we always have δ < 1, which verifies the soundness of Assumption 1. In Fig. 1, the convergences of gTop-k S-SGD are nearly consistent with SSGD, which validates our theoretical results and shows that gTop-k S-SGD can converge as fast as S-SGD.

5.3 Convergence Rate v.s. Compression Ratio

The second term in inequality (27) indicates that the convergence rate may be degraded by the compression ratio c. We evaluate the sensitivity of the convergence rates to c on training DNNs without changing hyper-parameters including the total number of iterations (i.e., a fixed number of epochs).

2 https://www.cs.toronto.edu/∼kriz/cifar.html

The results are shown in Fig. 2, which shows that with larger c, the convergence of the models would slowdown. Therefore, with large compression ratios, there is a trade-off between the communication size, which is directly related to the iteration time, and the convergence rate.

6 Conclusion

Top-k gradient sparsification is crucial for reducing the communication size in distributed S-SGD. The gTop-k scheme is a more communication efficient scheme than Top-k for gradient sparsification. In this study, we present a detailed convergence analysis for gTop-k S-SGD under some analytical assumptions, and we derive its convergence rate. Our theoretical results conclude that gTop-k S-SGD provides convergence guarantees for non-convex objective functions and it has the same convergence rate with vanilla mini-batch SGD with properly chosen learning rates. We derive and evaluate the impact of compression ratios on the convergence performance. We finally conduct experiments to verify the soundness of the analytical assumption and theoretical results.

Acknowledgements

The research was supported by Hong Kong RGC GRF grant HKBU 12200418.