[Paper Review] Train longer, generalize better: closing the generalization gap in large batch training of neural networks

Elad Hoffer, Itay Hubara+1 more|arXiv (Cornell University)|May 24, 2017

Domain Adaptation and Few-Shot Learning38 references418 citations

TL;DR

The paper argues the generalization gap in large-batch SGD is due to too few updates, not batch size, and shows how learning-rate scaling, Ghost Batch Normalization, and regime adaptation close the gap.

ABSTRACT

Background: Deep learning models are typically trained using stochastic gradient descent or one of its variants. These methods update the weights using their gradient, estimated from a small fraction of the training data. It has been observed that when using large batch sizes there is a persistent degradation in generalization performance - known as the "generalization gap" phenomena. Identifying the origin of this gap and closing it had remained an open problem. Contributions: We examine the initial high learning rate training phase. We find that the weight distance from its initialization grows logarithmically with the number of weight updates. We therefore propose a "random walk on random landscape" statistical model which is known to exhibit similar "ultra-slow" diffusion behavior. Following this hypothesis we conducted experiments to show empirically that the "generalization gap" stems from the relatively small number of updates rather than the batch size, and can be completely eliminated by adapting the training regime used. We further investigate different techniques to train models in the large-batch regime and present a novel algorithm named "Ghost Batch Normalization" which enables significant decrease in the generalization gap without increasing the number of updates. To validate our findings we conduct several additional experiments on MNIST, CIFAR-10, CIFAR-100 and ImageNet. Finally, we reassess common practices and beliefs concerning training of deep models and suggest they may not be optimal to achieve good generalization.

Motivation & Objective

• Motivate and characterize the generalization gap observed with large-batch training in neural networks.
• Propose a stochastic optimization model (random walk on a random potential) to explain weight dynamics early in training.
• Develop practical methods to close the gap: learning-rate scaling, Ghost Batch Normalization (GBN), and regime adaptation.
• Empirically validate on MNIST, CIFAR-10/100, and ImageNet across multiple architectures.
• Reassess training practices and highlight that generalization depends on update count rather than batch size alone.

Proposed method

• Model SGD as a random walk on a random potential to explain ultra-slow diffusion of weights.
• Derive that weight distance from initialization grows logarithmically with updates (approx. log t) and link diffusion rate to batch size.
• Propose learning-rate scaling with batch size (eta ∝ sqrt(M)) to preserve update statistics.
• Introduce Ghost Batch Normalization to compute BN statistics on small ghost batches within large batches.
• Advocate regime adaptation to keep the number of updates comparable across batch sizes by extending training iterations.
• Empirically validate using standard datasets and networks, reporting accuracy gains across SB/LB regimes.

Experimental results

Research questions

• RQ1 Can the generalization gap observed with large-batch training be eliminated without increasing total training time?
• RQ2 What mechanisms explain how weight updates during early training influence final generalization, and how do batch size and update counts interplay?
• RQ3 Do adjustments like learning-rate scaling and Ghost Batch Normalization consistently reduce or remove the generalization gap across architectures and datasets?
• RQ4 Is it possible to match small-batch generalization performance by extending the training regime for large batches?

Key findings

• The generalization gap with large batches can be largely eliminated by learning-rate scaling and Ghost Batch Normalization.
• Weight distance from initialization increases logarithmically with updates, consistent across batch sizes, indicating diffusion dynamics drive generalization more than batch size per se.
• Scaling the learning rate by the square root of the batch size helps preserve update statistics and improve generalization.
• Ghost Batch Normalization substantially reduces generalization error by using small ghost batches to compute batch statistics while training with large batches.
• Adapting the number of weight updates (regime adaptation) to match small-batch iteration counts eliminates the gap, yielding comparable or better validation accuracy.
• Experiments on MNIST, CIFAR-10/100, and ImageNet show consistent gains from +LR, +GBN, and +RA, often matching or surpassing SB results.