Variational Learning Finds Flatter Solutions at the Edge of Stability
This work provides insights into the training dynamics of VL for researchers in deep learning, but it is incremental as it extends existing EoS analysis to a specific method.
The paper tackled the problem of understanding the implicit regularization of Variational Learning (VL) in deep neural networks by analyzing it through the Edge of Stability (EoS) framework, showing that VL can find flatter solutions than gradient descent, with empirical validation on large networks like ResNet and ViT.
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.