Source-linked AI summary
Three Factors Influencing Minima in SGD
Stanisław Jastrzębski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, Amos Storkey
TL;DR
The paper asks how learning rate, batch size, and gradient covariance shape the minima and generalization behavior reached by SGD in DNNs. Using an SDE approximation, it finds that learning-rate-to-batch-size ratio governs key dynamics and minimum width, with experimental support for wider minima, better generalization, and batch-size schedule replacements. The analysis relies on restrictive assumptions that limit direct transfer to practical DNN training.
Problem
The paper addresses the open question of how learning rate, batch size, and gradient covariance relate to the properties of SGD’s final minima, including width and generalization.
Method
The paper approximates SGD as an SDE and analyzes how its learning-rate-to-batch-size ratio and gradient covariance relate to optimization dynamics and minima.
Results
Higher learning-rate-to-batch-size ratios correlate with wider minima and better generalization, while jointly rescaling learning rate and batch size approximately preserves dynamics; batch-size schedules can replace learning-rate schedules.
Takeaways & Limitations
The learning-rate-to-batch-size ratio is a key determinant of SGD dynamics and minimum width, and changing batch size offers an alternative way to vary training noise.
Takeaways & Limitations
The theoretical assumptions are too restrictive for the analysis to transfer directly to practical DNN training.
Abstract
from arXiv · showhide
We investigate the dynamical and convergent properties of stochastic gradient descent (SGD) applied to Deep Neural Networks (DNNs). Characterizing the relation between learning rate, batch size and the properties of the final minima, such as width or generalization, remains an open question. In order to tackle this problem we investigate the previously proposed approximation of SGD by a stochastic differential equation (SDE). We theoretically argue that three factors - learning rate, batch size and gradient covariance - influence the minima found by SGD. In particular we find that the ratio of learning rate to batch size is a key determinant of SGD dynamics and of the width of the final minima, and that higher values of the ratio lead to wider minima and often better generalization. We confirm these findings experimentally. Further, we include experiments which show that learning rate schedules can be replaced with batch size schedules and that the ratio of learning rate to batch size is an important factor influencing the memorization process.
1 Introduction
The paper examines why SGD reaches minima with different geometries and generalization properties, identifying learning-rate-to-batch-size ratio as a critical control parameter. It argues and experimentally verifies that this ratio influences SGD dynamics, minimum width, and generalization.
- DNNs generalize well despite massive over-parameterization and near-zero training error, leaving the source of their generalization ability unresolved.
- SGD can converge to regions with different geometries and generalization properties because DNN loss landscapes are non-convex and depend on hyperparameters and initialization.
- Prior work linked wide minima with better generalization and larger batch sizes with sharper minima and worse generalization, although sharp minima can sometimes have similar predictions.
- The paper identifies LR/BS, rather than batch size alone, as SGD’s critical control parameter: constant ratios yield similar behavior, while higher ratios lead to wider minima and often better generalization.
- SGD processes sharing LR/BS are treated as discretizations of the same SDE, and experiments verify similar dynamics under joint rescaling of learning rate and batch size.
2 Theory
The theory models SGD as an SDE whose path properties depend on η/S and gradient covariance, then relates this stochastic noise to minima width and expected loss under local assumptions.
- SGD minibatch gradients are approximately Gaussian with mean g(θ) and covariance Σ(θ) = C(θ)/S under IID sampling and sufficiently large batches.
- SGD is an Euler–Maruyama discretization of an SDE in which learning rate and batch size appear through the ratio η/S.The learning rate additionally controls the discretization stepsize and per-batch speed.
- Rescaling learning rate and batch size by the same factor preserves SGD dynamics on an epoch-time axis.With η′ = aη and S′ = aS, the corresponding epoch numbers satisfy e′ = e.
- The analysis defines minima width through Tr(H), with lower trace indicating a wider minimum, and assumes a locally quadratic zero-loss landscape.It also assumes gradient covariance and the Hessian are approximately equal, C = H.
- Under these assumptions, the η/S ratio determines the trade-off between expected loss and minimum width, with E(L) Tr(H) proportional to η/S.
- In the isotropic-covariance case, increasing η/S spreads the equilibrium distribution and makes minimum width more influential relative to loss depth.For fixed σ, η/S controls the thermodynamic temperature, and higher temperature favors wider minima.
3 Experiments
Experiments test whether SGD dynamics, minima geometry, generalization, cyclic schedules, and memorization depend primarily on the learning-rate-to-batch-size ratio. Across architectures and datasets, matched ratios produce similar dynamics and geometry, while higher ratios generally yield wider minima and better generalization, with scaling breakdown at large factors.
- Geometry and generalization: Higher η/S correlates with lower Hessian curvature, wider minima, and better generalization on Fashion-MNIST MLPs.The reported curvature measures are the largest Hessian eigenvalue and the Hessian Frobenius norm.
- Cyclic schedules: Cyclic batch-size schedules can approximately preserve the practical generalization benefit of cyclic learning-rate schedules by varying SGD noise.The cyclic schedules produced wider bowls and higher loss than the constant-learning-rate baseline in the reported comparison.
- Impact of SGD on memorization: After memorizing random labels, higher SGD noise leads to better validation generalization, while low noise steers optimization toward minima with low generalization ability.This result was reported for MNIST with 25% and 50% randomized labels, with and without momentum.
- Breakdown of η/S scaling: The η/S scaling approximation breaks down when the joint rescaling factor β becomes too large, especially for smaller datasets.The paper attributes this divergence to limitations of the SDE discretization approximation.
4 Related work
Prior work linked batch size and gradient noise to sharpness and generalization, while this paper emphasizes the joint role of learning rate and batch size through their ratio. It argues that equal LR/BS processes share dynamics and that this ratio relates to the width of SGD endpoints.
- Unlike analyses that omit batch size, this paper makes batch size explicit and identifies learning rate divided by batch size as the parameter setting SGD dynamics.
- The analysis concerns fixed-learning-rate SGD rather than using decreasing learning rates to sample from the actual posterior.
- SGD processes with the same LR/BS are different discretizations of one underlying SDE and therefore follow the same dynamics when the approximation is justified.
- The paper extends prior batch-size work by studying both learning rate and batch size, connecting their ratio to endpoint width and experimentally verifying matching learning curves under joint rescaling.
5 Conclusion
The paper concludes that approximating SGD by an SDE clarifies how learning rate, batch size, and gradient covariance shape optimization and final minima. Experiments support approximate LR/BS invariance, wider and often better-generalizing minima at higher ratios, and batch-size schedules as learning-rate alternatives.
- The SDE analysis identifies LR/BS as controlling SGD dynamics by scaling stochastic noise and, under the covariance–Hessian assumption, as a key determinant of minimum width.
- Across DNN models and datasets, higher LR/BS correlated with wider minima and better generalization, while rescaling learning rate and batch size produced approximately invariant behavior.
- Changing batch size during training was experimentally suggested as a viable alternative to changing the learning rate.
A When Covariance is Approximately the Hessian
The appendix analyzes when minibatch gradient covariance can approximate the Hessian. It uses sample-covariance identities and assumptions near a realizable optimum, concluding that covariance approaches the Hessian as the estimator becomes more accurate with larger N.
- The appendix seeks conditions under which the gradient covariance matrix C is approximately the Hessian H.
- The covariance C is approximated by the sample covariance K, using the mean gradient and expected squared gradient norm in the derivation.
- The squared norm of the mean gradient is assumed much smaller than the expected squared norm of the gradient, simplifying the sample-covariance expression.
- Near a realizable optimum, the second covariance term is negligible because the model distribution matches the training-data conditional distribution.
- With larger N making the estimator more accurate, the derivation concludes that the covariance is approximately the Hessian.
- The appendix relates the first covariance term to the Empirical Fisher and notes its connection to Hessian approximation and Natural Gradient near the optimum.