[Paper Review] Don't Use Large Mini-Batches, Use Local SGD

Tao Lin, Sebastian U. Stich+2 more|arXiv (Cornell University)|Aug 22, 2018

Advanced Neural Network ApplicationsComputer Science87 references162 citations

TL;DR

This paper shows that large-batch SGD suffer from poor generalization and introduces post-local SGD and hierarchical local SGD to improve generalization and efficiency, outperforming large-batch baselines across standard benchmarks.

ABSTRACT

Mini-batch stochastic gradient methods (SGD) are state of the art for distributed training of deep neural networks. Drastic increases in the mini-batch sizes have lead to key efficiency and scalability gains in recent years. However, progress faces a major roadblock, as models trained with large batches often do not generalize well, i.e. they do not show good accuracy on new data. As a remedy, we propose a \emph{post-local} SGD and show that it significantly improves the generalization performance compared to large-batch training on standard benchmarks while enjoying the same efficiency (time-to-accuracy) and scalability. We further provide an extensive study of the communication efficiency vs. performance trade-offs associated with a host of \emph{local SGD} variants.

Motivation & Objective

• Motivate the generalization issues associated with very large mini-batch SGD in distributed training.
• Systematically study the trade-offs of local SGD across workers, local steps, and mini-batch sizes.
• Propose post-local SGD to recover generalization while maintaining efficiency.
• Propose hierarchical local SGD to optimize system resource use in heterogeneous hardware environments.

Proposed method

• Define local SGD where each worker performs H local SGD updates with mini-batch B_loc before averaging (Eq. 2).
• Compare local SGD to mini-batch SGD under scenarios of communication efficiency and generalization performance.
• Introduce post-local SGD by switching from standard mini-batch SGD to local SGD after a phase t′, achieving large effective batch sizes with better generalization.
• Propose hierarchical local SGD to apply local updates at multiple levels of a system hierarchy to optimize computation-communication trade-offs.
• Relate local updates to stochastic noise injection and discuss implications for training dynamics and generalization.

Experimental results

Research questions

• RQ1 Can local SGD match or exceed mini-batch SGD in time-to-accuracy under communication constraints?
• RQ2 Does local SGD improve generalization compared to large-batch SGD at the same effective batch size?
• RQ3 Does post-local SGD close the generalization gap associated with large batches without sacrificing efficiency?
• RQ4 How can hierarchical local SGD optimize resource use in heterogeneous computing environments?

Key findings

• Local SGD can serve as a communication-efficient alternative to mini-batch SGD with favorable generalization on CIFAR-10/100 and ImageNet.
• Post-local SGD closes the generalization gap of large-batch training and can achieve better generalization than both small and large batch baselines.
• Post-local SGD provides at least 1.3× speedup over the whole training with improved generalization on CIFAR, and shows strong performance on ImageNet with large global batch sizes.
• Local SGD scales better in time-to-accuracy than mini-batch SGD as the number of workers increases, due to fewer communication rounds.
• Post-local SGD can be combined with sign-based compression to further improve communication efficiency without sacrificing accuracy.
• Post-local SGD tends to reach flatter minima than large-batch SGD, contributing to improved generalization.