PaperlayerView exact revision on arXiv

Source-linked AI summary

On the Linear Speedup Analysis of Communication Efficient Momentum SGD for Distributed Non-Convex Optimization

Hao Yu, Rong Jin, Sen Yang

arXiv:1905.03817v1math.OCcs.LG

TL;DR

Existing theory analyzes communication-efficient distributed SGD without momentum, despite momentum’s common use in distributed deep-network training and its practical importance for convergence and test accuracy. This paper studies parallel restarted SGD with momentum and proves linear speedup with reduced communication, including decentralized communication results.

  • Problem

    Theory for communication-efficient distributed SGD largely omits momentum, leaving unclear whether distributed momentum SGD retains linear speedup and reduced communication despite its widespread practical use.

  • Method

    The paper analyzes parallel restarted SGD with momentum, allowing Polyak’s or Nesterov’s local updates and periodic global averaging.

  • Results

    The method achieves O(1/NT) convergence with O(N 3/2T 1/2) or O(N 3/4T 3/4) communication rounds depending on whether workers access identical or non-identical data sets, and decentralized communication also achieves O(1/NT) convergence.

  • Takeaways & Limitations

    Distributed momentum SGD can possess the same linear speedup property as distributed SGD while reducing communication rounds.

  • Takeaways & Limitations

    The theory establishes linear speedup using constant learning rates, whereas practitioners often use decaying learning rates for good test accuracy.

Abstract

from arXiv · show

Recent developments on large-scale distributed machine learning applications, e.g., deep neural networks, benefit enormously from the advances in distributed non-convex optimization techniques, e.g., distributed Stochastic Gradient Descent (SGD). A series of recent works study the linear speedup property of distributed SGD variants with reduced communication. The linear speedup property enable us to scale out the computing capability by adding more computing nodes into our system. The reduced communication complexity is desirable since communication overhead is often the performance bottleneck in distributed systems. Recently, momentum methods are more and more widely adopted in training machine learning models and can often converge faster and generalize better. For example, many practitioners use distributed SGD with momentum to train deep neural networks with big data. However, it remains unclear whether any distributed momentum SGD possesses the same linear speedup property as distributed SGD and has reduced communication complexity. This paper fills the gap by considering a distributed communication efficient momentum SGD method and proving its linear speedup property.

1. Introduction

Distributed SGD can provide linear speedup, but synchronization creates communication bottlenecks, while existing theory largely analyzes methods without momentum despite momentum’s practical importance. The paper addresses this gap by analyzing communication-efficient distributed momentum SGD and establishing linear speedup with reduced communication.

  • Motivation: With N workers, parallel minibatch SGD can achieve O(1/NT) convergence, but synchronizing every iteration makes communication a bottleneck.Linear speedup means convergence requiring N times fewer iterations than single-worker SGD.
  • Motivation: Communication-efficient distributed SGD variants reduce synchronization rounds through local aggregation, compression, sparsification, or strategically skipped communication.Prior analyses established fast O(1/NT) convergence with fewer communication rounds for some SGD variants.
  • Research gap: Existing convergence analyses focus on parallel SGD without momentum, although momentum SGD is commonly used because it often converges faster and generalizes better.Momentum is also recommended for ResNet image classification and is widely used in distributed deep-network training.
  • Research gap: On CIFAR10, momentum SGD achieves roughly 2.5% higher final test accuracy than vanilla SGD for ResNet56 under the stated single-GPU settings.The comparison uses hyper-parameters suggested by He et al. (2016).
  • Contribution: The paper analyzes parallel restarted SGD with momentum, a momentum extension of local SGD that reflects model averaging with momentum in practice.Algorithm 1 supports both Polyak’s and Nesterov’s momentum for local worker updates.
  • Contribution: The method achieves O(1/NT) convergence with O(N 3/2T 1/2) communication rounds for identical data and O(N 3/4T 3/4) for non-identical data.The paper presents this as the first proof of linear speedup with communication reduction for distributed momentum SGD in non-convex stochastic optimization.
  • Contribution: The analysis also improves communication complexity results for parallel SGD without momentum and handles non-identical training sets, including sharded-data and federated-learning settings.It relaxes a bounded gradient moment assumption to a milder bounded variance assumption in the identical-data case.

