PaperlayerView exact revision on arXiv

Source-linked AI summary

On the Generalization Benefit of Noise in Stochastic Gradient Descent

Samuel L. Smith, Erich Elsen, Soham De

arXiv:2006.15081v1cs.LGstat.ML

TL;DR

The paper revisits whether minibatch noise improves generalization, a claim challenged by arguments about tuning and compute budgets. It combines carefully tuned experiments with an SDE account of SGD, finding that smaller batches can outperform very large batches and identifying an optimal-temperature explanation for generalization.

  • Problem

    Whether minibatch noise genuinely improves deep-learning generalization remains disputed, including whether reported small-batch benefits reflect tuning or compute-budget choices.

    Source: §1 · pp. 1–2

  • Method

    The authors use an SDE perspective to analyze SGD regimes and run learning-rate-tuned experiments across batch sizes, models, and compute budgets.

    Source: §1 · p. 2, §4 · p. 5

  • Results

    Smaller batches outperform very large batches on the test set in some models under constant epoch and step budgets, even when large batches achieve smaller training losses.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors find evidence for an optimal temperature that promotes generalization, often defined by the early-training learning-rate-to-batch-size ratio.

    Source: §8 · p. 8

    The conclusions are limited to the learning-rate schedule and architectures studied, though other schedules may perform equally well without following the paper’s claims.

    Source: §8 · p. 8

Abstract

from arXiv · show

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.

1. Introduction

The paper addresses whether minibatch noise genuinely improves generalization, despite claims that small-batch advantages reflect tuning or compute-budget choices. It combines an SDE account of SGD with experiments that find test-set advantages for smaller batches and different learning-rate behavior as training budgets grow.

1. Introduction

  • Earlier work left open whether a small-batch generalization gap persists under a constant step budget after careful learning-rate schedule tuning.Prior studies focused largely on convergence, often for convex losses, and did not resolve whether gradient noise promotes generalization.Source: §1 · pp. 1–2
  • The SDE perspective predicts distinct noise-dominated and curvature-dominated regimes, which experiments verify across models under a constant epoch budget.The authors use this perspective to make predictions about test-set behavior.Source: §1 · p. 2
  • 94.9% test accuracy at batch size 2048 versus 92.5% at batch size 16384 after 9725 updates; smaller and moderately large batches outperform very large batches in some models.The comparison holds after learning-rate grid searches and under the same number of iterations, even though large batches achieve smaller training losses.Source: §1 · p. 2
  • For a 16-4 Wide-ResNet on CIFAR-10 at batch size 64, a 128-fold increase in epoch budget reduces the test-optimal learning rate only twofold, versus 16-fold for the training-loss-optimal rate.The authors explain the slow decay of the test-optimal rate through an optimal early-training temperature.Source: §1 · p. 2

2. Preliminaries of Empirical Analysis of SGD

The preliminaries define the minibatch SGD update and explain how learning-rate tuning and compute-budget choices shape the comparisons. The study examines three common budget settings.

2. Preliminaries of Empirical Analysis of SGD

  • Each minibatch update uses batch size B, learning rate ϵ_i, and the loss of the examples in that batch; batches are randomly reshuffled without replacement.With B=N, the update becomes full-batch gradient descent.Source: §2 · p. 2
  • The experiments sweep batch sizes and tune the initial learning rate separately for each batch size using a shared decay schedule.This single-free-parameter schedule was selected to perform well across the architectures and datasets considered.Source: §2 · p. 2
  • The analysis compares constant epoch, constant step, and unlimited compute budgets, which respectively fix epochs, updates, or training until a target is reached.The corresponding computational cost varies with batch size for the constant-step budget.Source: §2 · p. 2

3. A Stochastic Differential Equation Perspective on the Two Regimes of SGD

The SDE perspective organizes SGD behavior into noise-dominated and curvature-dominated regimes, with different constraints on learning rates and different implications for generalization. The paper uses this framework to reason about batch size, schedules, and compute budgets.

3.1. Full batch gradients

  • For full-batch gradients, the optimal initial learning rate is just below ϵcrit=2/λmax, the stability limit set by the largest Hessian eigenvalue.Learning-rate decay can later reduce loss along high-curvature directions.Source: §3.1 · p. 3

