Variational Learning Finds Flatter Solutions at the Edge of Stability

We thank the reviewer for their time and effort in providing a critical and valuable feedback. We are happy to see the reviewer recognize the novelty and the strong empirical verification of the theorem. Below are our answer to the reviewer's comments to the weakness (W) and questions (Q).

W1. Theorem statements are slightly unclear; for example, Eq. (7) is defined after Theorem 1, whereas it should be defined within or before that theorem.

A1. Thanks for the suggestion. We will move Eq. 7 within the theorem.

W2. Mathematical link between Theorem 1 and Hypothesis 1 is unclear.

A2 Similar gaps exist in existing EoS analysis for GD [1], SGD [2], and SAM [3]. All such works and ours rely on the principle that stability limits can be faithfully determined by local dynamics on a quadratic approximation. This is mentioned and discussed in lines 97-115. A rigorous proof is complex and beyond this paper’s scope, but extensions can be built on self-stabilization theory [4] which is a rather new development and not yet extended to settings such as ours. If it helps, we will add a dedicated appendix section to elaborate on this connection.

W3: Hypothesis-1 assumes isotropic Gaussian noise, whereas Theorem-1 does not; the paper does not explain this discrepancy.

A3: Thanks for pointing this out. This does not change the overall conclusion but has a technical implication. We will add an explanation in the paper. Essentially, the stability threshold is determined by the alignment between the noise covariance ($\mathbf{\Sigma}$) and the Hessian ($\mathbf{Q}$), a relationship captured by the term (Equation-23). For isotropic noise, this term becomes an equality $Tr(\mathbf{\Sigma}\mathbf{Q}^3)=\sigma^2 \sum_{i} \lambda^3_{i}$, but otherwise it forms an upper bound ($Tr(\mathbf{\Sigma}\mathbf{Q}^3) \leq \sum_{i} \sigma_{i}^2 \lambda^3_{i}$) (lines 531-538). However, in both scenario, this provides a sufficient condition for expected descent so the difference is only in the technical details of the analysis.

W4 +Q3: "the setup resemble Langevin dynamics or SGLD rather than genuine VL, raising doubts about the scope of the analysis. …. 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?".

A4: There appears to be a misunderstanding. VL noise is added to the parameter which is fundamentally different from the SGLD noise added to the gradient. This is even true for a simple linear regression case, therefore it is not appropriate to view the method from a Langevin perspective. This difference is in fact highlighted in lines 203-214. We hope this clarifies the misunderstanding.

W5: 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.

A5: Thanks for raising an important point. Our theoretical results intentionally focus on full gradients to separate the effect of posterior-sampling noise from the minibatch-gradient noise. The study of the interplay between the two is left to our experiments (Section 4.3) which confirms that minibatch noise further helps to reduce the stability threshold. A formal theory for the combined setting is challenging (due to heavy-tailed nature of the noise) and left as future work.

W6: 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.

A6: Thanks for pointing out this distinction. The later extensions are used to check if the theory extends to these tricky cases which are difficult to analyze. We find that the theory still holds even when the assumptions are violated. This only increases the relevance of the theoretical results showing that they are valid under much less restrictions. We will add some clarification in the text to help the readers understand the reason to use more complex posteriors. Our hope is that future work will extend the theory to other non-Gaussian cases.

W7: NLP experiments that follow the theoretical Gaussian setting would strengthen the paper.

A7: To verify that Variational GD with isotropic Gaussian noise achieves lower sharpness compared to GD on NLP tasks, we add a classification head to a frozen BERT-mini [5, 6] backbone and finetune the head on SST-2 [7], which is a sentiment classification dataset. For both GD and Variational GD, we use the same learning rate of 0.05. Here we show the results:

Variational GD with isotropic Gaussian noise

GD

[1] Cohen, Jeremy M., et al. "Gradient descent on neural networks typically occurs at the edge of stability.", ICLR-21.

[2] Lee, Sungyoon, and Cheongjae Jang. "A new characterization of the edge of stability based on a sharpness measure aware of batch gradient distribution." ICLR-23.

[3] Long, Philip M., and Peter L. Bartlett. "Sharpness-aware minimization and the edge of stability." JMLR-24

[4] Damian et al, "Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability." ICLR-23.

[5] Devlin, Jacob, et al. "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding." NAACL-19

[6] Bhargava, Prajjwal, Aleksandr Drozd, and Anna Rogers. "Generalization in NLI: Ways (Not) To Go Beyond Simple Heuristics." arXiv

[7] Socher, Richard, et al. "Recursive deep models for semantic compositionality over a sentiment treebank." EMNLP-13

Q1. How does VL being less stable enables it to find flatter minima?

