Variational Learning Finds Flatter Solutions at the Edge of Stability

Rongrong Wang (Michigan State University) · Rio Yokota (Institute of Science Tokyo) · Thomas Möllenhoff (RIKEN AIP) · Mohammad Emtiyaz Khan (RIKEN) · Saiprasad Ravishankar (Michigan State University) · Molei Tao (Georgia Tech) · Avrajit Ghosh (University of California, Berkeley) · Bai Cong (Institute of Science Tokyo; RIKEN Center for AI Project)
deep neural networksedge of stabilityempirical validationeos dynamicsflat solutionsgradient descentimplicit regularizationmarginal likelihoodminimum description lengthpac-bayes boundsposterior samplesresnetvariational learningvariational posteriorvit

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.