The alignment property of SGD noise and how it helps select flat minima: A stability analysis

Lei WuMingze WangWeijie Su

article2022NeurIPS56 citations

Modern machine learning models are typically over-parameterized, possessing far more parameters than training samples. In this regime, numerous solutions achieve zero training error, yet they often yield vastly different real-world generalization performance. Stochastic gradient descent (SGD)—the standard optimization method used to train these models—consistently succeeds by finding solutions that generalize well, outperforming standard full-batch gradient descent without requiring explicit penalties or constraints. Understanding the underlying mechanism behind this implicit regularization is a critical open problem for optimizing model reliability and performance.

The article establishes a direct mathematical and empirical link between the internal structure of SGD noise and dynamical linear stability. Its primary objective is to demonstrate quantitatively why and how SGD inherently prefers flat minima over sharp ones during training, providing rigorous bounds on the geometric curvature of solutions that SGD can stably maintain.

To evaluate this behavior, the authors developed a theoretical stability framework for over-parameterized models trained with square loss. They analyzed how random mini-batch sampling generates noise that directly mirrors the local curvature of the loss landscape, measuring this relationship via a loss-scaled alignment factor. The theoretical findings were validated through extensive numerical simulations across linear networks, random feature models, fully connected networks, and convolutional architectures, as well as large-scale image classification experiments on the full CIFAR-10 dataset using standard deep neural networks such as VGG and ResNet architectures.

The analysis produced three primary findings. First, SGD noise is inherently geometry-aware: its magnitude scales with training loss, and its fluctuations concentrate predominantly along the sharp directions of the local parameter landscape. Second, any solution that is linearly stable under SGD has its sharpness—measured by the Frobenius norm of the Hessian matrix—strictly bounded by a factor proportional to the square root of the batch size divided by the learning rate. Crucially, this bound is independent of both the model size and the training sample size. Third, if a candidate solution is too sharp to satisfy this stability threshold, SGD escapes from it exponentially fast, typically within about 10 steps in empirical tests, driven purely by the positive feedback between loss and noise magnitude.

These findings provide clear practical implications for machine learning practitioners and engineering teams. They explain why using larger learning rates and smaller batch sizes systematically forces optimization away from sharp, poorly generalizing solutions toward broader, more robust minima. Unlike prior continuous-time differential equation approximations that struggled to capture behavior under large learning rates, this stability framework offers a reliable explanation of modern deep learning dynamics and demonstrates that SGD operates naturally near the edge of stability.

Practitioners should leverage these insights by carefully tuning the ratio between learning rate and batch size as a direct mechanism to control model sharpness and boost generalizability. When training large over-parameterized models, teams can favor smaller batch sizes and larger stable learning rates to accelerate the escape from sharp minima. As a next step, researchers should extend this stability framework beyond local quadratic approximations to study non-local trajectory dynamics and explicitly connect Hessian Frobenius norm bounds to formal generalization error bounds.

A primary limitation of this work is its reliance on local linearization around global minima and quadratic loss functions, which may not fully describe non-local training dynamics or non-convex regimes far from convergence. Nevertheless, the consistent match between theoretical predictions and empirical benchmarks across varied architectures provides high confidence in the core conclusion that noise geometry drives SGD toward size-independent flat minima.

asymptotical stability global stability SGD noise flat minima alignment property

Wu et al (2022).pdf

Paper: On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima, Nitish Shirish Keskar et al. (2016). Provides the foundational empirical evidence that small-batch stochastic gradient descent converges to flat minima while large-batch training leads to sharp, poorly generalizing solutions.
Paper: Train faster, generalize better: Stability of stochastic gradient descent, Moritz Hardt et al. (2015). Introduces the algorithmic stability framework for stochastic gradient descent that links optimization dynamics directly to generalization error bounds.
Paper: Sharpness-Aware Minimization for Efficiently Improving Generalization, Pierre Foret et al. (2020). Establishes the connection between loss landscape curvature, explicit flatness optimization, and improved generalization performance in over-parameterized models.
Paper: Visualizing the Loss Landscape of Neural Nets, Hao Li et al. (2017). Develops methods for analyzing and visualizing neural network loss landscapes and Hessian curvatures across deep architectures.
Paper: Understanding deep learning requires rethinking generalization, Chiyuan Zhang et al. (2017). Highlights the breakdown of classical generalization bounds in over-parameterized neural networks, motivating the need to analyze SGD's implicit regularization.
Paper: Averaging Weights Leads to Wider Optima and Better Generalization, Pavel Izmailov et al. (2018). Demonstrates the practical generalization benefits of finding broader, flatter optima within the parameter loss landscape.
Paper: Exploring Generalization in Deep Learning, Behnam Neyshabur et al. (2017). Examines how sharpness, parameter norms, and optimization dynamics interact to influence generalization in deep learning.
Paper: Optimization Methods for Large-Scale Machine Learning, Léon Bottou et al. (2016). Offers a comprehensive survey of the theoretical properties, convergence dynamics, and inherent noise behaviors of stochastic gradient methods in large-scale settings.

