On the Convergence Properties of a K-step Averaging Stochastic Gradient Descent Algorithm for Nonconvex Optimization

1 Introduction

Parallel and distributed processing have been adopted for stochastic optimization to solve large-scale machine learning problems. Efficient parallelization is critical to accelerating long running deep-learning applications. Derived from stochastic gradient descent (SGD), parallel solvers such as synchronous SGD (e.g., [Zinkevich et al., 2010] [Dekel et al., 2012]) and Asynchronous SGD (ASGD) (e.g., see [Dean et al., 2012; Recht et al., 2011]), have been proposed.

Beginning with the seminal paper [Robbins and Monro, 1951], the convergence properties of SGD and its variants have been extensively studied for the past 50 years (e.g. see [Robbins and Siegmund, 1971; Bottou, 1998; Nemirovski et al., 2009; Shamir and Zhang, 2013; Ghadimi and Lan, 2013]). The asymptotic optimal convergence rate of SGD was proved by [Chung, 1954] and [Sacks, 1958] to be O(1/N) with twice continuously differentiable and strongly convex objectives. N is the number of samples processed. The iteration complexity is O(1/√N) for general convex (see [Nemirovski et al., 2009]) and nonconvex (see [Ghadimi and Lan, 2013]) problems. Regarding parallel variants of SGD, [Dekel et al., 2012] extend these results to the setting of synchronous SGD with P learners and show that it has convergence rate of O(1/√NP ) for non-convex objectives, with N being the number of samples processed by each learner. Hogwild! is a lockfree implementation of ASGD, and [Recht et al., 2011] prove its convergence for strongly convex problems with theoretical linear speedup over SGD. Downpour is another ASGD implementation with resilience against machine failures [Dean et al., 2012]. [Lian et al., 2015] show that as long as the gradient staleness is bounded by the number of learners, ASGD converges for nonconvex problems. Due to its asynchronous nature that reduces communication cost, ASGD receives much attention in many recent studies.

∗ Supported in part by NSF Grants DMS-1509739 and CCF1523768

Although ASGD has the same asymptotic convergence rate as SGD when the staleness of gradient update is bounded, the learning rate assumed for proving ASGD convergence are usually too small for practical purposes. It is also difficult for an ASGD implementation to control the staleness in gradient updates as it is influenced by the relative processing speed of learners and their positions in the communication network. Furthermore, the parameter server presents performance challenges on platforms with many GPUs. On such platforms, a single parameter server oftentimes does not serve the aggregation requests fast enough. A sharded server alleviates the aggregation bottleneck but introduces inconsistencies for parameters distributed on multiple shards. Communication between the parameter server (typically on CPUs) and the learners (on GPUs) is likely to remain a bottleneck in future systems.

We adopt a distributed, bulk-synchronous SGD algorithm that allows delayed gradient aggregation to effectively minimize the communication overhead. We call this algorithm K-step average SGD (K-AVG). Instead of using a parameter server, the learners in K-AVG compute the average of their copies of parameters at regular intervals through global reduction. Rather than relying on asynchrony that reduces communication overhead but has adverse impact on practical convergence, the communication interval K is a parameter in K-AVG. The communication time is amortized among the data samples processed within each interval. On current and emerging computer platforms that support high bandwidth direct communication among GPUs (e.g., GPU-direct), global reduction does not involve CPUs and avoids multiple costly copies through the software layers. Similar averaging approaches have been proposed in the literature, see [Hazan and Kale, 2014; Johnson and Zhang, 2013; Smith et al., 2016; Zhang et al., 2016; Loshchilov and Hutter, 2016; Chen et al., 2016; Wang et al., 2017]. However, their convergence behavior is not well understood analytically for nonconvex objectives, and it is unclear how they compare with ASGD approaches. This is the part where our major contribution goes to.

