PaperlayerView exact revision on arXiv

Source-linked AI summary

A Bayesian Perspective on Generalization and Stochastic Gradient Descent

Samuel L. Smith, Quoc V. Le

arXiv:1710.06451v3cs.LGcs.AIstat.ML

TL;DR

The paper asks why models that memorize random labels can still generalize on informative labels, and why SGD finds minima that generalize well. It uses Bayesian evidence to explain both phenomena and analyzes SGD through its stochastic dynamics. The paper reports an optimal batch size whose scaling follows B_opt ∝ εN, while noting practical limitations for deep-network evidence evaluation and Laplace approximations.

  • Problem

    The paper addresses how to predict whether a minimum will generalize and why SGD finds well-generalizing minima despite models’ ability to memorize randomized labels.

  • Method

    The paper evaluates Bayesian evidence and interprets SGD as a stochastic differential equation governed by a mini-batch noise scale.

  • Results

    Linear models reproduce random-label memorization, Bayesian evidence penalizes sharp minima, and an optimal batch size scales as B_opt ∝ εN.

  • Takeaways & Limitations

    Mini-batch noise can drive SGD away from sharp minima, linking batch-size selection to evidence and test accuracy.

  • Takeaways & Limitations

    Evidence is difficult to evaluate for deep networks, and the Laplace approximation can change approximate evidence under reparameterization.

Abstract

from arXiv · show

We consider two questions at the heart of machine learning; how can we predict if a minimum will generalize to the test set, and why does stochastic gradient descent find minima that generalize well? Our work responds to Zhang et al. (2016), who showed deep neural networks can easily memorize randomly labeled training data, despite generalizing well on real labels of the same inputs. We show that the same phenomenon occurs in small linear models. These observations are explained by the Bayesian evidence, which penalizes sharp minima but is invariant to model parameterization. We also demonstrate that, when one holds the learning rate fixed, there is an optimum batch size which maximizes the test set accuracy. We propose that the noise introduced by small mini-batches drives the parameters towards minima whose evidence is large. Interpreting stochastic gradient descent as a stochastic differential equation, we identify the "noise scale" $g = ε(\frac{N}{B} - 1) \approx εN/B$, where $ε$ is the learning rate, $N$ the training set size and $B$ the batch size. Consequently the optimum batch size is proportional to both the learning rate and the size of the training set, $B_{opt} \propto εN$. We verify these predictions empirically.

1 INTRODUCTION

The paper argues that Bayesian principles explain why models can memorize random labels yet generalize on informative labels, and why mini-batch SGD favors well-generalizing minima.

  • Small over-parameterized linear models, like deep networks, can generalize on informative labels while memorizing randomized labels.This shows the phenomenon is not unique to deep learning.
  • Bayesian evidence explains these observations by penalizing sharp minima while remaining invariant to model parameterization.The paper contrasts evidence with curvature-based interpretations of generalization.
  • Mini-batch noise drives SGD toward minima with larger evidence, helping explain why smaller batches can generalize better.The proposed mechanism links stochastic-gradient noise to the geometry and evidence of selected minima.
  • The SGD noise scale is g ≈ εN/B, with ε the learning rate, N the training-set size, and B the batch size.This scale motivates the predicted relationship between optimization hyperparameters and generalization.
  • An optimal batch size maximizes test accuracy and is proportional to the learning rate and training-set size.The introduction also notes an equivalent optimal-learning-rate scaling relationship.

2 BAYESIAN MODEL COMPARISON

The paper develops Bayesian model comparison around posterior inference, evidence, and the Occam factor. Evidence combines fit and curvature near a minimum while preserving parameterization invariance, though practical evaluation is difficult for deep networks.

  • Bayesian model comparison infers a posterior over parameters from inputs, labels, and a model using Bayes’ theorem.The setup begins with a classification model and a parameter posterior conditioned on training data.
  • A Gaussian prior yields an L2-regularized cross-entropy cost function whose minimizer is the posterior mode.The regularization coefficient controls the prior’s contribution to the cost.
  • Model comparison uses the evidence ratio, which measures how training data changes prior beliefs, alongside a prior ratio describing model plausibility.The prior ratio is usually set to 1 to avoid unnecessary subjectivity.
  • The optimal batch size is proportional to the learning rate and training-set size.Equivalently, the optimal learning rate is proportional to batch size and training-set size.
  • The Laplace approximation estimates evidence from the cost at the minimum and the logarithm of curvature relative to the regularization constant.The approximation is obtained by Taylor expanding the cost near the minimum.
  • The Occam factor penalizes sharp minima because low-curvature minima require less parameter fine-tuning to fit the data.It represents the fraction of prior parameter space consistent with the data.
  • Evidence is invariant to model parameterization, unlike raw curvature, because it integrates out the parameters.The paper notes that the Laplace approximation can nevertheless lose this invariance under some reparameterizations.

