# Variational Learning Finds Flatter Solutions at the Edge of Stability

arXiv: https://arxiv.org/abs/2506.12903 · HTML: https://mlsanity.com/paper/2506.12903/

## 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 (most cited first, 30 of 56)

- [Weight Uncertainty in Neural Networks](https://mlsanity.com/paper/1505.05424/) (2015)
- [Computing Nonvacuous Generalization Bounds for Deep (Stochastic) Neural Networks with Many More Parameters than Training Data](https://mlsanity.com/paper/1703.11008/) (2017)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [How to Escape Saddle Points Efficiently](https://mlsanity.com/paper/1703.00887/) (2017)
- [User-friendly introduction to PAC-Bayes bounds](https://mlsanity.com/paper/2110.11216/) (2021)
- [A Tail-Index Analysis of Stochastic Gradient Noise in Deep Neural Networks](https://mlsanity.com/paper/1901.06053/) (2019)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [Practical Deep Learning with Bayesian Principles](https://mlsanity.com/paper/1906.02506/) (2019)
- [Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability](https://mlsanity.com/paper/2209.15594/) (2022)
- [The Break-Even Point on Optimization Trajectories of Deep Neural Networks](https://mlsanity.com/paper/2002.09572/) (2020)
- [A Bayesian Perspective on Generalization and Stochastic Gradient Descent](https://mlsanity.com/paper/1710.06451/) (2017)
- [Fast and Scalable Bayesian Deep Learning by Weight-Perturbation in Adam](https://mlsanity.com/paper/1806.04854/) (2018)
- [Understanding Gradient Descent on Edge of Stability in Deep Learning](https://mlsanity.com/paper/2205.09745/) (2022)
- [Non-Vacuous Generalization Bounds at the ImageNet Scale: A PAC-Bayesian Compression Approach](https://mlsanity.com/paper/1804.05862/) (2018)
- [The Heavy-Tail Phenomenon in SGD](https://mlsanity.com/paper/2006.04740/) (2020)
- [Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model](https://mlsanity.com/paper/1907.04164/) (2019)
- [Scalable Marginal Likelihood Estimation for Model Selection in Deep Learning](https://mlsanity.com/paper/2104.04975/) (2021)
- [Adaptive Gradient Methods at the Edge of Stability](https://mlsanity.com/paper/2207.14484/) (2022)
- [The Bayesian Learning Rule](https://mlsanity.com/paper/2107.04562/) (2021)
- [Understanding the Generalization Benefit of Normalization Layers: Sharpness Reduction](https://mlsanity.com/paper/2206.07085/) (2022)
- [On the Generalization Benefit of Noise in Stochastic Gradient Descent](https://mlsanity.com/paper/2006.15081/) (2020)
- [The Description Length of Deep Learning Models](https://mlsanity.com/paper/1802.07044/) (2018)
- [Variational Learning is Effective for Large Deep Networks](https://mlsanity.com/paper/2402.17641/) (2024)
- [The alignment property of SGD noise and how it helps select flat minima: A stability analysis](https://mlsanity.com/paper/2207.02628/) (2022)
- [Understanding Optimization in Deep Learning with Central Flows](https://mlsanity.com/paper/2410.24206/) (2024)
- [Understanding Edge-of-Stability Training Dynamics with a Minimalist Example](https://mlsanity.com/paper/2210.03294/) (2022)
- [Non-Vacuous Generalization Bounds for Large Language Models](https://mlsanity.com/paper/2312.17173/) (2023)
- [Learning threshold neurons via the "edge of stability"](https://mlsanity.com/paper/2212.07469/) (2022)
- [Large Learning Rate Tames Homogeneity: Convergence and Balancing Effect](https://mlsanity.com/paper/2110.03677/) (2021)
- [First Exit Time Analysis of Stochastic Gradient Descent Under Heavy-Tailed Gradient Noise](https://mlsanity.com/paper/1906.09069/) (2019)

## Built on by (most cited first, 2 of 2)

- [Generalization at the Edge of Stability](https://mlsanity.com/paper/2604.19740/) (2026)
- [Convergence of a class of gradient-free optimisation schemes when the objective function is noisy, irregular, or both](https://mlsanity.com/paper/2512.03225/) (2025)

## 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.