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

Hao Yu, Rong Jin+1 more|arXiv (Cornell University)|May 9, 2019

Stochastic Gradient Optimization Techniques147 citations

TL;DR

The paper proves that parallel restarted SGD with momentum achieves the same O(1/√(NT)) convergence (linear speedup) as vanilla distributed SGD, while greatly reducing communication rounds, under mild assumptions.

ABSTRACT

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.

Motivation & Objective

• Motivate the study of linear speedup with momentum SGD in distributed non-convex optimization.
• Analyze a parallel restarted SGD with momentum (PR-SGD-Momentum) and establish its convergence and communication efficiency.
• Show how momentum can be integrated with reduced communication while preserving convergence rates.
• Compare momentum variants (Polyak’s and Nesterov’s) and demonstrate they share similar convergence properties under the proposed framework.

Proposed method

• Study problem (1) under smoothness and bounded variance/heterogeneity Assumption 1.
• Propose Parallel Restarted SGD with Momentum (Algorithm 1) with two momentum options (Polyak’s and Nesterov’s).
• Show that node-averaged iterate bar{x}^{(t)} follows a momentum SGD-like dynamic when gradients are aggregated periodically.
• Derive convergence bound for Option I (Polyak’s momentum) showing dependence on learning rate γ, momentum β, and synchronization interval I.
• Extend analysis to Option II (Nesterov’s momentum) and show similar rate results.
• Provide two corollaries: (i) linear speedup with γ = √N/√T and I=1, (ii) reduced communication achieving O(N^{3/2}T^{1/2}) rounds for identical data and O(N^{3/4}T^{3/4}) for non-identical data.

Experimental results

Research questions

• RQ1 Can distributed momentum SGD attain the same linear speedup (O(1/√(NT))) as momentum-free distributed SGD in non-convex settings?
• RQ2 How does the communication interval I affect convergence and can it be reduced without sacrificing speedup?
• RQ3 Do Polyak’s and Nesterov’s momentum yield the same order of convergence under the proposed framework?
• RQ4 What are the communication round complexities in identical versus non-identical data scenarios for momentum-based distributed training?
• RQ5 How does decentralized communication affect the linear speedup property?

Key findings

• PR-SGD-Momentum achieves O(1/√(NT)) convergence, i.e., linear speedup, under Assumption 1 with suitable γ and I.
• With identical data (κ=0), T iterations require O(N^{3/2}T^{1/2}) communication rounds to retain linear speedup.
• With non-identical data (κ>0), T iterations require O(N^{3/4}T^{3/4}) communication rounds for linear speedup.
• Both Polyak’s and Nesterov’s momentum lead to the same convergence rate up to constants, yielding the same linear speedup properties.
• Decentralized communication (Algorithm 2) also achieves linear speedup under Assumptions 1 and 2, with appropriate γ, and shows convergence O(1/√(NT)) under standard mixing conditions (ρ).
• Experiments on ResNet-56 with CIFAR-10 validate faster convergence and demonstrate the practical benefits of communication-skipping momentum methods.