[Paper Review] On the Generalization Benefit of Noise in Stochastic Gradient Descent

Samuel Smith, Erich Elsen+1 more|arXiv (Cornell University)|Jun 26, 2020

Stochastic Gradient Optimization Techniques22 citations

TL;DR

This paper demonstrates that noise in stochastic gradient descent (SGD) enhances generalization in deep learning, even after rigorous hyperparameter tuning. Using a stochastic differential equation (SDE) framework, the authors show that small to moderate batch sizes outperform very large batches on test accuracy—even when training loss is lower for large batches—confirming a persistent generalization benefit from gradient noise.

ABSTRACT

It has long been argued that minibatch stochastic gradient descent can generalize better than large batch gradient descent in deep neural networks. However recent papers have questioned this claim, arguing that this effect is simply a consequence of suboptimal hyperparameter tuning or insufficient compute budgets when the batch size is large. In this paper, we perform carefully designed experiments and rigorous hyperparameter sweeps on a range of popular models, which verify that small or moderately large batch sizes can substantially outperform very large batches on the test set. This occurs even when both models are trained for the same number of iterations and large batches achieve smaller training losses. Our results confirm that the noise in stochastic gradients can enhance generalization. We study how the optimal learning rate schedule changes as the epoch budget grows, and we provide a theoretical account of our observations based on the stochastic differential equation perspective of SGD dynamics.

Motivation & Objective

• To resolve the debate on whether small-batch SGD generalizes better than large-batch SGD by conducting controlled experiments with rigorous hyperparameter tuning.
• To investigate how the optimal learning rate schedule evolves with increasing training budget (epoch count), particularly distinguishing between training loss minimization and test accuracy maximization.
• To provide a theoretical account of SGD dynamics using the stochastic differential equation (SDE) perspective, identifying noise-dominated and curvature-dominated regimes.
• To clarify the role of learning rate schedules in generalization, especially the importance of high initial learning rates and rapid decay in late training phases.
• To validate that the generalization benefit of noise persists under constant step budgets, challenging claims that it arises solely from suboptimal tuning or compute constraints.

Proposed method

• Employed a stochastic differential equation (SDE) framework to model SGD dynamics, distinguishing between noise-dominated and curvature-dominated regimes based on batch size and loss landscape curvature.
• Conducted controlled experiments across multiple standard models (e.g., Wide-ResNet on CIFAR-10) with fixed training iterations and systematic hyperparameter sweeps over learning rates at each batch size.
• Used a learning rate schedule parameterized by initial and final learning rates (with a fixed decay rate), enabling comparison of generalization performance across different learning rate trajectories.
• Performed grid searches over learning rates at each batch size to ensure optimal tuning, avoiding biases from suboptimal hyperparameter choices.
• Analyzed the decay rate of optimal learning rates as the epoch budget increased, comparing schedules that minimize training loss versus maximize test accuracy.
• Identified the boundary between noise-dominated and curvature-dominated regimes by estimating the largest stable learning rate for small-batch training.

Experimental results

Research questions

• RQ1 Does stochastic gradient noise provide a genuine generalization benefit in deep learning, independent of hyperparameter tuning or compute budget differences?
• RQ2 How does the optimal learning rate schedule for maximizing test accuracy differ from that for minimizing training loss as the training budget increases?
• RQ3 Can the SDE perspective of SGD accurately predict the existence of two distinct training regimes—noise-dominated and curvature-dominated—and their impact on generalization?
• RQ4 What is the relationship between batch size, learning rate, and test set performance when training for the same number of iterations?
• RQ5 Is there an 'optimal temperature' (ratio of learning rate to batch size) that promotes generalization, and how does it scale with compute budget?

Key findings

• Small and moderately large batch sizes (e.g., 2048) achieved significantly higher test accuracy (94.9%) than very large batches (16384) on a 16-4 Wide-ResNet trained on CIFAR-10, even after 9725 training updates and full hyperparameter tuning.
• The optimal learning rate for maximizing test accuracy decays very slowly with increasing epoch budget—only by a factor of 2 when the budget increases 128-fold—while the optimal learning rate for minimizing training loss decays rapidly.
• The optimal initial learning rate for test accuracy maximization is consistently higher than that for training loss minimization, and the optimal final learning rate for test accuracy is consistently lower.
• The generalization benefit of noise persists under constant step budgets, confirming that it is not an artifact of longer training or suboptimal tuning.
• The boundary between noise-dominated and curvature-dominated regimes can be estimated by identifying the largest stable learning rate in a small-batch sweep, suggesting a practical strategy for batch size selection.
• Learning rate schedules that maintain high initial learning rates and rapidly decay them in later phases generalize better, supporting the analogy to simulated annealing and the importance of early noise in shaping generalization.