We study the convergence behavior of K-AVG and the impact of the number of processors P on convergence. As the abundance of data is critical to the success of most machine learning tasks, training employs increasingly more learners. We show that K-AVG scales better than ASGD with P and K-AVG allows larger stepsizes than ASGD for the same P. We also analyze the impact of K on convergence. Since with K = 1 K-AVG becomes hard-sync SGD, and with larger K K-AVG can stimulate averaging after an epoch or many epochs, our analysis can be applied to many existing variants of synchronous SGD. Finding the optimal length of delay Kopt for convergence is of high importance to practitioners. Contrary to popular belief, Kopt is oftentimes not 1 and can be very large for many applications. Thus K-AVG is a good fit for large-scale distributed training as communication may not need to be very frequent for optimal convergence. Our analysis of convergence of K-AVG provides guidelines practitioners to explicitly balance the decrease of communication time and the increase of iterations through an appropriately chosen K.

Using an image recognition benchmark, we demonstrate the nice convergence properties of K-AVG in comparison to two popular ASGD implementations: Downpour [Dean et al., 2012] and EAMSGD [Zhang et al., 2015]. In EAMSGD, global gradient aggregation among learners simulates an elastic force that links the parameters they compute with a center variable stored by the parameter server. In both Downpour and EAMSGD, updates to the central parameter server can also have a K-step delay. On our target platform, when K is small, K-AVG significantly reduces the communication time in comparison to Downpour and EAMSGD while achieving similar training and test accuracies. The training time reduction is up to 50%. When K is large, K-AVG achieves much better training and test accuracies than Downpour and EAMSGD after the same amount of data samples are processed. For example, with 128 GPUs, K-AVG is up to about 7 and 2-6 times faster than Downpour and EAMSGD respectively, and achieves significantly better accuracy.

This rest of the paper is organized as follows: In section 2, we introduce the standard assumptions in optimization theory needed to analyze SGD methods and frequently used notations throughout the paper; In section 3, we formally introduce the K-AVG algorithm, and prove its standard convergence results with fixed and diminishing stepsize. Based on the convergence result, we analyze the scalability of KAVG and investigate the optimal choice of K; In section 4, we present our experimental results to validate our analysis.

2 Preliminaries and Notations

In this section, we introduce some standard assumptions used in the analysis of convex and non-convex optimization algorithms and key notations frequently used throughout this paper. We use ‖ · ‖2 to denote the ℓ2 norm of a vector in Rd ; 〈·〉 to denote the general inner product in Rd . For the key parameters we use:

• P denotes the number of processors;
• K denotes the length of the delay;
• Bn, ¯B, or B denotes the size of mini-batch for n-th update;
• γn, ¯ γ, or γ denotes the step size for n-th update;
• ξ jk,s denotes the i.i.d. realizations of a random variable ξ generated by the algorithm on different processors and in different iterations, especially, j = 1, ...,N, k = 1, ..., K, and s = 1, ..., B.

We study the following optimization problem:

where objective function F : Rm → R is continuously differentiable but not necessarily convex over X, and X ⊂ Rm is a nonempty open subset. Since our analysis is in a very general setting, F can be understood as both the expected risk F(w) = Ef(w; ξ) or the empirical risk F(w) = n-1 ∑ni=1 fi(w). As our approach for analysis is built upon smooth objectives, we introduce the following assumptions which are standard and fundamental.

Assumption 1. The objective function F : Rd → R is continuously differentiable and the gradient function of F is Lipschitz continuous with Lipschitz constant L > 0, i.e.

for all w, ˜w ∈ Rd .

Assumption 2. The sequence of iterates {wj} is contained in an open set over which F is bounded below by a scalar F ∗ .

Assumption 3. For any fixed parameter w, the stochastic gradient ∇F(w; ξ) is an unbiased estimator of the true gradient corresponding to the parameter w, namely,

Assumption 4. There exist scalars M ≥ 0 such that,

All results in this paper are based on the four assumptions.

3 Main Results

We present the distributed K-AVG algorithm in Alg. 1. Note that traditional parallel SGD algorithm is equivalent to KAVG with K = 1, as synchronization is required after each local update. However, K-AVG relaxes this requirement and allows for K individual updates before synchronization. Thus, K-AVG is a more general synchronous algorithm that contains parallel SGD. Surprisingly, as we show both analytically (section 3.4) and experimentally (section 4.2), more frequent synchronization does not always result in faster convergence for nonconvex objectives.

Algorithm 1: K-AVG Algorithm

