Variational Learning Finds Flatter Solutions at the Edge of Stability

• URL: http://arxiv.org/abs/2506.12903v3
• Date: Sun, 26 Oct 2025 03:59:44 GMT
• Title: Variational Learning Finds Flatter Solutions at the Edge of Stability
• Authors: Avrajit Ghosh, Bai Cong, Rio Yokota, Saiprasad Ravishankar, Rongrong Wang, Molei Tao, Mohammad Emtiyaz Khan, Thomas Möllenhoff,
• Abstract summary: We analyze the implicit regularization of Variational Learning (VL) through the Edge of Stability (EoS) framework.<n>This result is obtained by controlling the shape of the variational posterior as well as the number of posterior samples used during training.<n>We 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.
• Score: 47.33551032960102
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• On the Surprising Effectiveness of Large Learning Rates under Standard Width Scaling [11.168336416219857]
  Existing infinite-width theory would predict instability under large learning rates and vanishing feature learning under stable learning rates.<n>We show that this discrepancy is not fully explained by finite-width phenomena such as catapult effects.<n>We validate that neural networks operate in this controlled divergence regime under CE loss but not under MSE loss.
  arXiv  Detail & Related papers  (2025-05-28T15:40:48Z)
• Learning Dynamics of Deep Linear Networks Beyond the Edge of Stability [14.88048518927001]
  Deep neural networks trained using gradient descent with a fixed learning rate operate in the regime of "edge of stability"<n>We show that loss oscillations beyond EOS follow a period-doubling route to chaos.<n>Overall, our results contribute to explaining two key phenomena in deep networks.
  arXiv  Detail & Related papers  (2025-02-27T21:32:42Z)
• On the Dynamics Under the Unhinged Loss and Beyond [104.49565602940699]
  We introduce the unhinged loss, a concise loss function, that offers more mathematical opportunities to analyze closed-form dynamics. The unhinged loss allows for considering more practical techniques, such as time-vary learning rates and feature normalization.
  arXiv  Detail & Related papers  (2023-12-13T02:11:07Z)
• Convergence Analysis for Learning Orthonormal Deep Linear Neural Networks [27.29463801531576]
  We provide convergence analysis for training orthonormal deep linear neural networks. Our results shed light on how increasing the number of hidden layers can impact the convergence speed.
  arXiv  Detail & Related papers  (2023-11-24T18:46:54Z)
• Optimization Guarantees of Unfolded ISTA and ADMM Networks With Smooth Soft-Thresholding [57.71603937699949]
  We study optimization guarantees, i.e., achieving near-zero training loss with the increase in the number of learning epochs. We show that the threshold on the number of training samples increases with the increase in the network width.
  arXiv  Detail & Related papers  (2023-09-12T13:03:47Z)
• Stabilizing RNN Gradients through Pre-training [3.335932527835653]
  Theory of learning proposes to prevent the gradient from exponential growth with depth or time, to stabilize and improve training. We extend known stability theories to encompass a broader family of deep recurrent networks, requiring minimal assumptions on data and parameter distribution. We propose a new approach to mitigate this issue, that consists on giving a weight of a half to the time and depth contributions to the gradient.
  arXiv  Detail & Related papers  (2023-08-23T11:48:35Z)
• Fine-grained analysis of non-parametric estimation for pairwise learning [9.676007573960383]
  We are concerned with the generalization performance of non-parametric estimation for pairwise learning. Our results can be used to handle a wide range of pairwise learning problems including ranking, AUC, pairwise regression and metric and similarity learning.
  arXiv  Detail & Related papers  (2023-05-31T08:13:14Z)
• Globally Optimal Training of Neural Networks with Threshold Activation Functions [63.03759813952481]
  We study weight decay regularized training problems of deep neural networks with threshold activations. We derive a simplified convex optimization formulation when the dataset can be shattered at a certain layer of the network.
  arXiv  Detail & Related papers  (2023-03-06T18:59:13Z)
• A Unified Algebraic Perspective on Lipschitz Neural Networks [88.14073994459586]
  This paper introduces a novel perspective unifying various types of 1-Lipschitz neural networks. We show that many existing techniques can be derived and generalized via finding analytical solutions of a common semidefinite programming (SDP) condition. Our approach, called SDP-based Lipschitz Layers (SLL), allows us to design non-trivial yet efficient generalization of convex potential layers.
  arXiv  Detail & Related papers  (2023-03-06T14:31:09Z)
• SGD with Large Step Sizes Learns Sparse Features [22.959258640051342]
  We showcase important features of the dynamics of the Gradient Descent (SGD) in the training of neural networks. We show that the longer large step sizes keep SGD high in the loss landscape, the better the implicit regularization can operate and find sparse representations.
  arXiv  Detail & Related papers  (2022-10-11T11:00:04Z)
• Multi-scale Feature Learning Dynamics: Insights for Double Descent [71.91871020059857]
  We study the phenomenon of "double descent" of the generalization error. We find that double descent can be attributed to distinct features being learned at different scales.
  arXiv  Detail & Related papers  (2021-12-06T18:17:08Z)
• Mean-field Analysis of Piecewise Linear Solutions for Wide ReLU Networks [83.58049517083138]
  We consider a two-layer ReLU network trained via gradient descent. We show that SGD is biased towards a simple solution. We also provide empirical evidence that knots at locations distinct from the data points might occur.
  arXiv  Detail & Related papers  (2021-11-03T15:14:20Z)
• Subquadratic Overparameterization for Shallow Neural Networks [60.721751363271146]
  We provide an analytical framework that allows us to adopt standard neural training strategies. We achieve the desiderata viaak-Lojasiewicz, smoothness, and standard assumptions.
  arXiv  Detail & Related papers  (2021-11-02T20:24:01Z)
• Compressive Sensing and Neural Networks from a Statistical Learning Perspective [4.561032960211816]
  We present a generalization error analysis for a class of neural networks suitable for sparse reconstruction from few linear measurements. Under realistic conditions, the generalization error scales only logarithmically in the number of layers, and at most linear in number of measurements.
  arXiv  Detail & Related papers  (2020-10-29T15:05:43Z)
• Understanding Self-supervised Learning with Dual Deep Networks [74.92916579635336]
  We propose a novel framework to understand contrastive self-supervised learning (SSL) methods that employ dual pairs of deep ReLU networks. We prove that in each SGD update of SimCLR with various loss functions, the weights at each layer are updated by a emphcovariance operator. To further study what role the covariance operator plays and which features are learned in such a process, we model data generation and augmentation processes through a emphhierarchical latent tree model (HLTM)
  arXiv  Detail & Related papers  (2020-10-01T17:51:49Z)
• Layer-wise Conditioning Analysis in Exploring the Learning Dynamics of DNNs [115.35745188028169]
  We extend conditioning analysis to deep neural networks (DNNs) in order to investigate their learning dynamics. We show that batch normalization (BN) can stabilize the training, but sometimes result in the false impression of a local minimum. We experimentally observe that BN can improve the layer-wise conditioning of the optimization problem.
  arXiv  Detail & Related papers  (2020-02-25T11:40:27Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.