3.2. Minibatch gradients

  • Large batches with finite training noise are curvature dominated, while small batches are noise dominated and have learning rates controlled by update noise.In the curvature-dominated regime, the optimal initial learning rate is expected to be determined by the critical learning rate.Source: §3.2 · p. 3
  • When the SDE approximation holds, equal temperatures and epoch budgets predict similar performance, and learning rate typically scales linearly with batch size in the noise-dominated regime.This scaling may fail when noise is long-tailed or other assumptions are violated; batch-normalization statistics must also be computed over a fixed number of examples for the rule to hold.Source: §3.2 · pp. 3–4
  • As batch size rises, the increasing optimal learning rate can violate the critical-rate assumption, producing a transition to curvature domination that experiments often find sharp.The theory allows a transition phase at the boundary between the regimes.Source: §3.2 · p. 3

3.3. Consequences of the two regimes

  • Momentum is expected to match vanilla SGD in the noise-dominated regime but outperform it in the curvature-dominated regime; warmup can support larger stable learning rates when the critical rate rises early.The effective learning rate with momentum is ϵ/(1−m).Source: §3.3 · p. 4

3.4. On learning rate schedules and compute budgets

  • For practical schedules, the initial training phase generally determines whether a run is treated as noise- or curvature-dominated; very small compute budgets may remain curvature dominated at every batch size.Carefully tuned schedules could include both regimes, but the authors say such schedules are usually not identifiable within realistic computation budgets.Source: §3.4 · p. 4

3.5. The generalization benefit of noise

  • To test whether noise benefits generalization, the framework recommends comparing small batches with very large batches in the curvature-dominated regime, where constant temperature cannot be maintained below the critical rate.Convergence bounds predict smaller training losses with larger batches at equal step counts, while test accuracy may fall if noise benefits generalization.Source: §3.5 · p. 4
  • The authors conjecture that an optimal temperature is independent of epoch budget, implying that the test-optimal learning rate for a fixed batch need not decay as budgets grow.They add that noise benefits primarily arise early, while later learning-rate decay can improve test accuracy.Source: §3.5 · p. 4

4. Experimental Setup

The experiments use a shared learning-rate schedule and tune learning rates across batch sizes and compute budgets. The setup evaluates test accuracy and final training loss while extending the main-text experiments with additional architectures in the appendices.

4. Experimental Setup

  • The shared schedule holds the learning rate constant for half the epochs, then reduces it by γ every Nepochs/20 epochs; γ is usually fixed at 2.The schedule is used across architectures, and the initial learning rate is tuned.Source: §4 · p. 5
  • The main text reports Wide-ResNets on CIFAR-10, while appendices add ResNet-50, LSTMs, and autoencoders under the same schedule.Figure 1 specifically uses a 16-4 Wide-ResNet with ghost batch normalization on CIFAR-10 for 200 epochs.Source: §4 · p. 5, Figure 1 · p. 5
  • For each batch size, the Wide-ResNet is trained 15 times across a logarithmic learning-rate grid, and optimal test accuracy is the best mean across the best 12 runs.The procedure also reports mean final training loss at the optimal learning rate and uses error bars to characterize the optimal-rate range.Source: §4 · p. 5

5. SGD under a Constant Epoch Budget

Under a fixed epoch budget, test accuracy and the optimal effective learning rate change with batch size, with SGD behavior differing between small- and large-batch regimes. The observed results match convergence-bound predictions and the SDE analogy.

5. SGD under a Constant Epoch Budget

  • The batch-size sweep trains the 16-4 Wide-ResNet on CIFAR-10 for 200 epochs, matching the original papers’ epoch budget.The experiments use both SGD and SGD with Momentum, with and without batch normalization; without normalization, they use Regularized SkipInit for stability.Source: §5 · pp. 5–6
  • Test accuracy is batch-size independent when batches are small, but falls once batch size becomes sufficiently large; Momentum performs better than plain SGD at large batch sizes.For plain SGD, the optimal learning rate scales with batch size below B ≈512 and is constant beyond; Momentum can scale to larger effective rates when B ≳512.Source: §5 · p. 5
  • The observed batch-size results match predictions from convergence bounds and the SDE analogy.Source: §5 · pp. 5–6

6. SGD under a Constant Step Budget

Under a fixed number of updates, optimal test accuracy first improves with batch size and then declines at very large sizes, even as training loss generally falls. The experiments support a generalization benefit from stochastic gradient noise across several settings.