2. Parallel Restarted SGD with Momentum

The paper analyzes parallel restarted SGD with momentum under standard smoothness and bounded-variance assumptions, explaining how periodic synchronization controls worker disagreement. It proves linear speedup with reduced communication for both Polyak’s and Nesterov’s momentum.

  • Preliminary assumptions: The analysis assumes smooth local objectives and bounded stochastic-gradient variance and objective heterogeneity.The heterogeneity parameter κ is zero when workers access identical data sets, while bounded variance is milder than bounded second-moment assumptions used in earlier work.
  • Preliminary analysis: Independent stochastic gradients reduce averaged variance by a factor of N even when workers evaluate them at different local points.When synchronization is skipped, local points can diverge, so the analysis also controls gradient concentration through worker deviations.
  • Algorithm: Parallel restarted SGD with momentum lets workers perform local Polyak or Nesterov updates and periodically reset local solutions and momentum buffers to node averages.With synchronization interval I, communication occurs only once every I iterations; choosing β = 0 recovers local SGD.
  • Polyak’s momentum: Theorem 1 establishes a convergence bound for Polyak’s momentum under a learning-rate condition involving β and L.The bound contains terms depending on iteration count, worker number, stochastic variance, synchronization interval, and objective heterogeneity.
  • Linear speedup: Polyak’s momentum achieves O(1/√(NT)) convergence, providing linear speedup with N workers when T is sufficiently large.The result extends to communication skipping: T iterations require O(N 3/2 T 1/2) rounds for κ = 0 and O(N 3/4 T 3/4) rounds for κ ≠ 0.
  • Comparison with prior analyses: For parallel SGD without momentum, the paper’s communication complexity improves the state of the art under both identical and non-identical data settings.The analysis relaxes a bounded-gradient-moment assumption to bounded variance and reduces communication complexity for identical training data.
  • Nesterov’s momentum: Nesterov’s momentum has the same dependence on γ, N, and I as Polyak’s momentum and therefore achieves the same linear-speedup communication complexity.Theorem 2 provides the corresponding convergence analysis for Nesterov’s option.

3. Extension: Momentum SGD with Decentralized Communication

The paper extends momentum SGD to decentralized communication using neighborhood aggregation and establishes linear-speedup convergence under mixing-matrix assumptions.

  • Method: Algorithm 2 is distributed momentum SGD with decentralized communication, using a mixing matrix W to encode network connectivity.Each worker maintains local solutions and momentum buffers and updates them through neighborhood weighted averages.
  • Assumptions: The mixing matrix is assumed to be symmetric and doubly stochastic, with λ1(W) = 1 and spectral conditions governing network mixing.These assumptions support the convergence analysis for the decentralized setting.
  • Communication trade-off: Unlike Algorithm 1, Algorithm 2 aggregates at every iteration, increasing communication rounds by a factor of I but requiring only neighborhood averages.This trade-off makes the method more suitable for mobile distributed learning with heterogeneous network topologies.
  • Guarantee: With an appropriate learning-rate condition, Algorithm 2 achieves O(1/√(NT)) convergence with linear speedup.The guarantee is stated for problem (1) under Assumptions 1 and 2.
  • Relation to prior methods: Algorithm 2 with β = 0 reduces to decentralized SGD, while this work establishes the same linear-speedup property for decentralized momentum SGD.The paper identifies this as the first such result for decentralized momentum SGD.

4. Experiments

