Source-linked AI summary
Parallel Restarted SGD with Faster Convergence and Less Communication: Demystifying Why Model Averaging Works for Deep Learning
Hao Yu, Sen Yang, Shenghuo Zhu
TL;DR
Model averaging reduces communication in distributed deep-learning training, but why it works as well as parallel mini-batch SGD remains theoretically unclear. This paper analyzes parallel restarted SGD and shows that it preserves O(1/NT) convergence for non-convex optimization while reducing communication rounds by O(T 1/4).
Problem
The theory had not established whether model averaging universally achieves parallel mini-batch SGD’s linear speed-up for non-convex deep-learning optimization with less communication.
Method
The paper analyzes parallel restarted SGD, which periodically averages workers’ local models and restarts parallel local SGD from their average.
Results
O(1/NT) convergence is achieved for non-convex optimization while communication rounds are reduced by a factor of O(T 1/4) versus classical parallel mini-batch SGD.
Takeaways & Limitations
The analysis provides convergence guarantees and averaging-frequency guidelines for using model averaging without losing linear speed-up in non-convex optimization.
Takeaways & Limitations
Periodic averaging requires faster workers to wait for slower workers, wasting computation in heterogeneous networks.
Abstract
from arXiv · showhide
In distributed training of deep neural networks, parallel mini-batch SGD is widely used to speed up the training process by using multiple workers. It uses multiple workers to sample local stochastic gradient in parallel, aggregates all gradients in a single server to obtain the average, and update each worker's local model using a SGD update with the averaged gradient. Ideally, parallel mini-batch SGD can achieve a linear speed-up of the training time (with respect to the number of workers) compared with SGD over a single worker. However, such linear scalability in practice is significantly limited by the growing demand for gradient communication as more workers are involved. Model averaging, which periodically averages individual models trained over parallel workers, is another common practice used for distributed training of deep neural networks since (Zinkevich et al. 2010) (McDonald, Hall, and Mann 2010). Compared with parallel mini-batch SGD, the communication overhead of model averaging is significantly reduced. Impressively, tremendous experimental works have verified that model averaging can still achieve a good speed-up of the training time as long as the averaging interval is carefully controlled. However, it remains a mystery in theory why such a simple heuristic works so well. This paper provides a thorough and rigorous theoretical study on why model averaging can work as well as parallel mini-batch SGD with significantly less communication overhead.
Introduction
The paper studies model averaging as a communication-efficient alternative to parallel mini-batch SGD, whose per-iteration gradient exchange can become a bottleneck. It proves that, for non-convex optimization such as deep learning, carefully scheduled averaging preserves convergence and linear speed-up while reducing communication.
- Parallel mini-batch SGD: Parallel mini-batch SGD averages workers’ stochastic gradients each iteration and offers linear speed-up, but exchanging gradients among all workers can become the communication bottleneck.It is conceptually equivalent to single-node SGD with a batch size N times larger.
- Model averaging: Model averaging reduces communication by letting workers perform local SGD steps and periodically averaging their individual models instead of exchanging gradients every iteration.The approach has been widely used in practical deep-learning training.
- Open problem: Prior theory established linear speed-up with less communication for strongly-convex stochastic optimization, but left O(1/NT) convergence for non-convex optimization open.Non-convex optimization is the setting relevant to deep learning.
- Main contribution: I = O(T 1/4/N 3/4) averaging preserves O(1/NT) convergence for non-convex optimization while saving communication by a factor of O(T 1/4/N 3/4).The paper presents this as a provable guarantee with linear speed-up and less communication, together with guidance on averaging frequency.
- Additional advantage: Model averaging can also reduce privacy and security risks in federated learning because it passes models rather than workers’ raw data or gradients.The passage attributes this advantage to models’ preservation of good differential privacy.
Parallel Restarted SGD and Its Performance · Analysis
The paper analyzes Parallel Restarted SGD (PR-SGD), where workers independently run SGD from a shared epoch-start average and communicate only at synchronization points. Under smoothness and bounded-variance assumptions, it establishes convergence, linear speedup, and reduced communication for suitable synchronization intervals.
- Analysis: Under smoothness and bounded-variance assumptions, PR-SGD controls the deviation between each local solution and their average through the synchronization interval I.Lemma 1 relates E[∥x_t − x_t^i∥^2] to I and the bounded-second-moment constant G.
- Parallel Restarted SGD and Its Performance: Within each epoch, workers need not synchronize iterations, and inter-node communication is required only to compute the epoch’s initial average.The average can be obtained through a parameter server or all-reduce.
- Analysis: For 0 < γ ≤ 1/L, Theorem 1 gives a convergence guarantee for PR-SGD under Assumption 1, measured using the average expected squared gradient norm.For non-convex optimization, the paper uses this gradient-based measure rather than objective-value convergence.
- Analysis: For T ≥ N, Corollary 1 provides a convergence bound after selecting suitable learning-rate γ and synchronization-interval I values.The corollary follows directly by substituting suitable γ and I into Theorem 1.
- Analysis: For sufficiently large T, specifically T > N^3, PR-SGD achieves linear speedup in the number of workers, improving on the T > N^5 condition cited for D-PSGD.The paper also states that PR-SGD’s linear speedup is irrelevant to network topology.
- Analysis: Choosing I ≤ T^1/4 preserves the convergence rate while reducing communication rounds by a factor T^1/4 compared with parallel mini-batch SGD or D-PSGD.Communication occurs only at iterations that are multiples of I.
Extensions
The extensions remove the need to know the time horizon in advance and adapt parallel restarted SGD to heterogeneous workers. Time-varying rates preserve the original complexity, while worker-specific epoch lengths retain convergence and communication benefits under stated assumptions.
- Unknown time horizon: Algorithm 2 uses time-varying learning rates when T is unknown and achieves the same computation and communication complexity as Algorithm 1.Its accuracy improves automatically as the algorithm runs longer.
- Unknown time horizon: Algorithm 2 is mathematically equivalent to Algorithm 1 after sequential reindexing, except that learning rates vary across epochs.Each worker restarts every epoch from the node average of the previous epoch’s final solutions.
- Heterogeneous networks: With a poor initial point, Algorithm 2 can possibly converge faster than Algorithm 1, at least during the first few epochs of heterogeneous-network training.The comparison follows from the relative first and second error terms in the corresponding bounds.
Experiment
The experiment section highlights that model averaging has shown superior training speed-ups across diverse deep learning tasks and describes a ResNet20-over-CIFAR10 setup for comparing training loss across distributed schemes.
- Empirical evidence: Model averaging has empirically achieved superior training speed-ups across CNN, VGG, DNN-GMM, and LSTM scenarios.Reported applications include MNIST, CIFAR10, speech recognition, and language modeling.
- Empirical evidence: A thorough empirical study has examined ResNet performance over CIFAR10 and ImageNet.The supplied passage identifies this study as recent work by Lin, Stich, and coauthors, but the citation is truncated.
- Experimental setup: 8 P100 GPUs were used to measure ResNet20 training loss on CIFAR10, with each worker using batch size 32 and momentum 0.9.The initial learning rate was 0.1 and was divided by 10 after 8 workers accessed 150 and 275 epochs of training data.
Conclusion
The paper studies parallel restarted SGD as a theoretical abstraction of model averaging and shows that it achieves favorable non-convex convergence with substantially fewer communication rounds than classical parallel mini-batch SGD.
- Conclusion: Parallel restarted SGD provides a theoretical abstraction of the model-averaging practice used in deep neural network training.The paper frames its analysis around this distributed-training method.
- Conclusion: O(1/NT) convergence is achieved for non-convex optimization.This is the stated convergence rate for parallel restarted SGD.
- Conclusion: O(T 1/4) fewer communication rounds are required than with classical parallel mini-batch SGD.The reduction is stated relative to the communication required by classical parallel mini-batch SGD.
Supplement
The supplement establishes worker-iterate norm bounds for Algorithms 2 and 3 and sketches the proofs of Theorems 2 and 3 under the stated assumptions. Theorem 3’s analysis accounts for progressively fewer active workers and recovers the Algorithm 2 result when all workers perform equally many updates.
- Lemma 2: Lemma 2 bounds each Algorithm 2 worker iterate by 4(γ_sK_s)^2G^2 under Assumption 1.The proof restricts Lemma 1 to each particular epoch s; G is the constant defined in Assumption 1.
- Theorem 2: Theorem 2’s proof treats sufficiently large epochs separately, then sums the resulting bounds across epochs and inner iterations.It uses γ_s = N s^2/3 for sufficiently large s, bounded stochastic-gradient second moments for earlier epochs, and x_s,T_s = x_s+1,0.
- Lemma 3: Lemma 3 bounds each Algorithm 3 worker iterate by 4γ_k^2G^2 under Assumption 1.The bound uses the common initial point of epoch s, vector-norm inequalities, fewer than N summands, and Assumption 1.
- Theorem 3: Theorem 3’s proof models stragglers through j_k, the largest worker index performing the k-th update, which decreases as k approaches I_1.At iteration k, only the first j_k workers use true stochastic gradients; the remaining workers stop updating because they are too slow.
- Theorem 3: When I_1 = I_2 = · · · = I_N, j_k = N for every k and the Theorem 3 bound reduces to (11), matching Algorithm 1 and Algorithm 2.This serves as a sanity check of the analysis in the equal-update case.