6. SGD under a Constant Step Budget

  • Under a constant 9765-update SGD-with-Momentum budget on CIFAR-10, optimal test accuracy initially rises with batch size then falls sharply, from 94.9% at 2048 to 92.5% at 16384.Final training loss falls as batch size increases except at 32768, where it rises but remains below the loss at 2048, despite test accuracy being 5% lower.Source: §6 · p. 6
  • Under a constant step budget, the test-accuracy benefit of increasing batch size is non-monotonic: accuracy initially improves, then degrades for very large batches.The fixed update count gives large batches equal opportunity to minimize training loss; the reported schedule holds the rate for 4882 updates, then halves it every 488 updates.Source: §6 · p. 6
  • The large-batch test-accuracy drop also appears across CIFAR-100, autoencoder, and LSTM experiments, with and without batch normalization in the reported settings.Table 1 keeps ghost batch size at 64 so the noise from estimating batch statistics on a subset of training data does not change as overall batch size rises.Source: §6 · p. 6
  • The results clarify that constant-step test accuracy can initially improve with batch size yet decline for very large batches, contrary to a simple small-versus-large batch distinction.Shallue et al. (2018) argued that small batches perform better under constant epoch budgets and large batches under constant step budgets; they also reported that explicit regularization can reduce the gap in their setup.Source: §6 · p. 6

7. SGD with an Unlimited Epoch Budget

With batch size fixed at 64, the paper studies how optimal test performance and learning rates change as epoch budgets grow. It finds that test performance saturates or declines while training loss keeps falling, and independently tuning initial and final rates supports the same schedule conclusions.

7. SGD with an Unlimited Epoch Budget

  • In figure 2, as the epoch budget increases, optimal test accuracy initially rises and then saturates or begins to fall, while optimal training loss falls monotonically.Source: §7 · p. 7
  • The test-accuracy-maximizing learning rate changes little as the epoch budget grows, while the training-loss-minimizing rate falls rapidly.With batch normalization, the test-optimal rate falls by only a factor of 2 from 50 to 6400 epochs; without it, that rate is constant across the tested budgets.Source: §7 · p. 7

7.1. Checking the robustness of our conclusions

  • The robustness experiment breaks the usual coupling between initial and final learning rates by tuning them independently across epoch budgets.It uses the 16-4 Wide-ResNet with batch normalization, batch size 64, and SGD with Momentum, over budgets from 50 to 6400 epochs.Source: §7.1 · pp. 7–8
  • With independently tuned rates, the test-optimal initial learning rate is constant beyond 400 epochs, while the test-optimal final rate decays rapidly as the epoch budget rises.Figure 3 compares initial and final rates for maximizing test accuracy and minimizing training loss; its accuracy saturates beyond 800 epochs while training loss declines monotonically.Source: Figure 3 · p. 8, §7.1 · p. 8
  • The independently tuned schedule results are limited to one model, batch size, and normalization setting because the experiments require a very large compute budget.Source: §7.1 · p. 8
  • The test-optimal initial rate exceeds the loss-optimal initial rate, while the test-optimal final rate is lower than the loss-optimal final rate.The authors relate this pattern to maintaining high temperature early and decaying the learning rate rapidly late; they suggest rapid late decay helps preserve the influence of early noise.Source: §7.1 · p. 8

8. Discussion

The paper argues that gradient noise benefits generalization, with smaller batches outperforming very large batches under both constant-epoch and constant-step budgets after careful tuning. It also derives when linear learning-rate scaling should hold and gives practical tuning guidance, while limiting its conclusions to the schedules and architectures studied.

8. Discussion

  • Smaller batches can outperform very large batches on test sets under constant-epoch and constant-step budgets, even after careful hyperparameter tuning.With unlimited compute, the authors find evidence for an optimal early-training temperature, defined in most models by the learning-rate-to-batch-size ratio.Source: §8 · p. 8
  • The conclusions apply only to the tested learning-rate schedule and architectures, though the authors designed the schedule to resemble popular practitioner schedules.They note that enough compute might allow equally effective schedules that do not follow their main claims.Source: §8 · p. 8
  • With limited compute, the authors recommend tuning a schedule parameterized only by the initial learning rate, then using a small-epoch logarithmic sweep before increasing the epoch budget.With additional resources, tune the final learning rate too; the authors say it usually has less influence on test accuracy.Source: §8 · pp. 8–9
  • The authors suggest gradient noise may help explain why optimization methods designed for the noise-dominated regime can be difficult to make effective on test accuracy.They contrast typical batch sizes of 32–128, often noise-dominated, with large-batch settings where curvature-focused algorithms are more likely to help.Source: §8 · p. 9

