Title: Variational Learning Finds Flatter Solutions at the Edge of Stability
Venue: NeurIPS 2025
Accepted: yes
Keywords: variational learning

noise injection

edge of stability

generalization

deep neural networks

bayesian deep learning
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nVariational Learning (VL) is reported to implicitly steer the optimization toward flatter minima by applying Gaussian perturbations to the weights, and this study theoretically analyzes this phenomenon from the perspective of the Edge of Stability. In particular, from the theoretical analysis of quadratic losses, the authors derive that the stability threshold of VL is smaller than that of Gradient Descent (GD), and they demonstrate that this threshold can be further lowered by enlarging the posterior variance and reducing the number of Monte Carlo samples. In addition, experiments on various neural networks confirm that the sharpness consistently adheres to the theoretically predicted Variational Factor, achieving lower sharpness and higher test accuracy than GD.\n\nSTRENGTHS AND WEAKNESSES:\n## Strength\n- The paper’s primary contribution is extending Edge-of-Stability sharpness analysis, which was previously limited to SGD variants, to the broadly used class of Gaussian-noise methods in variational learning.  \n- Through this extension, the authors rigorously show that the eigenvalue-based quadratic optimality condition for VL is stricter than that for SGD; this is a first-time result, and the accompanying numerical experiments confirm that this condition indeed governs optimization behavior.  \n\n\n## Weakness\n- The theorem statements are mathematically slightly unclear; for example, Eq. (7) is defined after Theorem 1, whereas it should be defined within or before that theorem.  \n- Based on my reading of Appendix B (Sections 7 and 8 excluded), the mathematical link between Theorem 1 and Hypothesis 1 is unclear: Theorem 1 and Theorem 2 in Appendix merely give sufficient conditions for a one-step loss decrease under Gaussian-perturbed GD (and Theorem 3 adds necessity under special assumptions), yet Hypothesis 1 claims that the largest eigenvalue will hover around the derived threthold. So Theorem 1 is just the result for the one-step analysis, but the hypothesis considers the entire trajectory; this gap, central to the paper, requires a more rigorous discussion.  \n- Hypothesis 1 assumes isotropic Gaussian noise, whereas Theorem 1 does not; the paper does not explain this discrepancy.  \n- Because variational learning is characterized by adaptive variance, fixing both the noise covariance and the matrix \\(Q\\) makes the setup resemble Langevin dynamics or SGLD rather than genuine VL, raising doubts about the scope of the analysis.  \n- Practical VL employs minibatching, yet the paper offers no justification for ignoring minibatch-induced noise; if such noise is negligible, some justification should be provided.  \n- There is a substantial gap from Theorem 1 to the experimental sections: although Theorem 1 presumes fixed Gaussian noise, later experiments use Student-t noise and adaptive covariances that violate the theorem’s assumptions, leaving their relevance unclear.  As far as I understand, we need a significant modification to the proof of Theorem 1, which does not seem straightforward.\n- Experiments are restricted to small 32 × 32 image datasets; evaluations on larger benchmarks such as ImageNet or NLP tasks are absent. I believe experiments for such ImageNet and NLP that follow the theoretical Gaussian setting would strengthen the paper more than the current Student-t or adaptive cases that diverge from the theory.\n\nQUESTIONS:\nSee the weakness section. The following are additional questions.\n- Although the paper interprets Theorem 1 as showing that VL is guided toward flatter minima because a stricter eigenvalue bound must be satisfied, Theorem 1 is only a sufficient condition for a single-step decrease in the loss. One could therefore equally conclude that VL is *less* stable, since the stricter bound makes a loss decrease harder to achieve. Could the authors discuss which interpretation is accurate and supply theoretical or empirical evidence to support that position?  \n- Hypothesis 1 assumes isotropic Gaussian noise, whereas the setting of Theorem 1 does not. Is this isotropic assumption intended to match the isotropic variance condition used in Theorem 3 to derive a necessary-and-sufficient bound, or is there another reason for the mismatch?  \n- Because the covariance matrix and \\(Q\\) are fixed, the algorithm keeps injecting noise from the same Gaussian distribution at every iteration, a situation that resembles Langevin dynamics or SGLD more than adaptive variational learning. If the work is viewed through a Langevin lens, how do the results connect with existing Langevin analyses, and what is genuinely novel for VL?  \n- A minor issue: after Appendix B the document unexpectedly reverts to Section 7 and Section 8, which is confusing; unifying the section numbering would improve readability.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe main contribution of the paper is the idea of using the concept of Edge of Stability (EoS) to analyze training dynamics and implicit regularization of variational inference for deep neural networks. After an introduction, the paper first recaps the basics of EoS for the quadractic loss, and then extends the analysis to variational inference in the simplified setting, where the variational family is the mean-field Gaussians with fixed and known variances. \n\nTheorem 1 gives conditions for a decrease in the expected (quadratic) loss, and Lemma 2 gives conditions for a decrease in the actual loss (a high probability bouund). Using simple and well-designed set of numerical experiments, the paper provides insight into these results.\n\nNext, the paper switches focus to larger models and proposes a hypothesis of how the top Hessian eigenvalue (of the loss) is related to the \"stability threshold\" (Hypothesis 1). Using several models (ViT, MLP, ResNet20) trained on CIFAR10, they provide evidence for this hypothesis, e.g. Fig3 suggests that  lower sharpness leads to better generalization and that the top Hessian eigenvalue tracks the stability threshold (Fig4). Moveover, they provide several sensitivity studies for different design choices (e.g. the posterior variance, the number of posterior samples, the effect of fat-tails in the variational family etc). In the last part of the paper, the authors relate the discussion to the Variational Online Newton (VON) rule and show empirically how different hyperparameters (batch size, learning rate) affect the optimization process.\n\nSTRENGTHS AND WEAKNESSES:\nOverall, the paper is well-written and relatively easy to follow. The idea of using EoS in the context of variational inference is very interesting and novel to the best of my knowledge. The paper provides novel theoretical results in a simplified setting, and then extends this to more complicated problems empirically. The theoretical analysis as well as the numerical experiments appears to be well-designed and sound.\n\nIn my opinion, the paper lacks a proper discussion. The final chapter \"5 Conclusion and discussion\" appears to me much more as a summary than a discussion. While the paper is well-written, the structure and coherence good be approved. To me, the paper appears a little bit as many small pieces (e.g. theoretical insights, small numerical experiments etc) that have been written independently. While most of sections contains a bit of discussion, it would be a significantly improvement (in my opinion) if the authors would include a \"global\" discussion in the end of the paper (there is indeed enough room), connecting the findings from each subsection and putting it into a bigger perspective. For example,\n\n- The relation to PAC-Bayes. The abstract says: \"Part of its empirical success can be explained by theories such as PAC-Bayes bounds.....\" and  the conclusion says: \"Empirically, we observe that generalization in VL cannot be fully explained by PAC-Bayes bounds alone....\", but besides that PAC-Bayes is not mentioned in the paper.\n\n- The way I see paper is that the authors use the variational framework and approximate posterior as a means for regularization and stability, which is completely fine, but in many case the approximate posterior is also an object of interest in itself. Clearly, some of the results take advantage of the fact that the entropy is constant because the posterior covariance is fixed, but I believe the paper would benefit from a discussion of this and to what degree these results can be extended broader variational approximation, where the posterior covariance is not assumed to fixed.\n\n-  The paper uses the MSE as a loss for the classification models. While I understand the choice in the context of the theoretical results, this choice is maybe not the most obvious choice in a general probabilistic setting, and  therefore, the paper would also benefit from a discussion of to what degrees these results would extend for example cross-entropy etc.\n\n- Section 4.3 starts out by stating that: \"One of the messages of this paper is that minimizing the variational objective (1) alone does not guarantee good generalization\". I definitely agree with this statement, and I believe it is somewhat common knowledge in the field. Nevertheless, the section continues with an experiment, where it is shown that different optimization hyperparameter affects the optimization process and the final value of the variational objective. Yet, these results are not coupled in any to metrics quantifying generalization ability. Therefore, if the authors want to support their statement, they should indeed include some metric for generalization besides the variational objective itself.\n\nQUESTIONS:\nAs mentioned above, I believe the paper could be further improved, if \n\n- the authors would include a \"global discussion\" connecting the different observations and insights obtained through the paper and seeing into a bigger perspective.\n\n- the authors would elaborate on whether these insights generalize to a more general variational setting, where the posterior covariance is not fixed. What would happen if the variational family was the mean-field Gaussians, low-rank or full-rank Gaussians?\n\n- the authors would elaborate on the importance of the MSE loss function. To what degree do these findings extend to other losses?\n\nIdeally, a summary of the related works should also be included in the main paper.\n\n\n\n\nMinor details:\n- Line 121: Typo \"implmented\"\n- v_i is mentioned in Theorem 1, but not defined formally. (v_{1,t} is defined in section 4 and I assume v_i in Theorem has a similar definition)\n- IVON abbrevation does not seems to be defined (Although the reference is included)\n-\n\nLIMITATIONS:\nNone", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper argues theoretically why introducing fixed weight noise during training with gradient descent (GD) leads to convergence to flatter regions of the loss function, verifies these predictions analytically in deep-learning models, and then extends both theory and empirical evaluation to variational learning (VL; i.e., where the noise amplitude is learned as well by maximizing the ELBO), and beyond GD to adaptive optimizes such as Adam. Theoretical predictions are *quantitative* and match empirical observations.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\nThe writing is exceptionally clear. Edge of Stability theory is reviewed concisely but very clearly. Deriving the main result first in a simple setting (quadratic loss, Gaussian noise with fixed variance, GD with fixed learning rate and without preconditioning) makes it very clear where results are coming from and prepares readers for the later extensions.\n\n*Qualitatively* the main result almost seems obvious in hindsight (the authors also explain it in simple terms without math on lines 52-53). But I see this as a strength rather than a weakness because (i) it speaks to the clarity of the presentation, (ii) both the theoretical predictions and the empirical evaluation go beyond qualitative statements and analyze the stability region *quantitatively*, and (iii) the quantitative analysis is eventually extended beyond the simple original case of fixed noise amplitude and learning rate to VL with an Adam optimizer.\n\n## Minor weakness\n\nI really can't find any weakness apart from a trivial point:\n\n- When referring to Equations or Sections, only the number (without \"Equation\" or \"Section\") is rendered. This is somewhat confusing at first (e.g., on line 138: \"An exact expression ... is derived in 6 where we show ...\" --> I first checked Equation 6 before I realized that the authors are referring to *Section* 6)\n\nQUESTIONS:\nIn Figure 3, can you plot (in addition to the $2/\\rho$ line) also a line that includes the \"VF\" factor (i.e., by solving the right part of Eq. 6 with equality instead of \"<\" numerically for $\\lambda$)?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-07-03"}