Variational Learning Finds Flatter Solutions at the Edge of Stability

TL;DR

This paper demonstrates that Variational Learning (VL) tends to find flatter solutions at the Edge of Stability, which improves generalization, by controlling the variational posterior and number of samples, supported by empirical validation on large networks.

Contribution

It is the first work to analyze the Edge of Stability dynamics of Variational Learning, revealing its ability to find flatter solutions compared to standard gradient descent.

Findings

VL finds flatter solutions at the Edge of Stability.

Controlling the variational posterior influences solution flatness.

Empirical results on ResNet and ViT support the theoretical analysis.

Abstract

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.