A. Deriving the linear scaling rule for small batch sizes

  • When ϵ ≪ ϵcrit, n SGD updates at temperature T and learning rate ϵ are equivalent to one step at T with learning rate nϵ.Since T = ϵ/B, doubling both learning rate and batch size should produce the same parameter distribution after the same number of epochs.Source: A. Deriving the linear scaling rule… · p. 11
  • The multi-step Gaussian-noise approximation requires independent, short-tailed example gradients and, under the central limit argument, N ≫ 1, nB ≫ 1, and nB ≪ N.The derivation needs the aggregate noise over n updates to be Gaussian; it does not require the single-step noise itself to be Gaussian.Source: A. Deriving the linear scaling rule… · p. 11

B. Additional experimental details

The appendix specifies a shared learning-rate decay schedule and describes additional architectures and datasets used in experiments. Its figures also report learning-rate sweeps and batch-size trends for Wide-ResNets on CIFAR-10.

B.1. Our learning rate decay schedule

  • The default schedule holds the learning rate constant for half the epoch budget, then reduces it by γ every Nepochs/20, with γ = 2 unless stated otherwise.In constant-step experiments, epoch budgets scale with batch size so all batch sizes decay after the same number of steps.Source: §B.1 · p. 11, Figure 4 · p. 11

B.2. Additional models used

  • Additional experiments use ResNet-50 on ImageNet, 28-10 Wide-ResNet on CIFAR-100, LSTMs on Penn TreeBank, and autoencoders on MNIST.The ResNet-50 runs use the default schedule without warmup; the LSTM uses two 650-unit layers, and the autoencoder has layer widths 784–1000–500–250–30–250–500–1000–784.Source: §B.2 · pp. 11–12
  • For batch sizes 64 and 1024, SGD and Momentum have similar final performance at small learning rates; at large batch size, Momentum performs better.The figure reports final test accuracy and training loss for a batch-normalized 16-4 Wide-ResNet on CIFAR-10 after 200 epochs.Source: Figure 5 · p. 12
  • Without batch normalization, test accuracy and training loss begin to worsen beyond small-batch thresholds, while the optimal effective learning rate stops scaling with batch size.The reported thresholds for test accuracy are B ≳ 8 for SGD and B ≳ 256 for Momentum; the figure uses Regularized SkipInit and averages the best 12 of 15 runs.Source: Figure 6 · p. 12
  • The 28-10 Wide-ResNet on CIFAR-100 uses the default schedule, which reaches the test accuracy reported by Zagoruyko & Komodakis (2016).Source: §B.2 · p. 12

C. Additional results under constant epoch budgets

Across the additional constant-epoch results, the authors report a transition from a noise-dominated small-batch regime to a curvature-dominated large-batch regime, where performance degrades. Results across model settings show how this transition relates to learning-rate scaling and Momentum.

C. Additional results under constant epoch budgets

  • On ImageNet, ResNet-50 test accuracies match for SGD and Momentum at small batch sizes, while Momentum outperforms SGD at large batch sizes.For Momentum, the optimal effective learning rate scales with batch size across the tested range; for vanilla SGD, linear scaling breaks at large batches.Source: Table 2 · p. 13
  • For the PTB word-level LSTM, test perplexity begins rising above batch size 128 with Momentum and above 64 with vanilla SGD.The optimal effective learning rate scales with the square root of batch size when batches are small, then levels off; the authors note consecutive minibatch gradients are not independent.Source: Figure 7 · p. 13
  • In the noise-dominated regime, training loss and optimal test accuracy remain independent of batch size; both begin to degrade in the curvature-dominated regime.The transition also marks a change from learning rates that increase with batch size to rates that are independent of batch size; Momentum does not outperform SGD in the small-batch regime but does in the large-batch regime.Source: C. Additional results under constant… · p. 13

C.1. Learning rate sweep with batch normalization for two batch sizes on CIFAR-10 Wide-ResNet model

  • With batch normalization, SGD and Momentum have similar performance at small learning rates, but Momentum performs better at large learning rates and large batch sizes.For the CIFAR-10 16-4 Wide-ResNet, the two methods have the same optimal test accuracy at small batch size, while Momentum outperforms vanilla SGD at large batch size.Source: §C.1 · p. 13