Experiments train ResNet on CIFAR-10 and ImageNet with momentum, comparing communication-skipping Algorithm 1 against classical synchronized momentum SGD. Wall-clock measurements show the communication-skipping variants are faster.

  • Setup: Experiments train ResNet on CIFAR-10 and ImageNet using 8 NVIDIA P100 GPUs, local batch size 64, and momentum coefficient β = 0.9.The learning rate starts at 0.3 and is divided by 10 after 80 and 120 jointly accessed epochs.
  • Theory–practice caveat: The experiments use decaying learning rates for practical test accuracy, whereas the theory establishing linear speedup uses constant learning rates.Supplementary experiments with constant learning rates are reported to verify the theoretical behavior.
  • Figures: Figure 2 reports convergence versus epochs for ResNet56 over CIFAR10, while Figure 3 reports convergence versus wall-clock time for the same setting.Both figures track training loss and test accuracy.
  • Epoch-based comparison: Algorithm 1 is evaluated with I ∈ {4, 8, 16, 32} against classical parallel minibatch SGD with momentum, equivalent to I = 1.The comparison examines training-loss and test-accuracy convergence over jointly accessed epochs.
  • Wall-clock outcome: Algorithm 1 variants with I > 0 are much faster than classical parallel momentum SGD when measured by wall-clock time because they skip communication steps.Additional experiments cover ImageNet, model averaging, and Nesterov momentum in the supplement.

5. Conclusion

The paper concludes that parallel restarted SGD with momentum and decentralized momentum SGD achieve linear-speedup convergence with reduced communication under their respective settings.

  • Conclusion: Parallel restarted SGD with momentum achieves O(1/√(NT)) convergence with O(N^3/2T^1/2) or O(N^3/4T^3/4) communication rounds.The communication bound depends on whether workers access identical or non-identical objective functions.
  • Conclusion: The decentralized extension is also shown to achieve O(1/√(NT)) convergence with linear speedup.Together, these results extend linear-speedup guarantees from distributed SGD to momentum variants.

6. Supplement

The supplement provides a formal proof of a lemma using unbiasedness, independence, a variance decomposition identity, and Assumption 1.

  • Proof basis: The proof uses unbiased stochastic gradients independently sampled across workers.It invokes E[gi] = ∇fi(xi) and the identity E[∥Z∥2] = E[∥Z − E[Z]∥2] + ∥E[Z]∥2.
  • Proof basis: Independence and zero means justify summing the expected squared norms of the stochastic-gradient noise terms.The final proof step follows from Assumption 1.

6.2. Proof of Lemma 2

The lemma is justified through algebraic inequalities, smoothness of the component functions, convexity, and Jensen’s inequality.

  • The proof begins with simple algebraic manipulations.
  • A three-vector norm inequality bounds the squared norm of a sum.The bound is ∥a1 + a2 + a3∥2 ≤ 3∥a1∥2 + 3∥a2∥2 + 3∥a3∥2.
  • Smoothness bounds gradient differences by L times the corresponding distance between points.
  • Convexity of the squared norm and Jensen’s inequality, followed by component smoothness, provide the remaining bound.

6.3. Proof of Theorem 1

Theorem 1 is proved for Option I by introducing an auxiliary momentum sequence, bounding its deviations and consensus errors, and combining these bounds with smoothness and stochastic-gradient properties.

  • An auxiliary sequence z̄(t) is introduced from node-average iterates to facilitate the momentum convergence analysis.The sequence is related to standard momentum analyses and is defined using β and consecutive node averages.
  • The proof develops lemmas controlling the auxiliary sequence, momentum differences, and node-average disagreement under periodic restarts.Algorithm 1 restarts SGD with momentum every I iterations by resetting the momentum state.
  • The final argument imposes step-size and restart conditions, rearranges the resulting inequality, and divides by T to obtain the theorem’s rate bound.The proof explicitly uses conditions involving 12L2γ2I2 and a nonnegative coefficient ensured by the choice of γ.
  • The disagreement analysis separates iterations at multiples of I from intervening iterations and uses t − t0 < I to bound local deviations.
  • The proof combines the auxiliary-sequence bounds with smoothness, Jensen’s inequality, and stochastic-gradient independence relations.