A. You are correct that VL's stricter bound makes it less stable on a simple quadratic. However, in the complex landscapes of neural networks, this instability is precisely what guides it toward flatter solutions.

Consider a landscape with two minima: one sharp and one flat. Gradient Descent (GD) with a small learning rate will simply converge to the nearest minimum, even if it's the undesirable sharp one. Variational Learning (VL), in contrast, has a stricter stability condition that is violated by the sharp minimum's high curvature. This instability prevents convergence there, effectively "kicking" the optimizer out of that sharp basin. It can then find and settle in the broad, flat minimum where the stability condition is satisfied. This process acts as an implicit filter, enabling VL to escape poor local minima and preferentially find the flatter solutions associated with better generalization.

Q2: Use of isotropic assumption.

A: You're correct, the use of isotropic noise in Hypothesis 1 is intentional. We use it to create a setting where our theoretical bound is tightest, allowing for a clear and direct empirical validation. Here’s the technical reasoning:

General (Sufficient) Bound: Our main result (Theorem 3.1) provides a sufficient condition for descent that holds for any (diagonal) anisotropic noise. It's based on the trace inequality $Tr(\mathbf{\Sigma}\mathbf{Q}^3) \leq \sum_{i} \sigma_{i}^2 \lambda^3_{i}$, which is a conservative upper bound in the general case.

Isotropic Case: For isotropic noise, the trace inequality becomes an exact equality. This makes the stability threshold upper bound tight, which is what we validate empirically under Hypothesis 1.

Our condition is intentionally sufficient, as it's analytically tractable. It works by ensuring every term (line-541) in the expected loss change is negative. A necessary-and-sufficient condition (requiring the whole sum to be negative) is imposed on the whole Hessian eigen-spectrum, and other versions like Theorem 3 require strong assumptions (e.g., uniform alignment), making them less practical.

We have revised the manuscript and Appendix E to clarify this distinction and to demonstrate why our sufficient condition is adequate for analyzing the neural network experiments.

Q3 is addressed with W4.

Q4 Numbering issue.

A: Thank you for spotting the error in section number. We will fix it in our update.

If the reviewer's concerns are addressed, we would be happy to see an increase in score.

Thank you for the insightful questions. Here we answer the two questions.

Q1 In scenarios where the Hessian spectrum is highly ill-conditioned or dominated by a few large eigenvalues, how reliable is the stability condition derived in Theorem 1? Could this potentially bias the analysis?

A1 Short Answer: Theorem-1 holds for a general quadratic Hessian spectrum (even when ill-conditioned). In this case, for stability to hold (expected loss, $\mathbb{E} [\Delta \ell]$ to decrease), the noise-variance in the (direction of the eigenvector with large eigenvalue) needs to be small which is given in condition (6) in paper. For stability to hold in direction of eigenvector with larger $\lambda_{i}$, $\sigma_{i}$ needs to be small due to monotonicity of VF(.). Below we give a detailed explanation.

Mathematical explanation: The expected descent depends on the alignment of the noise covariance $\mathbf{\Sigma}$ and Hessian ($\mathbf{Q}$), through $Tr(\mathbf{\Sigma} \mathbf{Q}^3)$, that controls the stability, which is seen in the following equation:

$

\mathbb{E} [\Delta \ell] = - \rho \mathbf{g}^T \mathbf{g} + \frac{\rho^2}{2} \mathbf{g}^T \mathbf{Q} \mathbf{g} + \frac{\rho^2}{2N_{s}}Tr(\mathbf{\Sigma} \mathbf{Q}^3)

$