3.1 Convergence of K-AVG

In the following theorem, we prove an upper bound on the expected average squared gradient norms, which serve as a metric to measure the convergence rate for nonconvex objectives.

Theorem3.1. (Nonconvex objective, fixed stepsize, and fixed batch size) Suppose that Algorithm 1 is run with a fixed stepsize γn = ¯ γ, a fixed batch size Bn = ¯B satisfying

with some constant 0 < δ < 1. Then the expected average squared gradient norms of F satisfy the following bounds for all N ∈ N:

Proof. We denote ˜wα as the α-th global update in K-AVG, and denote wjα+t as t-th local update on processor j. The (α + 1)-th global average can be written as

according to our algorithm, the random variables ξ j t,s are i.i.d. for all t = 0, ..., K - 1, s = 1, ..., B, and j = 1, ..., P .

Under the assumption that a constant stepsize is implemented within each inner parallel step, i.e. γt = γ, for t = 0, ...,K1, one can immediately simplify the above inequality as

The goal here is to investigate the expectation of F(˜wα+1) - F(˜wα) over all random variables ξ j t,s . Under the unbiased estimation Assumption 3, by taking the overall expectation we can immediately get

for fixed j and t. Here we abuse the expectation notation E a little bit. Throughout this proof, E always means taking the overall expectation. We will frequently use the above iterative conditional expectation trick in the following analysis. As a result, we can drop the summation over s due to an averaging factor B in the dominator of (3.6). Next, we show how to get rid of the summation over j. Recall that wjα+1 = ˜wα - γB ∑Bs=1 ∇F(˜wα; ξ j0,s ). Obviously, wjα+1 , j = 1, ..., P are i.i.d. for fixed ˜wα because ξ j0,s , j = 1, ..., P , s = 1, ..., B are i.i.d. Similarly, wjα+2 = wjα+1 - γB ∑Bs=1 ∇F (wjα+1 ; ξ j1,s ), j = 1, ..., P are i.i.d. due to the fact that wjα+1 's are i.i.d., ξ j1,s 's are i.i.d., and wjα+1 's are independent from ξ j1,s 's. By induction, one can easily show that for each fixed t, wjα+t , j = 1, ..., P are i.i.d. Thus for each fixed t

We can therefore get rid of the summation over j as well in (3.6). By taking the overall expectation on both sides of (3.6) and (3.7), we have

We are going to bound (3.10) and (3.11) respectively. Obviously (we treat ˜wα as a constant vector at this moment),

where we used the fact that ˜wjα = ˜wα, for j = 1, ..., P for the last term and the assumption that gradient ∇F is Lipschitz continuous. Note that

where the first inequality is due to Cauchy-Schwartz inequality, the last equity is due to Assumption 3 and the last inequality is due to Assumption 4. We plug the above inequality back into (3.12) and get

On the other hand, to bound (3.11), we can apply the similar analysis,

Combine the results in (3.14) and (3.15), we have

Under the condition 1 ≥ L2 γ 2 (K+1)(K-2) 2 +LγK, the second term on the right hand side can be discarded. This condition also implies that

Then, together with the condition 1 - δ ≥ L2γ 2 for some 0 < δ < 1, we have

If we allow both batch size and step size to change after each averaging step, by taking the summation we have

Under Assumption 2, we have

Combining both, we can immediately get the following bound

If we employ a constant step size and batch size, we get the bound on the expected average squared gradient norms of F as following:

which completes the proof.

An immediate observation from (3.2) is that the expected average squared gradient norms of F converges to some nonzero constant as N →∞. The first term 2(F(˜w1)-F ∗ ) N(K-1+δ)¯ γ reflects the distance from initial weight to the solution. It eventually goes to zero as the number of iterations goes to infinity. The second term LK¯ γM ¯B(K-1+δ) ( KP + L(2K-1)(K-1)¯ γ 6 ) is not affected by the number of iterations. Compared with sequential SGD (see section 4.3 in [Bottou et al., 2016]), this term is scaled by 1/ ¯B, and 1/P or L(2K-1)(K-1)¯ γ 6 , which means larger batch size, and smaller stepsize, more learners or more frequent averaging reduces this term. Mini-batch method and averaging as variance reduction techniques contribute to the scaling of ¯B and 1/P, respectively. When P is large enough that L(2K-1)(K-1)¯ γ 6 dominates K/P, the effect of averaging diminishes.

