Variational Learning Finds Flatter Solutions at the Edge of Stability
Not on the map. Why?
Abstract (arXiv)
Variational Learning (VL) has recently gained popularity for training deep neural networks. Part of its empirical success can be explained by theories such as PAC-Bayes bounds, minimum description length and marginal likelihood, but little has been done to unravel the implicit regularization in play. Here, we analyze the implicit regularization of VL through the Edge of Stability (EoS) framework. EoS has previously been used to show that gradient descent can find flat solutions and we extend this result to show that VL can find even flatter solutions. This result is obtained by controlling the shape of the variational posterior as well as the number of posterior samples used during training. The derivation follows in a similar fashion as in the standard EoS literature for deep learning, by first deriving a result for a quadratic problem and then extending it to deep neural networks. We empirically validate these findings on a wide variety of large networks, such as ResNet and ViT, to find that the theoretical results closely match the empirical ones. Ours is the first work to analyze the EoS dynamics of VL.
Key facts. Cited by 2 papers in the map (data through 14 Oct 2026). Builds on 56 of its 74 references that are in the map. First posted on arXiv 15 Jun 2025.
Builds on
56 of 74 references.
References · 56
- 2015Weight Uncertainty in Neural Networkscited by 554
- 2017Computing Nonvacuous Generalization Bounds for Deep (Stochastic) Neural Networks with Many More Parameters than Training Datacited by 533
- 2021Gradient Descent on Neural Networks Typically Occurs at the Edge of Stabilitycited by 427
- 2017How to Escape Saddle Points Efficientlycited by 285
- 2021User-friendly introduction to PAC-Bayes boundscited by 239
- 2019A Tail-Index Analysis of Stochastic Gradient Noise in Deep Neural Networkscited by 214
- 2023Stabilizing Transformer Training by Preventing Attention Entropy Collapsecited by 161
- 2019Practical Deep Learning with Bayesian Principlescited by 158
- 2022Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stabilitycited by 155
- 2020The Break-Even Point on Optimization Trajectories of Deep Neural Networkscited by 150
- 2017A Bayesian Perspective on Generalization and Stochastic Gradient Descentcited by 135
- 2018Fast and Scalable Bayesian Deep Learning by Weight-Perturbation in Adamcited by 131
38 more
- 2022Understanding Gradient Descent on Edge of Stability in Deep Learningcited by 128
- 2018Non-Vacuous Generalization Bounds at the ImageNet Scale: A PAC-Bayesian Compression Approachcited by 126
- 2020The Heavy-Tail Phenomenon in SGDcited by 124
- 2019Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Modelcited by 117
- 2021Scalable Marginal Likelihood Estimation for Model Selection in Deep Learningcited by 106
- 2022Adaptive Gradient Methods at the Edge of Stabilitycited by 101
- 2021The Bayesian Learning Rulecited by 99
- 2022Understanding the Generalization Benefit of Normalization Layers: Sharpness Reductioncited by 86
- 2020On the Generalization Benefit of Noise in Stochastic Gradient Descentcited by 82
- 2018The Description Length of Deep Learning Modelscited by 72
- 2024Variational Learning is Effective for Large Deep Networkscited by 65
- 2022The alignment property of SGD noise and how it helps select flat minima: A stability analysiscited by 64
- 2024Understanding Optimization in Deep Learning with Central Flowscited by 64
- 2022Understanding Edge-of-Stability Training Dynamics with a Minimalist Examplecited by 56
- 2023Non-Vacuous Generalization Bounds for Large Language Modelscited by 53
- 2022Learning threshold neurons via the "edge of stability"cited by 51
- 2021Large Learning Rate Tames Homogeneity: Convergence and Balancing Effectcited by 50
- 2019First Exit Time Analysis of Stochastic Gradient Descent Under Heavy-Tailed Gradient Noisecited by 44
- 2022Anticorrelated Noise Injection for Improved Generalizationcited by 42
- 2023Implicit Bias of Gradient Descent for Logistic Regression at the Edge of Stabilitycited by 42
- 2022Second-order regression models exhibit progressive sharpening to the edge of stabilitycited by 40
- 2022Beyond the Edge of Stability via Two-step Gradient Updatescited by 38
- 2023(S)GD over Diagonal Linear Networks: Implicit Regularisation, Large Stepsizes and Edge of Stabilitycited by 36
- 2020Handling the Positive-Definite Constraint in the Bayesian Learning Rulecited by 31
- 2024Edge of Stochastic Stability: Revisiting the Edge of Stability for SGDcited by 30
- 2022Explicit Regularization in Overparametrized Models via Noise Injectioncited by 26
- 2023Good regularity creates large learning rate implicit biases: edge of stability, balancing, and catapultcited by 25
- 2023SAM operates far from home: eigenvalue regularization as a dynamical phenomenoncited by 24
- 2019Non-Gaussianity of Stochastic Gradient Noisecited by 22
- 2023Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaoscited by 22
- 2023Gradient Descent Monotonically Decreases the Sharpness of Gradient Flow Solutions in Scalar Networks and Beyondcited by 21
- 2024Flat-LoRA: Low-Rank Adaptation over a Flat Loss Landscapecited by 20
- 2019How noise affects the Hessian spectrum in overparameterized neural networkscited by 17
- 2022Quadratic models for understanding catapult dynamics of neural networkscited by 17
- 2023Sharpness-Aware Minimization and the Edge of Stabilitycited by 17
- 2023From Stability to Chaos: Analyzing Gradient Descent Dynamics in Quadratic Regressioncited by 17
- 2021Eliminating Sharp Minima from SGD with Truncated Heavy-tailed Noisecited by 14
- 2023Noise Stability Optimization for Finding Flat Minima: A Hessian-based Regularization Approachcited by 14
50 most cited of 56 All 56 (JSON)
Built on by
Citing papers per year: 2025: 1; 2026: 1.
2 papers.