Variational Learning Finds Flatter Solutions at the Edge of Stability

2025-06-15

AI Summary

This work investigates the implicit regularization mechanism of variational learning (VL) in deep neural network training, specifically examining whether VL converges to flatter, better-generalizing minima than standard gradient descent under the “Edge of Stability” (EoS) dynamics. We introduce the EoS framework into VL theory for the first time and prove that VL—by controlling posterior covariance structure and Monte Carlo sample size—can actively steer and stabilize optimization trajectories within flatter EoS regimes. Combining theoretical analysis via quadratic models, PAC-Bayes bounds, and marginal likelihood maximization, we derive predictions empirically validated across ResNet-18 and ViT-S: VL solutions exhibit significantly higher flatness and lower test error. Our core contribution is uncovering VL’s intrinsic EoS-aware flatness preference—a novel theoretical lens for understanding its generalization advantage.

Technology Category

• Computer Vision: Learning & Optimization for CV
• Machine Learning: Learning with Manifolds
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Responsible Web: Machine-in-the-loop, human agency and autonomy

Abstract

Variational Learning (VL) has recently gained popularity for training deep neural networks and is competitive to standard learning methods. Part of its empirical success can be explained by theories such as PAC-Bayes bounds, minimum description length and marginal likelihood, but there are few tools 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 VL to show that it can find even flatter solutions. This is obtained by controlling the posterior covariance and the number of Monte Carlo samples from the posterior. These results are derived in a similar fashion as 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 in VL.

Problem

• Analyzing implicit regularization in Variational Learning
• Extending Edge of Stability framework to Variational Learning
• Demonstrating flatter solutions in deep neural networks

Innovation

• VL finds flatter solutions via Edge of Stability
• Controls posterior covariance and Monte Carlo samples
• Validated on ResNet and ViT networks empirically