3 BAYES THEOREM AND GENERALIZATION

Logistic regression reproduces the random-label memorization phenomenon, while Bayesian evidence tracks when informative-label models generalize. The evidence penalizes sharp minima and remains invariant to parameterization.

  • Weak regularization lets logistic regression generalize on informative labels while perfectly memorizing random labels.
  • When weakly regularized, the mean training margin is roughly 50% larger for informative than random labels.
  • For random labels, weak regularization yields vanishing training cross-entropy but exploding test cross-entropy and strongly positive evidence ratios.
  • For informative labels, increasing regularization lowers test cross-entropy below ln 2, while the evidence ratio becomes negative before returning toward zero.
  • Test cross-entropy and Bayesian evidence are strongly correlated, with minima at the same regularization strength.
  • Bayesian comparison is quantitatively consistent with the linear-model results and qualitatively consistent with corresponding deep-network observations.

4 BAYES THEOREM AND STOCHASTIC GRADIENT DESCENT

The paper links SGD’s batch-size generalization gap to a balance between gradient-driven depth and noise-driven breadth. Experiments show an optimal batch size and its predicted dependence on learning rate.

  • Gradient noise is proposed to drive SGD toward broad, high-evidence minima, balancing the gradient’s drive toward deep minima.
  • With a constant learning rate, small batches eventually achieve higher test accuracy than large batches, although they converge more slowly.
  • The accuracy peak shifts right as learning rate ϵ increases, and the best observed batch size is proportional to ϵ across two orders of magnitude.
  • L2 regularization significantly shrinks the generalization gap between small- and large-batch training.
  • Test accuracy peaks at an intermediate batch size, while models with batch sizes B ⪅10 cannot train in the reported setting.

5 STOCHASTIC DIFFERENTIAL EQUATIONS AND THE SCALING RULES

The paper models SGD as a stochastic differential equation whose noise scale determines the batch-size tradeoff, deriving scaling rules for learning rate, dataset size, and momentum. Experiments support these predictions, including an optimal batch size that shifts with each parameter.

  • Optimal batch size: An optimal batch size maximizes test accuracy when the other SGD hyperparameters remain fixed.The peak is attributed to a tradeoff between depth and breadth in Bayesian evidence.
  • Noise scale: The SGD noise scale is g ≈ εN/B, so increasing batch size reduces the random fluctuations driving the dynamics.The derivation models mini-batch gradient error as Gaussian noise and matches its variance to a stochastic differential equation.
  • Learning rate and dataset size: B_opt ∝ εN: increasing the learning rate or training-set size shifts the accuracy peak toward larger batches.The learning-rate experiment shows the peak moving right as ε increases, while the training-set experiment finds a linear trend with N.
  • Practical scaling: Coupling learning rate and batch size allows higher learning rates without loss of test accuracy or increased computational cost.The simultaneous increase also permits greater parallelism across multiple GPUs.
  • Momentum: Increasing momentum shifts the optimal batch size right according to B_opt ∝ 1/(1 − m), with remarkably good empirical agreement.The momentum experiments compare test performance across batch sizes and fit the observed best batch sizes to this rule.

6 CONCLUSIONS

The paper argues that Bayesian evidence explains generalization behavior and SGD’s batch-size effects across linear and neural models. It derives noise-based scaling rules and reports empirical verification.

  • 6 CONCLUSIONS: Linear models can generalize on informative labels while memorizing random labels, as deep neural networks do.The paper explains both behaviors using Bayesian evidence, including an Occam factor that penalizes sharp minima while remaining parameterization-invariant.
  • 6 CONCLUSIONS: Mini-batch noise drives SGD away from sharp minima, producing an optimal batch size that maximizes test accuracy.The paper connects this behavior to Bayesian evidence and the preference for minima with larger evidence.
  • 6 CONCLUSIONS: The predicted scaling rules are B_opt ∝ εN and B_opt ∝ 1/(1 − m), and the paper verifies them empirically.The rules relate optimal batch size to learning rate, training-set size, and momentum coefficient.

A BAYESIAN POSTERIOR SAMPLING AND LANGEVIN DYNAMICS