Paper: Stochastic Training is Not Necessary for Generalization, Jonas Geiping et al. (2022). Challenges the indispensability of SGD's stochastic noise for selecting flat, generalizable minima by demonstrating comparable generalization via non-stochastic full-batch training.
Paper: Gradient Norm Aware Minimization Seeks First-Order Flatness and Improves Generalization, Xingxuan Zhang et al. (2023). Extends landscape flatness optimization by targeting gradient norm constraints to better control local curvature and escape sharp minima.
Paper: Sharpness-Aware Training for Free, Jiawei Du et al. (2022). Builds on the principles of seeking flat loss landscapes during training by proposing a computationally efficient method that removes the two-pass overhead of SAM.
Paper: Make Sharpness-Aware Minimization Stronger: A Sparsified Perturbation Approach, Peng Mi et al. (2022). Applies selective parameter perturbation to optimize sharpness along critical curvature dimensions without incurring the full computational cost of dense sharpness-seeking algorithms.

Topics

asymptotical stability global stability

Keywords

SGD noise flat minima alignment property stability analysis stochastic gradient descent implicit regularization loss-scaled alignment noise structure sharp minima

1 Introduction
1.1 Related work
2 The alignment property of SGD noise
2.1 Why does the alignment property hold?
2.2 Empirical validations
3 The linear stability analysis
3.1 The linear stability imposes size-independent flatness constraints
3.2 SGD escapes from sharp minima exponentially fast
3.3 The importance of the noise structure
4 Larger-scale experiments
5 Conclusion
Acknowledgements
References

Knowls