6.4. On the equivalence between (5) and (14)

The paper shows that the two update formulations in (5) and (14) generate identical solution sequences.

  • The equivalence proof initializes the auxiliary variables and verifies that both formulations produce the same first iterate.
  • For later iterations, substitution of the updates in (5) yields the auxiliary-variable recursion used in (14).
  • Therefore, formulations (5) and (14) produce the same solution sequences.

6.5. Proof of Theorem 2

Theorem 2 is proved for Option II by introducing a Nesterov-specific auxiliary sequence, establishing counterparts of the earlier lemmas, and reusing the main proof structure with corresponding substitutions.

  • A distinct auxiliary sequence ȳ(t) is introduced for Option II and Nesterov’s momentum.It is defined from node-average iterates and includes a momentum-related correction term.
  • The increment ȳ(t+1) − ȳ(t) matches the corresponding increment z̄(t+1) − z̄(t) for Option I.
  • Counterparts of Lemmas 3–5 establish bounds for Option II’s auxiliary sequence, momentum, and disagreement terms.
  • The disagreement proof again uses periodic momentum restarts and bounds iterations between consecutive multiples of I.
  • Theorem 2 follows the proof of Theorem 1 after replacing Lemmas 3–5 with Lemmas 6–8 and making minor changes.
  • The final derivation combines the new lemmas with smoothness and stochastic-gradient relations, then rearranges and divides by T.

6.6. Proof of Theorem 3

The proof analyzes node averages and disagreement dynamics for Algorithm 2 with Polyak momentum, then extends the resulting lemmas to establish convergence under a stepsize condition.

  • Algorithm 2’s global averages follow the same dynamics as Algorithm 1 despite decentralized local averaging.
  • The proof introduces auxiliary average sequences and shows their recursions match those used for Algorithm 1 with Option I.
  • The main technical challenge is obtaining a tight bound for the disagreement quantity needed as the counterpart of Lemma 5.
  • Matrix formulations track disagreement from node averages, using the mixing matrix W, averaging projector Q, and spectral contraction of (I − Q)W^k.
  • The resulting bound contains an additional term representing the cost of asymmetry in local averages.
  • Under the stated assumptions and a stepsize satisfying the corollary’s condition, the proof reuses the earlier convergence argument for Algorithm 2 with Option I.

6.7. More Experiments

Experiments test Algorithm 1 on CIFAR10 and ImageNet, comparing communication-skipping momentum variants with single-worker or classical distributed baselines. The reported results verify linear-speedup behavior and show differences between momentum options and model-averaging strategies.

  • Constant learning rates: With constant learning rates on ResNet56/CIFAR10, Algorithm 1 reaches the same loss while each GPU uses 1/N as many epochs as single-worker momentum SGD.The experiment uses N ∈ {2, 4, 8} GPUs and evaluates both jointly accessed epochs and wall-clock time.
  • Model averaging comparison: Algorithm 1 averages both local momentum buffers and local models, unlike cleared-momentum model averaging, and converges slightly faster.Both methods synchronize every I = 16 iterations.
  • Model averaging comparison: Algorithm 1 attains roughly 1% better test accuracy than model averaging with cleared momentum on ResNet56 over CIFAR10.
  • ImageNet: On the harder ResNet50/ImageNet task, experiments compare Algorithm 1 with I ∈ {4, 8, 16} against classical parallel momentum SGD.The tests use eight NVIDIA P100 GPUs, local batch size 128, and learning-rate reductions after jointly accessed epochs 30 and 60.
  • Option II: Figures 8 and 9 verify Option II’s convergence analysis, while preliminary results suggest Nesterov momentum is less robust to communication skipping.Even a small I can degrade test accuracy in these experiments.
Loading 1905.03817v1…