The paper relates Bayesian posterior sampling to Langevin dynamics and uses this connection to explain why minima with larger evidence are sampled more often. In this framework, SGD resembles noisy posterior exploration under specific covariance assumptions.

  • A BAYESIAN POSTERIOR SAMPLING AND LANGEVIN DYNAMICS: Bayesian posterior sampling is presented as an alternative to simply minimizing the cost function.The posterior assigns probability to parameter values rather than selecting only the cost-minimizing point.
  • A BAYESIAN POSTERIOR SAMPLING AND LANGEVIN DYNAMICS: Overdamped Langevin dynamics can sample the Bayesian posterior by combining cost gradients with Gaussian noise.At temperature T = 1, sufficiently small steps and enough iterations produce valid posterior samples.
  • A BAYESIAN POSTERIOR SAMPLING AND LANGEVIN DYNAMICS: SGD resembles overdamped Langevin dynamics when gradient covariances are stationary and proportional to the identity, with temperature proportional to its noise scale.The paper notes that its SGD analysis does not require this stationarity assumption throughout training.
  • A BAYESIAN POSTERIOR SAMPLING AND LANGEVIN DYNAMICS: The probability of sampling a local minimum is proportional to the posterior integral over its basin, which corresponds to its local evidence.The evidence is defined through integrals over parameter space or local attraction bowls.

B THE EFFECT OF REGULARIZATION ON THE GENERALIZATION GAP

Regularization changes the comparison between small-batch and full-batch training. In the reported experiments, it slows full-batch convergence, reduces the generalization gap, and yields similar long-time performance.

  • B THE EFFECT OF REGULARIZATION ON THE GENERALIZATION GAP: λ = 0.1 regularization is used for the training-curve comparison, whereas the main section 4 experiments use λ = 0.The regularized experiments otherwise match figure 3.
  • B THE EFFECT OF REGULARIZATION ON THE GENERALIZATION GAP: Full-batch regularized training takes longer to converge than small-batch training but achieves similar long-time performance.The comparison concerns mean test accuracy and mean test cross-entropy during training.
  • B THE EFFECT OF REGULARIZATION ON THE GENERALIZATION GAP: Small-batch training exhibits fluctuating performance because of its inherent noise.Figure 8 reports this behavior for both mean test accuracy and mean test cross-entropy.
  • B THE EFFECT OF REGULARIZATION ON THE GENERALIZATION GAP: Regularization significantly reduces the generalization gap between small-batch and large-batch training.Large-batch regularized training has slightly lower test accuracy than unregularized small-batch training but lower test cross-entropy.

C THE GAUSSIAN APPROXIMATION TO THE MINI-BATCH ERROR

The paper approximates mini-batch gradient error as Gaussian noise, using a central-limit-theorem rationale and an empirical proof-of-principle examination of gradient distributions.

  • The mini-batch gradient error is approximated by a Gaussian random variable, enabling derivation of scaling rules later verified empirically.
  • The Gaussian approximation is motivated by the central limit theorem when N and B grow while B remains much smaller than N.This requires that gradients over individual training examples do not have heavy tails.
  • The paper cautions that practical training violates the approximation’s asymptotic conditions because N and B are finite and gradients may be heavy tailed.
  • A proof-of-principle experiment presents the gradient distribution immediately after random initialization for a shallow neural network.The distribution is taken over individual training examples for one randomly chosen softmax output-layer matrix element.

D DERIVING THE SCALING RULES FOR SGD WITH MOMENTUM

The paper derives SGD-with-momentum scaling by matching a generalized Langevin description to momentum updates and equating their noise statistics. The resulting rules link batch size, learning rate, momentum, and fluctuation scale.

  • Momentum is modeled as a generalized Langevin equation with structured fluctuations, damping, and parameter-dependent Gaussian noise.
  • The noise scale g controls random fluctuations, while F(ω) describes gradient covariances between parameters.The noise variance includes the damping coefficient λ to satisfy the fluctuation-dissipation theorem.
  • The Langevin dynamics are rewritten as two coupled first-order differential equations so they can be compared with momentum updates.
  • The mini-batch gradient error is approximated as Gaussian with zero mean and covariance approximately N^2F(ω)/B.
  • Keeping random-fluctuation scale constant requires B ∝ εN and additionally predicts B ∝ 1/(1 − m).The effective learning rate can be interpreted as ε_eff = ε/(1 − m).

E HOW TO ACHIEVE LARGE BATCH TRAINING

The paper proposes a heuristic for tuning batch size, learning rate, and momentum to obtain large-batch training while maintaining validation accuracy. It presents this as worthwhile mainly when models are retrained repeatedly.

  • The proposed heuristic seeks to maximize test accuracy and batch size, enabling parallel training across many machines.The paper notes that this objective is worthwhile if a model is expected to be retrained many times.
  • The procedure begins with learning rate 0.1 and momentum coefficient 0.9, then searches logarithmically over batch sizes for the best validation accuracy.If training is unstable, the learning rate is reduced and the search repeated.
  • The heuristic repeatedly triples batch size while scaling ε ∝ B until validation accuracy starts to decline.
  • The authors believe the procedure can increase test accuracy, reduce hyper-parameter tuning cost, and reduce the final number of gradient updates.
Loading 1710.06451v3…