Knowl 1 — Loss-scaled alignment of SGD noise
definition
For scalar-output regression with training examples $(x_i,y_i)_{i=1}^n$ , model $f(x;\theta)$ , per-example squared loss $L_i(\theta)=\tfrac12(f(x_i;\theta)-y_i)^2$ , and empirical loss $L(\theta)=\tfrac1n\sum_{i=1}^n L_i(\theta)$ , mini-batch SGD can be written as
$\theta_{t+1}=\theta_t-\eta\bigl(\nabla L(\theta_t)+\xi_t\bigr), \qquad \mathbb E[\xi_t\mid\theta_t]=0, \qquad \mathbb E[\xi_t\xi_t^\top\mid\theta_t]=\frac{\Sigma(\theta_t)}{B},$
where $\eta>0$ is the learning rate, $B$ is the batch size, and
$\Sigma(\theta)=\frac1n\sum_{i=1}^n\nabla L_i(\theta)\nabla L_i(\theta)^\top-\nabla L(\theta)\nabla L(\theta)^\top$
is the per-example gradient covariance. Let $g_i(\theta)=\nabla f(x_i;\theta)$ and let the Fisher matrix be $G(\theta)=\tfrac1n\sum_i g_i(\theta)g_i(\theta)^\top$ . The paper defines
$\alpha(\theta)=\frac{\operatorname{Tr}(G(\theta)\Sigma(\theta))}{\|G(\theta)\|_F\|\Sigma(\theta)\|_F}, \qquad \beta(\theta)=\frac{\|\Sigma(\theta)\|_F}{2L(\theta)\|G(\theta)\|_F},$
and the loss-scaled alignment factor
$\mu(\theta)=\alpha(\theta)\beta(\theta) =\frac{\operatorname{Tr}(\Sigma(\theta)G(\theta))}{2L(\theta)\|G(\theta)\|_F^2},$
for $L(\theta)>0$ . The factor $\alpha$ measures whether noise is concentrated in Fisher-sharp directions, while $\beta$ measures noise magnitude relative to the loss. Thus, a positive lower bound on $\mu$ means that SGD noise is both sufficiently large and directed toward sharp directions. Under the heuristic approximation $\Sigma(\theta)\approx2L(\theta)G(\theta)$ , all three quantities $\alpha$ , $\beta$ , and $\mu$ equal $1$ . At a global zero-loss minimum of the squared-loss model, the Hessian satisfies $H(\theta^*)=G(\theta^*)$ .
Knowl 2 — Gradient-norm uniformity guarantees noise alignment
theoretical result
Let $g_i(\theta)=\nabla f(x_i;\theta)$ , $G(\theta)=\tfrac1n\sum_i g_i(\theta)g_i(\theta)^\top$ , and define
$\chi_i(\theta)=g_i(\theta)^\top G(\theta)g_i(\theta), \qquad \bar\chi(\theta)=\frac1n\sum_{i=1}^n\chi_i(\theta), \qquad \gamma(\theta)=\min_i\frac{\chi_i(\theta)}{\bar\chi(\theta)}.$
The quantity $\gamma(\theta)$ measures the uniformity of the model-gradient norms in the Fisher metric. Define
$\Sigma_1(\theta)=\frac1n\sum_i\nabla L_i(\theta)\nabla L_i(\theta)^\top, \qquad \Sigma_2(\theta)=\nabla L(\theta)\nabla L(\theta)^\top,$
so that $\Sigma=\Sigma_1-\Sigma_2$ , and define
$\mu_1(\theta)=\frac{\operatorname{Tr}(\Sigma_1(\theta)G(\theta))}{2L(\theta)\|G(\theta)\|_F^2}, \qquad \mu_2(\theta)=\frac{\operatorname{Tr}(\Sigma_2(\theta)G(\theta))}{2L(\theta)\|G(\theta)\|_F^2}.$
For squared loss, the paper proves
$\mu_1(\theta)\ge \gamma(\theta), \qquad \mu(\theta)=\mu_1(\theta)-\mu_2(\theta).$
Consequently, uniformity of the model-gradient norms is sufficient for a positive $\mu_1$ -alignment even when the full-batch gradient contribution $\Sigma_2$ is not negligible. This result identifies gradient-norm uniformity, rather than uniformity of the individual fitting errors, as the structural reason alignment can hold.
Knowl 3 — Exact alignment for over-parameterized linear models
theoretical result
Consider an over-parameterized re-parameterized linear model
$f(x;\theta)=F(\theta)^\top x,$
where $F(\theta)\in\mathbb R^d$ , and suppose the input is Gaussian, $x\sim\mathcal N(0,S)$ , with online SGD corresponding to infinitely many samples. Let $L$ be the population squared loss, $G$ its Fisher matrix, and $\Sigma$ the covariance of the per-example gradient. If
$\Sigma_1=\mathbb E[\nabla L_x(\theta)\nabla L_x(\theta)^\top], \qquad \mu_1=\frac{\operatorname{Tr}(\Sigma_1G)}{2L\|G\|_F^2}, \qquad \mu=\frac{\operatorname{Tr}(\Sigma G)}{2L\|G\|_F^2},$
then the noise covariance has the exact form
$\Sigma(\theta)=\nabla L(\theta)\nabla L(\theta)^\top+2L(\theta)G(\theta),$
and satisfies
$\mu_1(\theta)\ge\mu(\theta)\ge1.$
Thus, for this whole class of re-parameterized linear models, SGD noise is aligned with the Fisher geometry throughout parameter space, and the alignment strength does not depend on the model size. The result is proved for the infinite-sample setting; the paper reports that finite-sample and low-sample behavior is subsequently examined empirically.
Knowl 4 — Alignment in random ReLU feature models
theoretical result
For a feature model $f(x;\theta)=\sum_{j=1}^m\theta_j\phi_j(x)=\langle\theta,\Phi(x)\rangle$ , the model gradients $g_i=\Phi(x_i)$ , Fisher matrix, and Hessian are constant:
$G=H=\frac1n\sum_{i=1}^n g_i g_i^\top.$
With $\chi_i=g_i^\top Gg_i$ , $\gamma=\min_i\chi_i/(n^{-1}\sum_i\chi_i)$ , and
$\tau(G)=\frac{\lambda_1(G)^2}{\sum_j\lambda_j(G)^2},$
where $\lambda_1(G)$ is the largest eigenvalue, the SGD covariance satisfies
$\mu_1(\theta)\ge\gamma, \qquad \mu_2(\theta)\le\tau(G), \qquad \mu(\theta)\ge\gamma-\tau(G).$
For the random ReLU features $\phi_j(x)=\operatorname{ReLU}(w_j^\top x)$ with independent $w_j\sim\operatorname{Unif}(\sqrt d\,S^{d-1})$ and inputs $x\sim\operatorname{Unif}(S^{d-1})$ , the paper proves that for any $\delta\in(0,1)$ , if $m\ge n\gtrsim d^5\log(1/\delta)$ , then with probability at least $1-\delta$ ,
$\mu_1(\theta)\ge1, \qquad \mu(\theta)\gtrsim d^{-1}.$
Here $\gtrsim$ hides an absolute constant. The result supplies a positive alignment guarantee for random feature models and supports applying the same local reasoning to nonlinear networks through their linearization around a minimum.
Knowl 5 — Linearized SGD converts aligned noise into a stability constraint
theoretical result
Let $\theta^*$ be a global minimum of a squared-loss model and use local coordinates with $\theta^*=0$ . The linearized model has quadratic loss
$L(\theta)=\frac12\theta^\top H\theta,$
where $H\succeq0$ is the Hessian at the minimum. Consider a general local SGD recursion
$\theta_{t+1}=\theta_t-\eta\bigl(H\theta_t+\xi_t\bigr), \qquad \mathbb E[\xi_t\mid\theta_t]=0, \qquad \mathbb E[\xi_t\xi_t^\top\mid\theta_t]=S(\theta_t).$
For $\nu(\theta)=\tfrac12\operatorname{Tr}(HS(\theta))$ , the expected one-step loss obeys
$\mathbb E[L(\theta_{t+1})] =\mathbb E\left[r(\theta_t)L(\theta_t)+\eta^2\nu(\theta_t)\right],$
where, for $\theta^\top H\theta>0$ ,
$r(\theta)=1-\frac{2\eta\,\theta^\top H^2\theta}{\theta^\top H\theta} +\frac{\eta^2\,\theta^\top H^3\theta}{\theta^\top H\theta},$
and $r(\theta)\ge0$ ; moreover, $r(\theta)\le1$ whenever $\eta\le2/\lambda_1(H)$ . The first term is the full-batch-gradient contribution and the second is the noise contribution. Because the latter is proportional to $\operatorname{Tr}(HS)$ , noise aligned with sharp eigendirections increases local loss more strongly.
Knowl 6 — Stable SGD minima have size-independent Frobenius flatness
theoretical result
For a global minimum $\theta^*$ of an over-parameterized squared-loss model, consider the linearized mini-batch SGD dynamics with learning rate $\eta$ , batch size $B$ , Hessian $H(\theta^*)$ , and per-example noise covariance $\Sigma(\theta)$ . A minimum is called linearly stable if there is a constant $C>0$ such that the expected loss of the linearized dynamics satisfies $\mathbb E[L(\theta_t)]\le C\mathbb E[L(\theta_0)]$ for every $t\ge0$ . Suppose the loss-scaled alignment factor
$\mu(\theta)=\frac{\operatorname{Tr}(\Sigma(\theta)G(\theta))}{2L(\theta)\|G(\theta)\|_F^2}$
is bounded below by a positive constant $\mu_0$ in the local region relevant to the linearized dynamics. Then linear stability requires
$\|H(\theta^*)\|_F\le\frac1\eta\sqrt{\frac{B}{\mu_0}}.$
Therefore, when $\mu_0$ is independent of model size and sample size, the Frobenius-norm sharpness of minima accessible to SGD is also independent of those quantities. The bound grows as $\sqrt B$ and decreases as $1/\eta$ . For comparison, linear stability of full-batch gradient descent only gives $\lambda_1(H(\theta^*))\le2/\eta$ , implying the weaker Frobenius bound $\|H(\theta^*)\|_F\le2\sqrt p/\eta$ for a model with $p$ parameters.
Knowl 7 — A weaker alignment condition is sufficient for flatness selection
theoretical result
For the same linearized mini-batch SGD setting, define
$\mu_1(\theta)=\frac{\operatorname{Tr}(\Sigma_1(\theta)G(\theta))}{2L(\theta)\|G(\theta)\|_F^2},$
where $\Sigma_1(\theta)=n^{-1}\sum_i\nabla L_i(\theta)\nabla L_i(\theta)^\top$ excludes the negative full-batch-gradient term from the covariance. If $\mu_1(\theta)\ge\mu_1>0$ , then a linearly stable global minimum satisfies
$\|H(\theta^*)\|_F \le \min\left\{ \frac{B}{\sqrt{(B-1)\mu_1}},\frac{2B}{\mu_1} \right\}\frac1\eta.$
For large batch size, the first term behaves as
$\frac{B}{\eta\sqrt{(B-1)\mu_1}} \approx \frac1\eta\sqrt{\frac{B}{\mu_1}},$
matching the scaling obtained from the stronger $\mu$ -alignment condition. Thus, the full-batch gradient need not be negligible for the qualitative conclusion that SGD-accessible minima must be flat; local uniformity of model-gradient norms, which guarantees $\mu_1$ -alignment, is enough.
Knowl 8 — SGD escapes sufficiently sharp minima exponentially fast
theoretical result
Let $\theta^*$ be a global minimum whose linearized mini-batch SGD dynamics have learning rate $\eta$ , batch size $B$ , Hessian $H(\theta^*)$ , and alignment factor satisfying $\mu(\theta)\ge\mu_0>0$ . If the minimum violates the stability threshold,
$\|H(\theta^*)\|_F>\frac1\eta\sqrt{\frac{B}{\mu_0}},$
then the expected loss of the linearized SGD iterates obeys
$\mathbb E[L(\theta_t)]\ge\gamma_0^t\,\mathbb E[L(\theta_0)], \qquad \gamma_0=\frac{\eta^2\mu_0}{B}\|H(\theta^*)\|_F^2>1.$
Consequently, escaping from a region with loss of order $\varepsilon$ to one with loss of order $1$ takes roughly
$\log_{\gamma_0}(1/\varepsilon)$
steps. The escape time depends only logarithmically on the loss barrier and not on the parameter-space dimension. The result is a noise-driven effect: the full-batch gradient can be stable at the sharp minimum, while geometry-aligned SGD noise makes the expected loss grow geometrically.
Knowl 9 — Loss-dependent, geometry-aware noise explains efficient escape
theoretical result
The paper separates two structural properties of SGD noise. First, for the local quadratic model with Hessian $H$ , if the noise covariance $S(\theta)$ satisfies $\eta\le2/\lambda_1(H)$ and
$\mathbb E[\operatorname{Tr}(HS(\theta))]\le2\sigma^2,$
then the expected noise-driven loss increase is at most linear:
$\mathbb E[L(\theta_t)-L(\theta_0)]\le\eta^2\sigma^2 t.$
Thus, the exponential escape result relies on the fact that SGD noise magnitude grows with the loss. Second, compare two noises with equal total-variance normalization for a $p$ -parameter model:
$S_1(\theta)=2L(\theta)H \quad\text{and}\quad S_2(\theta)=2\sigma^2L(\theta)I_p, \qquad \sigma^2=\frac{\operatorname{Tr}(H)}{p}.$
The geometry-aware noise $S_1$ yields the size-independent constraint $\|H\|_F\le\sqrt B/\eta$ , whereas the isotropic noise $S_2$ yields only
$\operatorname{Tr}(H)\le\frac{\sqrt{pB}}\eta.$
Isotropic noise spends most of its variance in flat directions, which barely increases loss, while geometry-aware noise concentrates perturbations in sharp directions. The combination of loss-dependent magnitude and alignment with sharp directions therefore explains both geometric escape and size-independent flatness selection.
Knowl 10 — Empirical validation across model classes and CIFAR-10 networks
empirical result
Experiments examined random feature models, linear networks, fully connected networks, convolutional networks, VGG networks, and ResNets. In the small-scale experiments, the alignment factor $\alpha$ stayed close to $1$ during training, $\beta$ remained bounded away from zero, and therefore $\mu$ stayed positive. Alignment persisted for the linear network even in a low-sample regime with $n=100$ and input dimension $d=50$ . Across increasing over-parameterization, $\mu$ was essentially independent of model size for linear networks and random feature models, and was nearly independent for nonlinear networks; for the CNN experiment it decreased only from approximately $1.05$ to $1.0$ while model size grew by more than two orders of magnitude. The gradient-norm uniformity factor $\gamma$ also remained bounded below along the training trajectories.