The rate of convergence (also referred as iteration complexity) after N step updates is established in the following corollary from Theorem 3.1.

Corollary 3.1. Under the condition of Theorem 3.1, take

Then for any

the following bound is achieved

Proof. At first, we assume that K/P > L(2K - 1)(K - 1)¯ γ/6. Then we can rewrite the bound (3.2) as

Set

By taking f ′ = 0, one immediately get

and

By plugging in the value of ¯ γ1, K/P > L(2K - 1)(K - 1)¯ γ/6 implies that

At last, we need to check condition (3.1) is hold. It is sufficient to set

This completes the proof of Corollary 3.1.

Condition (3.23) and bound (3.25) imply when the number of updates N is large enough, K-AVG eventually achieves a similar rate of convergence as classical SGD method for nonconvex objectives. Indeed, the rate of convergence of classical SGD methods is N -1/2 after N samples processed. Note that N updates in K-AVG means that N ∗ K ∗ B ∗ P samples have been processed. Taking a closer look at bound in (3.25), the right hand side is of the order O((N ∗ B ∗ P)-1/2 ). KAVG loses a factor of 1/√K as a result of communication saving. However, this doesn't mean that the smaller K the better. Since there is an extra multiplicative factor 4K K-1+δ which is monotone decreasing with respect to K. We will have a more detailed discussion on the impact of K in section 3.4.

Deploying diminishing stepsizes and/or dynamic batch sizes makes the expected average squared gradient norms converge to zero for non-convex optimization. In the following theorem, we establish the convergence result under such conditions.

Theorem 3.2. (Nonconvex objective, diminishing step size, and growing batch size) Suppose that Algorithm 1 is run with diminishing step size γj , and growing batch size Bj satisfying

with some constant 0 < δ < 1. Then the weighted average squared gradient norms satisfies

Especially, if

Then

Proof. If we use a diminishing step size γj and growing batch size Bj, Thus, from (3.21), we divide both sides with ∑Nj=1 γj , we have

If the following restriction of step size is satisfied,

it immediately implies the result of the theorem when N → ∞.

As we can see, by adopting a diminishing sequence of stepsizes instead of a fixed one, the expected average squared gradient norms of K-AVG converges to 0 instead of a nonzero constant.

3.2 K-AVGAllows Larger Stepsize than ASGD

Compared with the classical stepsize schedule for both sequential SGD (proposed by [Robbins and Monro, 1951]) and ASGD:

the stepsize schedule proposed in (3.34) allows choices of larger γj. On one hand, ∑∞j=1 γ 3 j < ∞ itself is a more relaxed constraint than ∑∞j=1 γ 2 j < ∞. On the other hand, when P is large, ∑∞j=1 γ 2 j /BjP also allows larger γj. Intuitively, averaging brings variance reduction and leads to relaxation of stepsize constraints. In our experiments (section 4.1), larger stepsizes work well in K-AVG but result in divergence in ASGD implementations.

3.3 K-AVG Scales Better than ASGD

We analyze the bound on expected average squared gradient norms in (3.2) to show K-AVG algorithm scales better with P than ASGD. We first establish the following theorem on the scalability of K-AVG.

Theorem 3.3. K-AVG scales better than ASGD.

Proof. In ASGD, one key parameter is the maximum staleness, generally assumed to be upper bounded by the number of processors, i.e. P in K-AVG. With fixed stepsize, the expected average gradient norms is (see also [Lian et al., 2015], Theorem 3) is

where C0 and C1 are constants independent of P . Compared with the bound in (3.2),

P serves as a scaling factor in the denominator, which implies K-AVG scales better than ASGD.

The comparison of practical scalability between K-AVG and ASGD implementations is shown in section 4.1.

3.4 Optimal K for Convergence Is Not Always 1

