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

Elad Hoffer; Itay Hubara; Daniel Soudry

Source-linked AI summary

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

Elad Hoffer, Itay Hubara, Daniel Soudry

arXiv:1705.08741v2May 24, 2017stat.MLcs.LG

TL;DR

Large-batch training was associated with a persistent generalization gap, whose origin and mitigation remained unresolved. The paper analyzes the initial high-learning-rate phase using an ultra-slow random-walk model and experiments, finding that update count rather than batch size explains the gap and that Ghost Batch Normalization can reduce it without increasing updates.

Problem
Large batch sizes had been observed to degrade generalization, raising an unresolved question about the gap’s origin and how to reduce it.
Method
The paper models initial SGD training as a random walk on a random potential and evaluates learning-rate, batch-normalization, and training-regime adaptations.
Results
The generalization gap is attributed to the number of updates rather than batch size, while Ghost-BN significantly improves large-batch generalization without increasing training time.
Takeaways & Limitations
Large mini-batches can generalize as well as small mini-batches when the training regime supplies sufficient updates and uses the proposed adaptations.
Takeaways & Limitations
The sharpness-based explanation must exclude degenerate parameter directions, which can alter curvature without affecting the loss.

Abstract

from arXiv · show

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.

1 Introduction

Large-batch training was associated with a persistent generalization gap, but the paper argues that the gap reflects too few updates rather than batch size itself. It proposes an ultra-slow diffusion account and training adaptations that reduce or eliminate the gap.

Large-batch training had been reported to increase generalization error, even when training continued until the loss stopped improving.
Understanding the gap mattered because larger batches can increase parallelization and potentially reduce learning time.
The paper models the initial learning phase as a high-dimensional random walk on a random potential with ultra-slow logarithmic weight-distance growth.
Adjusting the learning rate and batch normalization reduced the gap from 5% to 1% −2%.
Generalization keeps improving during the initial high-learning-rate phase despite no observable changes in training or validation errors, and relates to distance from initialization.
There is no inherent generalization gap: adapting the number of iterations lets large-batch training generalize as well as small-batch training.

2 Training with a large batch

The paper studies momentum SGD under a conventional fixed-learning-rate and annealing regime, while contrasting it with adaptive optimizers and prior explanations of the large-batch generalization gap.

The standard training regime uses gradient steps with momentum, a fixed learning rate, and exponential annealing every few epochs.
Adaptive per-parameter methods such as Adam, Rmsprop, and Adagrad are presented as alternatives known to benefit SGD optimization convergence.
The study focuses on momentum SGD and evaluates most results using the Resnet44 topology.
Prior work reported that large batches increase generalization error and that this gap persists when training is continued until the loss ceases to improve.
A prior hypothesis attributed the gap to small-batch estimation noise steering weights away from sharp minima toward flatter minima with better generalization.

3 Theoretical analysis

The theoretical analysis treats SGD as a random walk on the loss landscape and uses ultra-slow diffusion to explain logarithmic growth in weight distance. Experiments across batch sizes support this account and emphasize update count and diffusion rate.

SGD and landscape model: SGD updates the parameter vector using a mini-batch mean gradient computed from randomly selected training samples.
SGD and landscape model: The weight vector can be viewed as a particle performing a random walk on the loss landscape, with momentum analogous to inertia.
Caveat: Certain degenerate parameter directions can change curvature without changing the loss, so they must be excluded from the sharpness-based explanation.
Ultra-slow diffusion: The proposed random-walk-on-random-potential model predicts ultra-slow diffusion, with distance scaling as ∥w_t−w_0∥∼(log t)^(2/α).
Empirical comparison: The observed weight distance from initialization increases logarithmically with training updates, matching the model with α = 2.
Empirical comparison: Similar logarithmic distance curves appear for all batch sizes, while slopes differ and smaller batches reach greater distances after the same number of epochs because they perform more updates.
Implications: Under the flat-minima assumption, reaching wider minima requires both high diffusion rates that avoid instability and many training iterations.

4 Matching weight increment statistics for different mini-batch sizes

The paper matches large-batch and small-batch update statistics by adjusting learning rates and, optionally, adding multiplicative noise or using Ghost Batch Normalization. These interventions target diffusion-rate differences while preserving the mean update where possible.

Learning-rate adjustment: Increasing the learning rate by the square root of mini-batch size keeps the update covariance approximately constant across batch sizes.The authors note that the resulting increase in mean step size is typically negligible relative to the update standard deviation.
Noise matching: Multiplicative Gaussian noise can match both first- and second-order statistics of small-batch updates while keeping covariance constant across batch sizes.The added-noise construction changes the covariance without changing the mean steps.
Noise matching: Dropout, dropconnect, and label noise do not accurately match small-batch updates because they alter covariance structure rather than only its scale.The authors report that these noise types did not reduce the large-batch generalization gap.
Ghost Batch Normalization: Ghost Batch Normalization computes training statistics on small virtual batches while using full-batch statistics during inference.This modification substantially reduces generalization error in the authors’ experiments.
Combined adjustments: Combining learning-rate and Ghost-BN adjustments improves generalization and produces more closely matched logarithmic weight-distance slopes across batch sizes.The authors attribute an observed constant shift to gradient clipping and assume it does not harm performance.