The Hessian experiments found that the Frobenius norm $\|H\|_F$ remained nearly unchanged as model size increased, whereas the Hessian trace grew substantially. The theoretical upper bound was valid throughout the observed training trajectories and became tighter for smaller batches. In noise-only escape experiments initialized at minima found by gradient descent, the loss grew exponentially; about ten SGD steps sufficed to reach a high-loss region in each examined model, using averages over 200 trajectories.

For full CIFAR-10 classification without explicit regularization, VGG-19, ResNet-110, and additional VGG and ResNet depths exhibited significantly positive, nearly size-independent alignment factors, exponentially fast escape from sharp minima, and nearly size-independent convergent flatness. The empirical upper bounds were relatively tight for models such as VGG-16 and VGG-19 and became tighter as batch size decreased. Error bars for reported small-scale comparisons were estimated over five runs.
Knowl 11 — Scope limitation of the stability explanation
limitation
The rigorous analysis is local: it studies the linearized dynamics around a global minimum and requires the alignment condition to hold in the corresponding local region. The paper argues that alignment makes perturbations in sharp directions dominate, which supports the relevance of the local quadratic approximation despite degeneracy in flat directions, but it does not provide a full proof that the nonlinear dynamics remain equivalent to the linearized dynamics in all flat directions. Experiments show exponential loss growth even after entering high-loss, nonlocal regions, but the paper does not explain this nonlocal escape behavior. It also leaves open whether the Hessian Frobenius norm directly controls generalization and how stability governs the entire SGD trajectory rather than only behavior near global minima.

