Saved in:
| Main Authors: | , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2506.12903 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866918170844987392 |
|---|---|
| author | Ghosh, Avrajit Cong, Bai Yokota, Rio Ravishankar, Saiprasad Wang, Rongrong Tao, Molei Khan, Mohammad Emtiyaz Möllenhoff, Thomas |
| author_facet | Ghosh, Avrajit Cong, Bai Yokota, Rio Ravishankar, Saiprasad Wang, Rongrong Tao, Molei Khan, Mohammad Emtiyaz Möllenhoff, Thomas |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2506_12903 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Variational Learning Finds Flatter Solutions at the Edge of Stability Ghosh, Avrajit Cong, Bai Yokota, Rio Ravishankar, Saiprasad Wang, Rongrong Tao, Molei Khan, Mohammad Emtiyaz Möllenhoff, Thomas Machine Learning 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. |
| title | Variational Learning Finds Flatter Solutions at the Edge of Stability |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2506.12903 |