Unlike convex optimization problems, different learners may converge to different local optimums in nonconvex case. As a consequence, the frequency of averaging for nonconvex problems may be different from that of convex cases. [Zhang et al., 2016] expressed the same concern, their experimental results showed that periodic averaging tends to improve the solution quality. Contrary to popular belief that more frequent averaging i.e. smaller K speeds up convergence, we show that the optimal frequency K for convergence is not always 1. We consider the case that the amount of samples processed N ∗ K is constant, which means that the computational time remains as a constant given a fixed number of processors. If every other parameter stays the same, larger K means longer delay and fewer updates of global parameter ˜wn. The following theorem discusses the impact and optimal choice of K in K-AVG under such an assumption.

Theorem 3.4. Let S = N ∗ K be a constant. Suppose that Algorithm 1 is run with a fixed stepsize γn = ¯ γ, a fixed batch size Bn = ¯B satisfying

with some constant 0 < δ < 1. If

Then the optimal choice of K is greater than 1.

Proof. Under the assumption S = N ∗K, we can rewrite the bound (3.2) as

To move on, we set

where

To minimize the right hand side of (3.43), it is equivalent to solve the following integer programing problem

. Integer programming is hard. In addition Kopt depends on quantities such as L, M and (F(˜w1) - F ∗ ) that are oftentimes unknown. Instead of solving for the optimal K, we evaluate the monotonicity of B(K). It turns out that (α + βK + η(2K - 1)(K - 1)) is monotonically increasing for all K ≥ 1, and ( K K-1+δ ) is monotonically decreasing for K ≥ 1. Thus, there exists a unique Kopt. Then a sufficient condition for Kopt > 1 is that B(2) < B(1), which implies

Theorem 3.4 indicates that more frequent averaging is not necessarily helpful for nonconvex optimization, The sufficient condition (3.42) implies that larger value of (F(˜w1) -

F ∗ ) requires larger K thus longer delay to decrease the bound in (3.43). The intuition is that if the initial weight is too far away from F ∗ , less frequent synchronizations can lead to faster convergence rate. Less frequent averaging implies higher variance in general. If the search is still far away from the solution, a stochastic gradient direction with larger variance may be preferred.

Note that Kopt > 1 doesn't necessarily mean that Kopt is very close to 1. In our experiments (see section 4), Kopt can be as large as 16 or 32.

4 Experiments

We conduct experiments to validate our analysis on the scalability of K-AVG vs. ASGD implementations, the optimal delay in averaging ( optimal value of K), the convergence comparison with the sequential algorithm, i.e., SGD, and the comparison of the learning rates allowed with ASGD.

In our application gradient descent is implemented with Torch, and the communication is implemented using CUDAaware openMPI 2.0 through the mpiT library. All implementations use the cuDNN library for forward propagation and backward propagation. Our experiments are done on a cluster of 32 Minsky nodes interconnected with Infiniband. Each node is an IBM S822LC system containing 2 Power8 CPUs with 10 cores each, and 4 NVIDIA Tesla P100 GPUs.

We experiment with the CIFAR-10 data set using the vgg and nin models. CIFAR-10 contains 50, 000 training images and 10, 000 test images, each associated with 1 out of 10 possible labels.

4.1 Comparison with ASGD

Theorem 3.3 shows that the convergence bound of K-AVG is not affected much by the scaling of P but rather by K. On the contrary, the convergence bound of ASGD increases linearly with P . Thus we expect poorer convergence behavior of ASGD implementations at large P in comparison with KAVG.

Figures 1 and 2 compare the performance of K-AVG with two ASGD implementations, Downpour and EAMSGD, for two neural networks, vgg [Simonyan and Zisserman, 2014] and nin [Lin et al., 2013], respectively. We use P = 8, 16, 32, 64 and 128 learners, and show the test accuracies. All implementations use the same initial learning rate (γ0 = 1) and learning rate adaptation schedule (reduce γ by half after 50 epochs). The batch size is fixed at ¯B = 16. We run for 600 epochs with K = 16 for K-AVG.

In both figures, K-AVG always achieves better test accuracy than Downpour and EAMSGD. The test accuracies for Downpour and EAMSGD decreases as P increases (the effect is more pronounced for vgg). When P reaches 128, the accuracies of Downpour and EAMSGD both degrade to around 10%, i.e, random guess. The ASGD implementations do not converge with γ0 = 1 at P = 128.