Coverage note — No substantial contributed material was omitted; proof derivations and routine appendix implementation details were excluded because the requested knowls retain results and experimental findings rather than proofs or reproducibility minutiae.

References

1.Vladimir Igorevich Arnold. Geometrical methods in the theory of ordinary differential equations, volume 250. Springer Science & Business Media, 2012. 2
2.Nachman Aronszajn. Theory of reproducing kernels. Transactions of the American mathematical society, 68(3):337–404, 1950. 22
3.Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019. 5
4.Shahar Azulay, Edward Moroshko, Mor Shpigel Nacson, Blake E Woodworth, Nathan Srebro, Amir Globerson, and Daniel Soudry. On the implicit bias of initialization shape: Beyond infinitesimal mirror descent. In International Conference on Machine Learning, pages 468–477. PMLR, 2021. 5
5.Francis Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, 18(1):629–681, 2017. 25
6.Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an Ornstein-Uhlenbeck like process. In Conference on learning theory, pages 483–513. PMLR, 2020. 1, 3
7.Youngmin Cho and Lawrence K Saul. Kernel methods for deep learning. In Proceedings of the 22nd International Conference on Neural Information Processing Systems, pages 342–350, 2009. 25
8.Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations, 2020. 2, 10
9.Alex Damian, Tengyu Ma, and Jason D Lee. Label noise SGD provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 34:27449–27461, 2021. 3
10.Yu Feng and Yuhai Tu. The inverse variance–flatness relation in stochastic gradient descent is critical for finding flat minima. Proceedings of the National Academy of Sciences, 118(9), 2021. 1
11.Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2020. 1
12.Crispin Gardiner. Stochastic methods, volume 4. springer Berlin, 2009. 3
13.Jeff Z HaoChen, Colin Wei, Jason Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. In Conference on Learning Theory, pages 2315–2357. PMLR, 2021. 5
14.Hangfeng He and Weijie Su. The local elasticity of neural networks. In International Conference on Learning Representations, 2020. 1
15.Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016. 9, 15
16.S. Hochreiter and J. Schmidhuber. Flat minima. Neural Computation, 9(1):1–42, 1997. 1
17.P Izmailov, AG Wilson, D Podoprikhin, D Vetrov, and T Garipov. Averaging weights leads to wider optima and better generalization. In 34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018, pages 876–885, 2018. 1
18.Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems, pages 8571–8580, 2018. 21
19.Stanisław Jastrz˛ebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in SGD. arXiv preprint arXiv:1711.04623, 2017. 3
20.Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point on optimization trajectories of deep neural networks. In International Conference on Learning Representations, 2019. 2
21.N. S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In In International Conference on Learning Representations (ICLR), 2017. 1, 3
22.Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images, 2009. 9
23.Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278–2324, 1998. 15
24.Michel Ledoux and Michel Talagrand. Probability in Banach Spaces: Isoperimetry and processes. Springer Science & Business Media, 2013. 20
25.Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and adaptive stochastic gradient algorithms. In Proceedings of the 34th International Conference on Machine Learning, volume 70, pages 2101–2110. PMLR, Aug 2017. 2
26.Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling SGD with stochastic differential equations (SDEs). In Advances in Neural Information Processing Systems, volume 34, 2021. 3
27.Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? –a mathematical framework. In International Conference on Learning Representations, 2022. 3
28.Tengyuan Liang, Tomaso Poggio, Alexander Rakhlin, and James Stokes. Fisher-Rao metric, geometry, and complexity of neural networks. In Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics, volume 89, pages 888–896. PMLR, 2019. 3
29.Chao Ma and Lexing Ying. On linear stability of SGD and input-smoothness of neural networks. Advances in Neural Information Processing Systems, 34, 2021. 2
30.Takashi Mori, Liu Ziyin, Kangqiao Liu, and Masahito Ueda. Logarithmic landscape and power-law escape rate of SGD. arXiv preprint arXiv:2105.09557, 2021. 1, 3, 8
31.Rotem Mulayoff, Tomer Michaeli, and Daniel Soudry. The implicit bias of minima stability: A view from function space. Advances in Neural Information Processing Systems, 34, 2021. 3, 7
32.Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv:1412.6614, 2014. 1
33.Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of SGD for diagonal linear networks: a provable benefit of stochasticity. Advances in Neural Information Processing Systems, 34, 2021. 3, 5
34.Dino Sejdinovic and Arthur Gretton. What is an RKHS? Lecture Notes, 2012. 25
35.Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014. 20
36.Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. 9, 15
37.Umut Simsekli, Levent Sagun, and Mert Gurbuzbalaban. A tail-index analysis of stochastic gradient noise in deep neural networks. In International Conference on Machine Learning, pages 5827–5837. PMLR, 2019. 3
38.Weijie Su. Neurashed: A phenomenological model for imitating deep learning training. arXiv preprint arXiv:2112.09741, 2021. 1
39.Yusuke Tsuzuku, Issei Sato, and Masashi Sugiyama. Normalized flat minima: Exploring scale invariant definition of flat minima for neural networks using PAC-Bayesian analysis. In International Conference on Machine Learning, pages 9636–9647. PMLR, 2020. 3
40.Xingyu Wang, Sewoong Oh, and Chang-Han Rhee. Eliminating sharp minima from SGD with truncated heavy-tailed noise. In International Conference on Learning Representations, 2022. 3
41.Stephan Wojtowytsch. Stochastic gradient descent with noise of machine learning type. part II: Continuous time analysis. arXiv preprint arXiv:2106.02588, 2021. 1, 3
42.Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pages 3635–3673. PMLR, 2020. 5
43.Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust generalization. Advances in Neural Information Processing Systems, 33, 2020. 1
44.Jingfeng Wu, Wenqing Hu, Haoyi Xiong, Jun Huan, Vladimir Braverman, and Zhanxing Zhu. On the noisy gradient descent that generalizes as SGD. In International Conference on Machine Learning, pages 10367–10376. PMLR, 2020. 1
45.Lei Wu and Jihao Long. A spectral-based analysis of the separation between two-layer neural networks and linear methods. Journal of Machine Learning Research, 23(119):1–34, 2022. 25
46.Lei Wu, Chao Ma, and Weinan E. How SGD selects the global minima in over-parameterized learning: A dynamical stability perspective. Advances in Neural Information Processing Systems, 31:8279–8288, 2018. 2, 3, 7, 10
47.Lei Wu, Zhanxing Zhu, and Weinan E. Towards understanding generalization of deep learning: Perspective of loss landscapes. arXiv preprint arXiv:1706.10239, 2017. 1
48.Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. In International Conference on Learning Representations, 2020. 1, 3, 8
49.Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017. 1
50.Hongyi Zhang, Yann N. Dauphin, and Tengyu Ma. Residual learning without normalization via better initialization. In International Conference on Learning Representations, 2019. 15
51.Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, et al. Towards theoretically understanding why SGD generalizes better than Adam in deep learning. Advances in Neural Information Processing Systems, 33, 2020. 3, 8
52.Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In International Conference on Machine Learning, pages 7654–7663. PMLR, 2019. 3, 8, 9
53.Liu Ziyin, Kangqiao Liu, Takashi Mori, and Masahito Ueda. Strength of minibatch noise in SGD. In International Conference on Learning Representations, 2022. 3