5 Adapting number of weight updates eliminates generalization gap

The authors argue that large-batch training’s generalization gap reflects too few weight updates rather than batch size itself. They stretch the training regime to match small-batch update counts and combine this with learning-rate adjustment.

Training-regime adaptation: Continuing the initial learning rate after validation error plateaus can improve final accuracy, while later learning-rate drops produce a sharp validation-error decrease.This observation challenges the practice of reducing the learning rate solely when validation performance appears flat.
Training-regime adaptation: The authors attribute the generalization gap to the relatively small number of updates rather than the batch size.Their adaptation is motivated by the logarithmic growth of weight distance during the initial high-learning-rate phase.
Training-regime adaptation: Stretching each training period by the large-to-small batch-size ratio makes the number of optimization steps identical to the small-batch regime.The modified schedule increases the number of epochs proportionally to the relative batch size.
Training-regime adaptation: Combining regime adaptation with learning-rate adjustment completely eliminates the previously observed generalization gap.Figure 3 compares large-batch regimes adapted to match small-batch performance.

6 Experiments

Experiments span multiple image-classification datasets and neural-network architectures. Validation results show progressively better large-batch generalization with learning-rate tuning, Ghost Batch Normalization, and regime adaptation.

Experimental setting: The evaluation covers MNIST, CIFAR-10, CIFAR-100, and ImageNet classification tasks.The datasets range from 10-class digit and natural-image tasks to 1000-class ImageNet.
Experimental setting: The experiments use fully connected, convolutional, VGG, ResNet44, Wide-ResNet16-4, and AlexNet models with momentum SGD.Architectures are selected across the listed datasets, including ResNet44 for CIFAR-10 and AlexNet for ImageNet.
Results: Large-batch learning-rate tuning diminishes the generalization gap, Ghost-BN provides additional improvement, and regime adaptation removes it completely in Table 1.The comparison uses small-batch and large-batch validation accuracy, with GBN and RA denoting Ghost Batch Normalization and regime adaptation.

7 Discussion

The paper argues that large-batch generalization problems arise from too few updates rather than batch size itself, and proposes training adjustments to close the gap. It also questions standard learning-rate annealing practices while noting that training speedup from large batches remained open in this work.

Discussion: Large batches offer parallelization potential, but this work leaves whether they reduce training wall-clock time unresolved.The paper distinguishes closing the generalization gap from achieving a speedup through larger batches.
Discussion: Large-batch training can match small-batch generalization when the training regime is adapted appropriately.The proposed remedies include momentum SGD, gradient clipping, batch-size learning-rate adaptation, Ghost Batch Normalization, and sufficient high-learning-rate iterations.
Discussion: Later ImageNet work reported significant large-batch speedups using similar training practices, with linear rather than square-root learning-rate scaling.The paper notes that linear scaling worked less well on CIFAR-10 and later on some ImageNet architectures.
Discussion: The authors suggest that common optimization and learning-rate annealing rules may be suboptimal because generalization can improve during extended updates without apparent validation-error change.They cite logarithmic weight growth and slowly improving margin as related observations in separable logistic regression.
Discussion: The proposed explanation attributes the generalization gap to the amount of updates rather than an inherent problem with large mini-batches.The conclusion links performance degradation to insufficient updates during training.

A Derivation of eq. (6)

This derivation analyzes mini-batch selector statistics under sampling with and without replacement to establish the covariance relation used for SGD gradients. Without replacement, cross-mini-batch covariance becomes much smaller when the batch is small relative to the dataset.

A Derivation of eq. (6): The mini-batch gradient is introduced as the quantity whose covariance is derived from selector-variable statistics.The surrounding derivation states that the mini-batch gradient is an unbiased estimator under the relevant condition.
A Derivation of eq. (6): Sampling with replacement makes different mini-batches uncorrelated, yielding the stated selector-statistics relation.The derivation explicitly treats mini-batches as uncorrelated in this sampling case.
A Derivation of eq. (6): Sampling without replacement correlates selector variables across mini-batches because previously selected samples cannot be chosen again.The derivation introduces second-order statistics for selector variables indexed by different mini-batches.
A Derivation of eq. (6): For M ≪N, covariance between different mini-batches is much smaller than the within-batch contribution.This behavior is used to confirm the covariance expression in eq. (6).

B Estimating α from random potential

The appendix tests the random-potential model's α = 2 prediction by measuring loss variation against weight distance. It also compares large-batch training adjustments through learning-rate, normalization, and update-distance behavior.

B Estimating α from random potential: α = 2 predicts that the standard deviation of loss differences increases linearly with weight distance, and the experiment confirms this behavior.The model equivalently predicts loss auto-covariance increasing with the square of weight distance.
B Estimating α from random potential: The experiment samples 1000 random directions and scalar displacements from an initialized Resnet44, then records weight distance and loss.Weights are constructed as w = w0 + zv, with the maximum distance set near 10.
B Estimating α from random potential: Figure 5 compares regime-adapted large-batch training with unadapted batch-2048 training.The caption identifies the comparison but does not state the plotted outcome.
B Estimating α from random potential: The appendix also includes a comparison of L2 distance from initialization across batch sizes.The supplied figure label identifies the measured quantity and comparison dimension without reporting an outcome.