C. Additional results under constant epoch budgets

  • For the MNIST autoencoder, test MSE begins rising above batch size 128 with Momentum and above 16 with vanilla SGD.The optimal effective learning rate scales with batch size for small batches, then becomes batch-size independent; it is larger in the curvature-dominated regime for Momentum.Source: Figure 8 · p. 14

C.2. Results without batch normalization on CIFAR-10 Wide-ResNet

  • Without batch normalization, performance is batch-size independent only up to B ≲ 8 for SGD and B ≲ 256 for Momentum; above those thresholds, both methods degrade rapidly.Momentum's larger maximum stable learning rate enables it to scale to larger batch sizes.Source: §C.2 · p. 14

C.3. Results from additional models

Across the additional models, batch-size effects and the relative performance of SGD and Momentum vary by architecture. The LSTM and autoencoder results also show distinct learning-rate scaling and batch-size sensitivity.

ResNet-50 on ImageNet

  • On ImageNet ResNet-50, SGD and Momentum have similar test accuracy at small batch sizes, but Momentum outperforms vanilla SGD at large batch sizes.With Momentum, the optimal effective learning rate is proportional to batch size across the tested range; vanilla SGD does not show this scaling.Source: §C.3 · p. 14

LSTM on PTB

  • On the PTB LSTM, performance begins degrading above batch size 64 for SGD and above 128 for SGD with Momentum.At small batch sizes, the optimal learning rate rises roughly with the square root of batch size before leveling off; correlated consecutive samples are offered as a possible explanation.Source: §C.3 · p. 14

Autoencoder on MNIST

  • On the MNIST autoencoder, performance is batch-size independent at small batches but degrades at smaller batch sizes than in normalized residual networks, with Momentum outperforming SGD sooner.The authors suggest the model's bottleneck structure and resulting poor conditioning may explain this pattern.Source: §C.3 · p. 14

D. Additional results under constant step budgets

Under constant step budgets, the experiments compare test performance and training losses across batch sizes in several architectures. Very large batches often reduce test performance even when training losses fall or remain comparatively low.

Wide-ResNet on CIFAR-10

  • For the 16-4 Wide-ResNet on CIFAR-10, test accuracy falls for very large batches although batch size 2048 still has lower training loss than batch size 128, where test accuracy was maximized.The comparison uses 156,250 updates and a learning-rate grid search.Source: Table 3 · p. 15, D. Additional results under constant… · p. 15

Wide-ResNet on CIFAR-100

LSTM on PTB

Autoencoder on MNIST

  • For the MNIST autoencoder, test-set MSE increases slightly as batch size rises while training-set MSE falls, with the authors describing the generalization benefit of noise as smaller than for the Wide-ResNet or LSTM.At batch size 4096, training loss is similar to that at batch size 64, but test-set MSE is worse; the optimal effective learning rate is independent of batch size.Source: D. Additional results under constant… · p. 15

E. Additional results with a fixed batch size and variable epoch budget

With a fixed batch size and varying epoch budgets, the LSTM and autoencoder show different test-loss and optimal-learning-rate patterns. The autoencoder results also point to early stopping as more influential than stochastic gradient noise in that setting.

Autoencoder on MNIST

  • For the MNIST autoencoder at batch size 32, test MSE initially falls and then rises as the epoch budget increases, while training MSE falls monotonically.Unlike the LSTM result, the test-MSE-minimizing learning rate decreases with the epoch budget while the training-MSE-minimizing rate remains roughly constant.Source: E. Additional results with a fixed… · p. 16, Figure 10 · p. 17
  • In the autoencoder, a larger learning rate can cause faster overfitting before the first decay at 400 epochs, making early stopping more influential on test performance than stochastic gradient noise.The authors report this behavior consistently for epoch budgets over 200 epochs.Source: E. Additional results with a fixed… · p. 16, Figure 10 · p. 17

LSTM on PTB

  • For the PTB LSTM at batch size 64, test perplexity initially falls but rises beyond 56 epochs, while training perplexity keeps falling and the test-optimal learning rate varies little once budgets are not too small.The training-loss-minimizing learning rate falls as the epoch budget rises.Source: E. Additional results with a fixed… · p. 16, Figure 9 · p. 17
Loading 2006.15081v1…