Citation

MLA

Wu, L., et al. “The Alignment Property of SGD Noise and How It Helps Select Flat Minima: A Stability Analysis”. Advances in Neural Information Processing Systems, vol. 35, 2022, pp. 4680–93, https://proceedings.neurips.cc/paper_files/paper/2022/file/1e55c38dd7d465c2526ae29d7ec85861-Paper-Conference.pdf.

APA

Wu, L., Wang, M., & Su, W. (2022). The alignment property of SGD noise and how it helps select flat minima: A stability analysis. Advances in Neural Information Processing Systems, 35, 4680–4693. https://proceedings.neurips.cc/paper_files/paper/2022/file/1e55c38dd7d465c2526ae29d7ec85861-Paper-Conference.pdf

Chicago

Wu, L., M. Wang, and W. Su. 2022. “The Alignment Property of SGD Noise and How It Helps Select Flat Minima: A Stability Analysis”. Advances in Neural Information Processing Systems 35: 4680–93. https://proceedings.neurips.cc/paper_files/paper/2022/file/1e55c38dd7d465c2526ae29d7ec85861-Paper-Conference.pdf.

Harvard

Wu, L., Wang, M. and Su, W. (2022) “The alignment property of SGD noise and how it helps select flat minima: A stability analysis”, Advances in Neural Information Processing Systems. Curran Associates, Inc., pp. 4680–4693. Available at: https://proceedings.neurips.cc/paper_files/paper/2022/file/1e55c38dd7d465c2526ae29d7ec85861-Paper-Conference.pdf.