The value $Tr(\mathbf{\Sigma} \mathbf{Q}^3)$ is high when the noise is larger in the direction of eigenvector with large eigenvalue, hence maximizing $Tr(\mathbf{\Sigma} \mathbf{Q}^3)$ and making $\mathbb{E} [\Delta \ell]>0$. This is because maximum of $Tr(\mathbf{\Sigma} \mathbf{Q}^3)$ is $\sum_{i} \sigma_{i} \lambda^3_{i}$ (equality condition of Von Neuman's trace inequality) and is achieved when $\mathbf{\Sigma}$ and $\mathbf{Q}$ are aligned. Theorem-1 provides a sufficient condition by using this maximum alignment.

Q2 Although the experiments show that the theory appears to generalize beyond its assumptions, are there any specific examples or regimes where the theory breaks down (e.g., highly non-isotropic covariance)?

A2 Here are two cases we found where there is a slight mismatch between theory and experiment:

1. In experiments with vision transformers, the sharpness limit is usually observed to be slightly higher than the stability thresholds. This is not only the case with VL, but even for GD, the sharpness is higher than $\frac{2}{lr}$. This is a due to a phenomenon called attention entropy collapse commonly occurring in attention layers [1].

2. For non-isotropic noise, a small gap between our theoretical bound and empirical results can occasionally appear. This is because the theory is based on a worst-case analysis using the Von Neumann trace inequality $Tr(\mathbf{\Sigma} \mathbf{Q}^3) ≤ \sum_{i} \sigma^2_{i} \lambda^3_{i}$ , which considers maximal alignment between the noise and the Hessian. If the alignment during training is not exact, the bound is not perfectly tight, which can explain a slight mismatch. This differs from the isotropic case, where this bound is always an equality $Tr(\mathbf{\Sigma} \mathbf{Q}^3) = \sigma^2 \sum_{i}\lambda^3_{i}$, leading to the perfect agreement between theory and experiments seen in our work.

[1] Zhai, Shuangfei, et al. "Stabilizing transformer training by preventing attention entropy collapse." International Conference on Machine Learning. PMLR, 2023.

We thank you again for your thorough review and time that helped improve this manuscript. We will add these detailed discussions in our revised version.

We thank the reviewer for their detailed and constructive feedback. We’re encouraged that the reviewer found the paper to be well-written, novel, and supported by strong experiments.

We also thank the reviewer for suggesting to add a “global” discussion in the end of the paper. We are happy to revise it along the lines that reviewer suggested. We will do the following:

1. We will include a concrete discussion regarding PAC-Bayes, emphasizing the fact that PAC-Bayes focuses on generalization bounds [1] but don’t take an optimizer’s implicit regularization. Our work fills this gap by using the Edge of Stability.

2. We will add a discussion on how the results can be extended to broader variational approximation, where the posterior covariance is not assumed to fixed (including low and full-rank Gaussians connecting them to Shampoo and MuON optimizers).

   a) We will specifically highlight our numerical results using IVON (where covariance is learned) but also those with more flexible posterior (e.g., t-distribution). We will also highlight the difficulty of obtaining such theoretical results by using “preconditioned sharpness”. For VON, preconditioned sharpness is below 2/lr, while for Adam it’s around 2/lr [2]. As you noted, analyzing adaptive covariance is challenging because it requires jointly tracking the mean and preconditioner dynamics (Equation 8), even for a simple quadratic. We will highlight this finding and designate it as a direction for future work.

3. We will add a discussion on extending MSE to cross-entropy etc. We will also add some numerical results.

   a) For, CE loss, the key difference lies in the peak sharpness. Due to its stricter stability bound, VL's sharpness peaks below 2/lr, whereas for GD, sharpness peaks to 2/lr. Following this peak, the sharpness for both methods decays to zero. Our preliminary results on a 2-layer MLP with CIFAR-10 confirm this: VL consistently exhibits lower sharpness than GD and decays faster. This decay to zero is a known property of the cross-entropy loss. The sharpness is governed by the Hessian, which in turn depends on the second derivative of the loss, $\ell^{''} \propto p_{i}(1-p_{i})$ where $p_{i}$ is the model's predicted probability. As training progresses and the model becomes confident, its predictions $p_{i}$ approach either 0 or 1. In both cases, the term $p_{i}(1-p_{i})$ approaches zero. This causes the entire Hessian, and thus the sharpness, to vanish. Here is an experiment with lr=0.05 trained with CE loss where the sharpness is reported across several epoch windows.

4. We will add generalization metric to support the statement “One of the messages of this paper is that minimizing the variational objective (1) alone does not guarantee good generalization”.

   a) We will clarify that what we aimed to say with that sentence is that we want to find good minimizers of the objective, and just blindly running a method without thinking about implicit bias is not enough. We will also add a sharpness plot in the figure alongside the variational objective to make it more clear and we report some of the metrics here alongside the original figure.

[1] Alquier et al, “User-friendly introduction to PAC-Bayes bounds”.

[2] Cohen et al, “Adaptive Gradient Methods at the Edge of Stability”. https://arxiv.org/abs/2207.14484 Many thanks for pointing out typos. We will fix them!

Many thanks for pointing out typos. We will fix them!

We thank the reviewer for their detailed feedback. We are glad that the reviewer found the writing to be exceptionally clear, results clear, and the experimental results perfectly validate the theory. Below is our response to some of your suggestions.

Q1: Fixing the reference issue in numbering equations and sections.

A1: Thanks for catching this. In the updated manuscript, we will properly refer to whether it is equation or section.

Q2: Can you plot (in addition to the $2/\rho$ line) also a line that includes the "VF" factor?

A2: Thanks for the suggestion! We will add it in the final version.