When we set the initial learning rate γ0 = 0.1, Downpour achieves around 80% and 87% test accuracies for vgg and nin, respectively; EAMSGD still does not converge.

K-AVG achieves better test accuracy than ASGD implementations, and in our experiment it is also faster. We measure wall-clock times for all implementations after 600 epochs. Naturally for K-AVG, K impacts the ratio of communication vs. computation. We still use K = 16.

Figures 3 and 4 show the speedups of K-AVG over Downpour and EAMSGD, with vgg and nin, respectively. In Figure 3, when P = 8, the ASGD implementations are slightly faster than K-AVG. As P increases, the speedup increases. When P = 128, the speedups are around 2.5 and 2.6 over Downpour and EAMSGD, respectively. Similar behavior is observed in Figure 4 for nin.

4.2 The Optimal Delay in Averaging for K-AVG

For K-AVG, K regulates its behavior. From the execution time perspective, larger K results in fewer communications to process a given number of data samples. From the convergence perspective, smaller K reduces the variances among learners. In the extreme case where K = 1, K-AVG is equivalent to synchronous parallelization of SGD. People tend to think that smaller K results in faster convergence in terms of number of data samples processed. As discussed in Section 3.4, there are scenarios where Kopt is not 1.

We evaluate the convergence behavior of K-AVG with different K values for vgg and nin. We experiment with K = 1, 2, 4, 16, 32, and 64. Figures 5 and 6 show the test accuracies achieved after 600 epochs for P = 8, 16, 32, 64, and 128, for vgg and nin, respectively. Again we use the initial γ0=1, and after every 50 epochs, γ is reduced by half. The batch size is fixed as ¯B = 32.

In Figure 5, strikingly, none of the experiments show the optimal value of K for K-AVG is 1. Kopt ranges from 32 (when P = 8) farthest away from 1, to 2 (when P = 64), closest to 1. In this set of experiments, as P increases, Kopt tends to decrease. Also with smaller P, K-AVG is more forgiving in terms of the choices of K. For example, when P = 8, test accuracies for different K are similar. With larger P , however, choosing a K that is too large has severe punishing consequences. For example, when P = 128, Kopt = 4, and the test accuracy degrades rapidly with the increase of K beyond 4.

In Figure 6, almost all experiments showKopt=1. The exception is with p = 8, and the accuracy is slightly higher (by 0.27%) at K = 8 than K = 1. Again we see for small p the choices of K is not critical, while for large p the degradation in accuracy is rapid with the increase of K beyond Kopt.

Since in our experiments the same hyper-parameters are used for vgg and nin, it is likely that the Lipschitz constant L largely determines the differences in Kopt between the two cases.

4.3 Convergence Comparison with SGD

We compare the performance of K-AVG against the sequential implementation, that is, SGD. We evaluate the test accuracies achieved and the wall clock time used.

Fig. 7 shows the accuracy gap between K-AVG and SGD. With 8 and 16 learners, K-AVG is slightly worse than sequential SGD for vgg but better than nin. K-AVG and SGD achieve comparable accuracies with 32 learners. As the number of learners reaches 64 and 128, significant accuracy degradation, up to 8.8% is observed for vgg. Interestingly, the accuracy degradation for nin is still within 1.3% with 128 learners. Compared with sequential SGD, K-AVG achieves almost linear speedup in terms of epoch processing time.

5 Conclusion

In this paper, we adopt and analyze K-AVG for solving large scale machine learning problems with nonconvex objectives. We establish the convergence results of K-AVG under both fixed and diminishing stepsize, and show that with a properly chosen sequence of stepsizes, K-AVG achieves similar convergence rate consistent with its sequential counterpart. We show that K-AVG scales better than ASGD with properly chosen K when P is large. K-AVG allows larger stepsizes that still guarantees convergence while ASGD may fail to converge. Contrary to popular belief, we show that the length of delay to average learning parameters among parallel learners is not necessarily to be 1. Although, a proper choice of Kopt remains unknown, we analytically explain the dependence of Kopt on other parameters which we hope can serve as a users' guide in practical implementations.