Vancouver

1. Wu L, Wang M, Su W (2022) The alignment property of SGD noise and how it helps select flat minima: A stability analysis. In: Advances in Neural Information Processing Systems. Curran Associates, Inc., pp 4680–4693

BibTeX

@inproceedings{wu2022the,
  title = {The alignment property of SGD noise and how it helps select flat minima: A stability analysis},
  author = {Wu, Lei and Wang, Mingze and Su, Weijie},
  year = {2022},
  booktitle = {Advances in Neural Information Processing Systems},
  publisher = {Curran Associates, Inc.},
  volume = {35},
  pages = {4680-4693},
  url = {https://proceedings.neurips.cc/paper_files/paper/2022/file/1e55c38dd7d465c2526ae29d7ec85861-Paper-Conference.pdf}
}

Metadata:DOI registry

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: Authors

The alignment property of SGD noise and how it helps select flat minima: A stability analysis

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Loss-scaled alignment of SGD noise

Knowl 2 — Gradient-norm uniformity guarantees noise alignment

Knowl 3 — Exact alignment for over-parameterized linear models

Knowl 4 — Alignment in random ReLU feature models

Knowl 5 — Linearized SGD converts aligned noise into a stability constraint

Knowl 6 — Stable SGD minima have size-independent Frobenius flatness

Knowl 7 — A weaker alignment condition is sufficient for flatness selection

Knowl 8 — SGD escapes sufficiently sharp minima exponentially fast

Knowl 9 — Loss-dependent, geometry-aware noise explains efficient escape

Knowl 10 — Empirical validation across model classes and CIFAR-10 networks

Knowl 11 — Scope limitation of the stability explanation

References

Citation

BibTeX

Access the Paper