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
Venue: NeurIPS 2025 spotlight ()
Keywords: variational learning
noise injection
edge of stability
generalization
deep neural networks
bayesian deep learning
Primary area: deep_learning
TL;DR: Variational Learning Finds Flatter Solutions at the Edge of Stability

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.

Variational Learning Finds Flatter Solutions at the Edge of Stability
Avrajit Ghosh !."
Bai Cong2 Rio Yokota²
Saiprasad Ravishankar'
Rongrong Wang!
Molei Tao 4
Mohammad Emtiyaz Khan
Thomas Mollenhoff
Michigan State University
Institute of Science Tokyo
3
RIKEN Center for AI Project
*Georgia Institute of Technology
"Work performed in part during internship the RIKEN Center for Al Project.
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 VI. 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 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.
1 Introduction
Variational Learning (VL) has been used to perform deep learning from early on [Graves, 2011, 1, Blundell et al., 2015] and recently also started to show good results at large scale. It has been shown to outperform state-of-the-an optimizers without any increase the cost. For example, on ImageNet, VL substantially improves overfitting commonly seen in AdamW and for pretraining GPT-2 from scratch, VL achieves lower validation perplexity than AdamW [Shen et al., 2024]. For low-rank fine-tuning of Llama-2 (7B), VL improves both accuracy (by 2.8%) and calibration (by 4.6%) [Cong et al., 2024, Li et al., 2025], Moreover, VL methods explicitly derived by using PAC-Bayes bounds [Wang et al., 2023b, Zhang et al., 2024b] have shown consistent improvements over AdamW. All such results confirm the importance of VL for deep learning.
Despite these successes, the theoretical mechanisms behind the good performance of VL remain poorly understood. It is often assumed that simplistic Gaussian posteriors, such as those used currently for deep learning, may not be enough because they are poor approximations of the true posterior; lack of good prior another issue. Despite these concerns, VL shows good performance in practice. Theories such as minimum description length [Hinton and Van Camp. 1993, Hochreiter and Schmidhuber, 1997, Blier and Ollivier, 2018), PAC-Bayes [Dziugaite and Roy. 2017, Zhou et al., 2019, Lotfi et al., 2024, Alquier et al., 2024] and marginal likelihood [Smith and Le, 2017, Immer et al., 2021] can partially explain the success. In particular, PAC-Bayes theory provides a natural explanation for why flatter solutions may generalize better, see for instance the discussion in [Alquier ct al., 2024, Scction 3.3]. However these theories say little about the regularizing properties of the
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Posterior variance
0.1
0.5
2.0
Number MC MC MC
with variance 0.05
10
VL with variance 0.1
5000
Epoch
Theory
Variance
0.05
Theory
Practice
Practice
Variance 0.1
1000
3050
4000
5000
Epoch
(a)
(b)
Figure I: Panel (a): The left figure shows trajectory traces of VL on quadratic problem with an isotropic variational posterior whose mean is learned but variance is set to a fixed value. The trajectory becomes more unstable as the posterior variance increased and number of Monte-Carlo samples decreased. We provide an exact expression to compute the stability threshold at which the iterations become unstable (Theorem 3.1). Panel (b): We show the validity of the threshold neural network training. The right figure (top) shows this on CIFAR-10 for an MLP where VL. achieves lower sharpness than GD when posterior variance is increased. The bottom figure shows that the sharpness (solid line) matches the stability threshold obtained by our theorem (dashed line).
learning algorithm. Similarly deep learning, the presence of implicit regularization is likely also at play in VL, but few tools exist to unravel these effects.
In this work, we analyze the implicit regularization of VI. algorithms at the Edge of Stability. The EoS analysis has previously been used to show that Gradient Descent (GD) with constant learning rate p implicitly biases the trajectories towards flatter solutions, where the sharpness (defined the operator norm of the loss Hessian) hovers around 2/p. We extend this analysis to VL and show that sharpness can be further lowered by controlling the posterior covariance and the number of Monte-Carlo samples used to compute posterior expectations; see an illustration Figure 1a. Similarly to the standard EoS technique, we first derive exact expression for the stability threshold for VL on a quadratic problem and then propose extensions to general loss functions. We then empirically validate these finding on a wide variety of deep networks, including Multi- Perceptron, ResNet, and Vision Transformers; see an example in Figure 1b. We observe similar results when posterior shape automatically learned, for instance, by using diagonal covariance Gaussians and heavy-tailed posteriors. Code to replicate these results is available at
Theoretical Tools to Analyze Variational Learning
Variational Learning optimizes the variational reformulation of Bayesian learning [Zellner, 1988], where the goal is to find good approximations of the Gibbs distribution exp(-((0))/ Z with partition function over parameters E Rd. Specifically, we seek the closest distribution (0)) in set of distributions Q that minimizes the expected loss regularized by the entropy:
arg min (e))) - H(g)-
v
This is an instance of the maximum-entropp principle. The objective equivalent to minimizing the KL divergence to the Gibbs posterior and naturally encourages higher-entropy (flatter) solutions due to the entropy H(g); for example, see the illustrative example in Khan and Rue [2023, Figure 1]. Similar arguments can also be made with the mininum-descrrption--eee
principles [Hinton and Van Camp, 1993, Graves, 2011, Blier and Ollivier, 2018]. In practice, Variational Learning has started to show good results achieving generalization performance better than the state of the art

---

Sharp Posterior finds sharp minimum
posterior finds
Lons
ritinum
Sharp
Figure 2: VL's mechanism for flatter minima: The posterior variance determines the minima's location. small variance settles the posterior in sharp minima (left), while larger variance allows it to explore and find flat minima (right).
optimizers across several tasks [Shen et al., 2024, Cong et al., 2024). This matches with the intuition: Because the variational objective prefers wider distributions, expect the posterior to be located in the region where the loss {()) is flatter, see Figure 2.
Despite this intuition, there is little work to analyze the implicit regularization that could help us understand how VL favors flatter regions and how we can control it. Existing theories do not sufficiently address this; review of the related work is in Appendix A. We suspect the implicit regularization to be related to the shape of the posterior but currently there are no results explicitly characterizing this. VL is also closely connected to weight-noiss
weight-perturbation methods where the goal is to optimize Ex(e) [ + <)] with g(c) being a fixed distribution to inject weight noise. This can be seen as a special case of variational learning where the shape of the posterior is fixed and only the location is learned. Multiple works [Zhu et al., 2019, Nguyen et al., 2019, Zhang al., 2019, Jin et al., 2017, Simsekli al., 2019] have analyzed generalization behavior of such weight-noise variational methods but, even for this simple case. there are no studies connecting them the EoS results for GD. In this paper, we will address these gaps and provide both theoretical and empirical results regarding the edge of stability phenomenon of such variational GD (VGD).
2.1 Edge of Stability for Gradient Descent
We will briefly review the EoS result for GD. The standard EoS literature relies result for quadratic problem and then extends it to deep neural networks. For instance, consider the following
quadratic loss: (0)) = 2 0 To0, where Q = Σ V.V.V. 1=1
(2)
Here, is positive definite matrix with A; being its 4th largest eigenvalue and v; being the corresponding eigenvector: The following result states the condition under which one step of GD leads to a decrease the loss.
Lemma 2.1. (Descent Lemma) For GD update Oc+1 =. -Ve((()) on the quadratic loss (2), the loss decreases at each step, that is, we have
E(0 E(O.) < if and only if A S for all i.
This lemma implies that GD converges on a quadratic loss if all cigenvalues satisfy A; < 2/p. This is different way of writing the standard condition that maximum eigenvalue X < 2/p. The result extends to any smooth {(e) ) and for such cases, the condition implies that the learning rate p < 2/8. where 3 sup, (V2((0))e is the bound the Hessian norm [Nesterov et al., 2018, Chapter 2]. This condition is necessary and sufficient for descent of GD on quadratic.
While the descent lemma is predictive of convergence of GD for smooth functions, deep learning exhibits a more complex behavior. When training deep neural networks with constant learning rate the Hessian's operator norm (or sharpness) tends to settle around the value 2/p. This phenomenon

---

referred to as the Edge of Stability. Deep neural networks often operate at this edge and converge in monmoonttonic. unstable manner. comprehensive empirical investigation of this phenomenon is given by Cohen et al. [2021], where two phases of training neural networks were noticed. In the first phase referred to as "progressive sharpening', the Hessian's operator norm, [| increases and slowly approaches 2/p. This is followed by the second, EoS phase, where sharpness hovers around 2/p and the loss continues to decrease an oscillatory fashion.
This phenomenon does not happen on a quadratic loss, but only for certain losses with nonzero third- order derivative. As shown by Damian et [2023], once sharpness reaches 2/p p. local quadratic approximation is insufficient to capture the dynamics because the third order Taylor expansion term of the loss becomes significant. This cubic term represents the gradient of the sharpness, which serves as negative feedback to counteract progressive sharpening and stabilize the sharpness around 2/ p in GD. Instead of divergence, the iterates exhibit oscillatory or non-monotonic behavior even when the sharpness reaches p/p. This is the reason why 2/p is also called the 'Stability T'reshold'.. Increasing p reduces the edge, which could then drive the iterates toward flatter minima that may generalize better. EoS analysis can help understand such implicit regularization during training, especially for nonconvex problems.
Different optimizers have different stability thresholds which depends on the sharpness value [[V)(()))]l For instance, Sharpness Aware Minimization (SAM) leads different, smaller stability threshold [Long and Bartlett, 2024] compared to p/p for GD. In this work, we show a similar result where VL has smaller threshold than GD. We show this by deriving the stability threshold for a simpler quadratic problem first and analyze several factors such as posterior covariance and the number of posterior samples that influence the threshold. Then, we empirically demonstrate that similar results hold for the case when VL is used to train deep neural networks.
3 Stability Threshold for a Simple Case of Variational Learning
We start with a simple VL setting where the goal is to estimate a Gaussian g = NO m, Σ) whose mean is unknown and covariance E is fixed. We assume Σ = o²l an isotropic covariance matrix, with scalar variance o². Khan and Rue [2023] [Section 1.3.1] show that this can be estimated by the following Variational GD algorithm:
M++1 + m. -E [[(m + t)].
(4)
This algorithm can be implemented by the following version where the expectation is estimated by drawing N, Monte Carlo samples to approximate the expectation as follows
1 m f- m. N
vi(m; + es), e, ~~0,00,
(5)
The updated iterate mf41 depends the N, Monte Carlo samples = [E1, E2, Compared to the standard gradient descent, the update step in Variational GD (VGD). is determined by gradients averaged over local neighborhood of perturbed weights. As result, the update introduces two interacting effects influencing its stability threshold: (1) a perturbation effect, originating from the perturbation covariance Σ, and (2) smoothing effect, resulting from averaging gradients across N, Monte Carlo samples. Similarly to Lemma 2.1, we can derive the stability threshold for the above VGD and show that it is smaller than that of GD.
Theorem 3.1. Consider the VGD update (5) the quadratic loss (2), then
Ee Ee(((m(+1)) {(m:) <0 if A 2 VF
N (
for all i,
(6)
where Cie (A,m, v.)?, and VF(() denotes the Variational Factor given by
VF(z) = p. v 1 sinh (1 arcsinh
(
3 -
See Appendix B for proof where we also discuss the case where Q is low rank. The theorem is analogous to the descent lemma for GD and it states that the Variational GD (5) also decreases the

---

2.0
Unstable
A
6 1.0
VF(z)
Stable
0.1
25 N
50
Nylo²
1/0²
(a) Variational Factor (VF)
(b) Descent probability
(c) Stability
Figure 3: (a) Solid black curve shows the theoretical stability-thresholdd of VGD as a function of N./o². The curve is clearly lower than the stability threshold of GD, shown with the horizontal dashed, gray line. (b) Empirical verification on scalar quadratic problem with curvature À where we plot the empirically computed probability of descent for VGD runs with different values of o². We show heatmap for (A, 1/o²) values where lighter colors indicate higher probability of descent. We overlay the heatmap with the stability thresholds of VGD (red solid curve), clearly showing that theoretical limit shown in (6) matches the empirical probability.
(c) The figure further includes N, and marks the region where pair (Na, 2²) will either lead to descent or not (marked with 'stable" and 'unstable' respectively).
expected loss over if each eigenvalue X; is less than 2/p times a function called the Variational Factor (VF). The VF function is strictly less than 1, therefore the stability threshold is strictly less than p/p. Unlike GD, the condition here only sufficient but not necessary. A necessary condition can also be derived but the one above is sufficient for this paper.
Figure 3a shows the stability threshold 2/p (dashed horizontal line). The exact value of VF depends on its argument z which in Theorem 3. mainly depends on the number of Monte-Carle Carlo samples N, and posterior variance o², but also on the loss-dependent constant Ca which is essentially a function of X and m, v,.
function of where is clear that it always lower bounds
Theorem guarantees a decrease in the expected loss. Below, we state another result to show that the actual loss also decreases with high probability if the expected loss decreases by a margin > 0. Lemma 3.2. In the same setting Theorem 3.1, when the expected loss at next iteration is smaller than the previous loss by some margin > O, that is,
Tee ()))(+))) < <((m.) - 5,
then (m{+1))((m) < O occurs with probability at least 1- 1-2 exp (-c, min for constants C1, C2 > depending only on p. Q, and E.
The result also shows that the probability with which this happens increases with the number of Monte Carlo samples N..
So far, we show that the stability threshold for Variational GD is smaller than that of GD, but can we also ensure it to be smaller in practice? The answer is yes, and it can be done by controlling o² and N₁. We will now demonstrate this on 1D quadratic loss {(m) 2 Amm For such cases, GD with step-size p descends whenever the curvature À is bounded by 2/p, but we can show that VGD descends only for lower curvature values. To show this, we run VGD for many (X, pairs. For each run, we use 10 random realizations of then record how often the loss decreases, and finally compute the approximate descent probability (E((mf+1) < f(mm)).
Figure 3b shows this probability as a heatmap where lighter color indicate a higher probability of descent; a white pixel indicates a probability of 1 and a black one indicates probability of O. We overlay this with the solid red curve showing the theoretical stability- of Variational GD as dictated by (6), that is, À = (2/p) VF(N,/o²) The theoretical curve closely matches the transition boundary where probabilities transition from 1 to O. The figure clearly shows that by increasing o² we can reduce the stability threshold practice too. Figure further illustrates the effect of changing N., showing the regions where a (N,, a²)) pair lead to descent (marked as 'stable") or otherwise (marked with "unstable"). As expected, the relationship linear and the same effect obtained by either increasing N. or decreasing o².
5

---

3.1 Nature of Stability in GD and VGD
The introduction of noise in VGD fundamentally changes the nature its stability compared the deterministic behavior of GD. The descent lemma for GD on a quadratic guarantees that the iterates are asymptotically stable, which is defined as follows:
Definition 3.3 (Asymptotic Stability Lyapunov [1992], Chapter 2). Let O" be the minimum. An iterate 0% is asymplotic stable if is Lyapunov stable and there also exists ô such that if the algorithum starts within a 8-neighborhood of the fixed point, the iterate will converge to the fixed point. Formally, if ||0 6" < 8 for a finite k, then:
lim 0
8-11 =
(8)
8-000
GD on quadratic is asymptotically stable because the condition learning-rate p < 2/ ensures that the iteration matrix - (Q) has all eigenvalues less than 1, guaranteeing convergence of the iterates to its fixed point. Analogously VGD, the iterates Stochastically Stable which is defined as follows Definition 3.4 (Stochastic Stability Kushner [1972], Chapter 2). Let 0" be the minimum. The VGD iterates O is said be stochastically stable in the mean-square sense if there exists a constant C > such that for any initial point Oo, the iterates 0f satisfy:
Ee [100 - 0" 11 - 0 11, for all t > O.
This form of stability ensures that the iterates remain bounded in the mean-square error sense, preventing them from diverging. Practically, this corresponds to the behavior where the distribution of the iterates converges to a stationary distribution, rather than the iterates themselves converging to single fixed point as in asymptotic stability. The condition for probabilistic descent, established in Theorem 3.1, is what characterizes this stable behavior. By ensuring the loss decreases on average, it prevents the divergence of the iterates, confining them to stable random walk that converges in distribution to a stationary state around the minimum. Accordingly, any reference to "stability" or a "stability threshold" throughout for VGD in this paper refers to this notion of stochastic stability.
3.2 Comparison with Regularization Effect of SGD
Our work differs fundamentally from studies on the regularization effects of mini-batch noise in SGD, such as Wu et al. [2022], Zhu et al. [2019]. Wu et al. [2018), Ibayashi and Imaizumi [2023), Mulayoff and Michaeli [2024], which analyze gradient noise and its role in promoting flatter solutions. In contrast, VGD introduces noise directly in the weights, inducing a structured and anisotropic form of gradient noise shaped by the curvature of the loss landscape, something not captured by existing SGD analyses. While most SGD-based studies assume Gaussian gradient noise, despite empirical evidence of heavy-tailed behavior [Gurbuzbalaban et al., 2021, Nguyen al., 2019], we show that -iitriibttdt gradients, that
for the quadratic loss, Gaussian perturbations in weights lead to is, & ~ (>((V((m.)) N, QEQ), where the covariance is amplified along sharper directions. This formulation requires no assumptions on the shape of gradient noise (unlike, e.g., Lee and Jang [2023]). To complement our work, we further perform an empirical study using weight perturbations drawn from a heavy-tailed distribution in deep neural networks.
4 Experiments on Deep Neural Networks
Similarly to the GD case, we expect the stability threshold for the quadratic to serve as an EoS limit. That is, expect that, when using variational learning for deep neural networks, the sharpness should hover around the stability limit. By controlling the covariance shape and number of Monte Carlo samples we can stere optimization into regions where sharpness is much lower than GD. The following hypothesis formally states this intuition which we verify through extensive experiments. Hypothesis 1. For twice-differentiable loss {m) optimized using the update rule (5), the top Hessian eigenvalue (()))) hovers around the stability threshold,
2 VF
N₂ 0²
C1. c..).
where Ci.t (((m))) and V1, denotes the top eigenvector of the Hessian VPE(m.), and (((m)) is the gradient evaluated at m.

---

GD
GD with Gaussian weight perturbation
2/p
Train Accuracy
Train Accuracy
Train Accuracy
1.00
1.00
1.00
0.75
0.75
0.75
0.50
0.50
0.50
0.25
0.25
0.25
50000 Train steps
100000
50000 Train steps
100000
5000 Train steps
10000
Test Accuracy
Test Accuracy
Test Accuracy
0.6
0.4
0.2
0.2
50000 Train steps
100000
50000 Train steps
100000
5000 steps
10000
Sharpness
Sharpness
Sharpness
60
400
40
200
20
20
50000 Train steps
100000
50000
100000
5000 Train steps
10000
Train steps
VIT
(b) MI.P
(c) ResNet-20
Figure 4: Smaller sharpness corresponds to higher test accuracy for network architectures trained on CIFAR-10 10. Panels show (a) ViT, (b) MLP, and (c) ResNet-20. Full batch GD with Variational GD achieves lower sharpness and better test accuracy.
Normalized Norralizee-Shhrrpnes Sharpness (a 0.05) Variational Factor (o 0.05)
Normalizee--Shrroees to 0.001)
Variational Factor 0.001)
1.100
1.305
3.300
0.825
0.623
0.825
0.52
0.000
1000
2005
3000
4000
5000
0.000
1003
2000
3000
4000
5000
0.002
1003
2000
3000
4000
5000
Epoch
Eooch
(a) 0 0.01
(b) 0.02
(c) 0.05
Figure 5: Normalized Sharpness (2((m))la/(2//p)) hovers around the Variational Factor in MLP
To validate the hypothesis, we conduct extensive experiments on standard architectures (MIPs, ResNets) and modern ones such as Vision Transformers. In Figure 4, we plot the full training dynamics of Variational GD, including test accuracy, training accuracy, and sharpness. Across all three architectures on the CIFAR- 10 classification task using MSE loss, Variational GD with isotropic Gaussian noise consistently achieves lower sharpness and higher test accuracy compared to GD.
In Figure 5, we plot the normalized sharpness (()))((me))l divided by 2/p) for MLPs, evaluated at each training step, alongside the corresponding Variational Factor (VF) for various learning rates and posterior variances o². The results support Hypothesis 1, showing that the normalized sharpness closely tracks the predicted VF across settings. This further validates the use of local quadratic approximation to analyze stability dynamics, consistent with several prior works. Unlike the stability analysis for GD quadratic where the threshold 2/ is constant, the stability condition in VGD
7

---

MC Sample
100 MC Samples
2/p
20
30
Starpness
1.5
Snarpness
20
20
Variance (x10-5)
Variance (x10
Variance (x10-5)
(a) 2/p =
(h) 2/p = 20
(c) 2/p = 10
Figure 6: Sharpness of the final iterate for MLP trained CIFAR-10 for variational GD across learning rate noise covariance Σ, and posterior samples N.. Higher variance and smaller samples lead lower sharpness. For panel (c), training did not converge for variance values 2 and 5, therefore not shown in the plot.
dynamic since it depends on the gradient evaluated at the mean. This explains the fluctuation of the stability threshold across iterations in VGD. In the Appendix (Figure 14), we demonstrate that similar phenomenon holds for ResNet-20. Next, we isolate and examine the roles of the posterior variance and the number of samples in determining this stability threshold.
Perturbation effect of Gaussian variance: In Figure 6, we present experiments on VGD for a classification task using multi-layer perceptron (MLP), where we plot the sharpness of the final iterate as a function the variance o² of an isotropic Gaussian. For three different learning rates, p 0.05, 0.1, and 0.2, we VGD with different a². Across all learning posterior samples N., we observe that larger variance consistently leads to lower sharpness. This exactly aligns with Hypothesis 1, larger perturbations reduce the stability threshold, promoting escape from sharper regions of the loss landscape.
p and numbers of
Smoothing effect of posterior samples N.: Larger posterior sample size N. reduces the vari- ance of the perturbed gradient estimator g ((VE((m)) N. QEQ) under a quadratic loss. As N₁, decreases, the variance increases (scaling as 1/N.), lowering the stability threshold and enabling escape from sharper regions. This effect is shown in Figure where smaller N, consistently yields lower sharpness across learning rates and variances. Additional discussion on the role of posterior samples, is provided in Appendix C.
EN Vf(m, <;), which satisfies g
4.1 Edge of Stability under Heavy-Tailed Noise
In variational gradient descent (VGD), the stability threshold depends not only on the perturbation covariance and sample size, but also the posterior's tail behavior. We empirically investigate this by drawing weight perturbations from a Student-t distribution with degrees of freedom o, where smaller a induces heavier tails and larger deviations in the perturbed gradients. The distribution becomes heavier-tailed as a decreases, with the Gaussian recovered as + co. In Figure 7, we train an MLP on CIFAR-I0 and observe that heavier-tailed perturbations yield lower sharpness and better test accuracy. Similar trends are confirmed for ResNet-20 and ViT in Appendix F. While theoretical analysis under heavy-tailed noise challenging, our results highlight that the posterior shape plays critical role in generalization.
4.2 Adaptive Edge of Stability and Variational Online Newton Methods
Recent work by Cohen et al. [2022] studies the dynamics of adaptive gradient methods, introducing the concept of Adaptive Edge of Stability (AEoS). In this section, we will show that similar phenomenon happens vatural-gradient variational learning.
Cohen et al. [2022] show that for adaptive gradient methods instead of sharpness [(m ( modified quantity, termed preconditioned sharpness -V2 l| hovers around p/p. where is preconditoner. Adaptive gradient methods differ from standard gradient descent as the

---

d=3
a=6
(Gaussian perturbation)
Train Accuracy
Test Accuracy
Sharpness
1.00
40
0.4
0.3
30
0.50
0.2
0.25
20
0.10
5000
10000
5000
10000
5000 Train steps
10000
Train steps
Train steps
Figure 7: Sharpness dynamics with noise injection from a heavy tailed Student t posterior parameter-
ized by a. Perturbations from heavier tails (smaller a)) lead smaller sharpness.
Adam
Weight perturbation,
Weight perturbation,
2/p
Train Accuracy
Test Accuracy
Pre-conditioned Sharpness
1.0
0.8
0.6
10a
0.4
0.2
0.2
10a
0.1
25000
50000
25000
50000
25000
50000
Train steps
Train
Train
Figure 8: Preconditioned sharpness for Adam and IVON across temperatures т. -mmaller 7 shrinks the posterior and yields larger preconditioned sharpness.
former can adapt their preconditioner P. and move into high curvature regions. For example in Adam, a preconditioner P. is updated in an exponential moving average (EMA) fashion with parameter B2:
Vt+1 -2 + (1 B2) V(mm)), P++1
V
Vi+1 82
a
M
m, Vi(mi)/P:+t. (10)
Here, all operations such squaring or division of vectors are performed elementwise.
Natural-gradient VL methods which learn a complete Gaussian posterior g: = ((',,,P7)) take a similar form to the above adaptive gradient methods. An instance of this is the Variational Online Newton (VON) update rule [Khan and Rue, 2023, Eq. 12).
M4+1 f m p P E, V (0)) and P.+1 < -1- P. 32 E. (2((O)
(1))
Here, the posterior covariance P. is learned using the loss Hessians. The variational GD in Eq. (5) corresponds the special case where P, fixed across iterations. Adaptive optimizers such as Adam, RMSProp, and Adadelta can be seen as special cases of VON, as shown in [Khan and Rue, 2023, Section 4.2).
Here, we study the AEoS for a recent large-scale implementation of Equation (11) by Shen et al. [2024] called IVON (Improved Variational Online Newton). There, the update for the preconditioner is approximated using Stein's identity to estimate the Hessian from N. samples:
N₂ Σ
E (diag( (VE(()))
V((O.)
m 0²
0, ~ q.-
(12)
N
In Figure 8, we compare the preconditioned sharpness of IVON and Adam on MLPs by varying the temperature т, which linearly scales covariance of the posterior distribution qt, see [Shen et al., 2024, Algorithm line 8]. We observe that IVON consistently yields lower preconditioned sharpness than the 2/p threshold typically observed in Adam. We conjecture the lower sharpness is due to the noise injection and smoothing effects estimating both the gradient and curvature. Additional comparisons for ResNet-20 and VIT are provided in Appendix F. Computing the exact stability threshold for IVON is nontrivial, as it requires a joint analysis the coupled dynamics in Eq. 11, and an interesting direction for future work.

---

p= 0.05
p= 0.005
p= 0.0005
B=128
B=256
B=512
14.79
Vartational Objective
Variational Objective
14.76
14.79
14.78
14.72
5000 Steps
10000
0
25000
50000
Steps
Sharpness
Sharpness
75
50
40
25
20
0
5000 Train Steps
10000
25000
50000
Train Sleps
(a) Varying learning rate
(b) Varying batch size
Figure 9: Minimizing the variational objective depends on the choice of hyperparameter and its implicit regularization effeet, visible in both the objective value and sharpness reduction.
4.3 Effect of Batch Size
For nonconvex problems such as deep learning, minimizing the variational objective (1) is not, by itself, sufficient to guarantee flat minima low objective value. The choice optimization dynamics and hyperparameters play crucial role well. We demonstrate this by rumning IVON with mini-batching across varying learning rates and batch sizes B. As shown in Figure 9, larger learning rates and smaller batch sizes consistently lead to better local minima of the variational objective. The present work offers an explanation why large learning rates work better, but we also speculate that smaller batch sizes encourages broader posteriors in flatter minima allowing smaller objective value. These results further highlight the importance of choice of hyperparameter and optimizers in variational learning. Poor performance of variational methods as claimed in the literature may not be due to flaws the variational formulation itself but due to unfavorable optimization dynamics or choices of hyperparameters.
Conclusion and Discussion
In this work, we study the regularization effect in Variational Learning (VL) that enables it to find solutions with better generalization than Gradient Descent (GD). We show that the sharpness dynamics in VL can be accurately tracked through its instability mechanism under a local quadratic approximation of the loss. We argue that fully explain generalization in VL, must look beyond theoretical frameworks like PAC-Bayes bounds and instead understand the optimization dynamics and the implicit regularization they induce. We hope our work takes positive step this direction and motivates further investigation into the role of training dynamics in Bayesian deep learning.
Acknowledgements
This work is supported by the Bayes duality project, JST CREST Grant Number JPMICR21 12. A. G.Ghosh and R. Wang acknowledge support from NSF Grant CCF-2212065. S.Ravishankar acknow]- edge support from NSF CAREER CCF-2442246 and NSF Grant CCF-221 2065. M. grateful for partial supports by NSF Grant DMS- 1847802, Cullen-Peck Scholarship. and Emory-GT Almmnnit Award. We acknowledge Zhedong Liu for his help with the experiments on heavy-tailed distribution and Pierre Alquier for helpful discussions.
10

---

Impact Statement
This work improves understanding of how learning algorithms can find solutions that generalize better. It may lead to more reliable and robust Al systems. While primarily theoretical, the insights could support better model training in practical settings. We see minimal risk of misuse, but encourage responsible use in sensitive applications.
References
A. Agarwala and Y. Dauphin. SAM operates far from home: eigenvalue regularization as a dynamical phenomenon. In International Conference Machine Learning (ICML), 2023.
A. Agarwala, F. Pedregosa, and J. Pennington. Second-order regression models exhibit progressive sharpening to the edge of stability. In International Conference on Machine Learning (ICML), 2023.
K. Ahn, S. Bubeck, S. Chewi, Y. T. Lee, F. Suarez, and Y. Zhang. Learning threshold neurons via edge of stability. Advances in Neural Information Processing Systems (NeurIPS), 2024.
P. Alquier et al. User-friendly introduction to PAC-Bayes bounds. Foundations and Trends@ in Machine Learning, 2024.
S.-I. Amari. Natural gradient works efficiently in learning. Neural computation, 1998.
A. Andreyev and P. Beneventano. Edge of stochastic stability: Revisiting the edge of stability for SGD. arXiv:2412.20553, 2024.
S. Arora, Z. Li, and A. Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In International Conference Machine Learning (ICML), 2022.
L. Blier and Y. Ollivier. The description length of deep learning models. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
C. Blundell, J. Cornebise, K. Kavukcuoglu, and D. Wierstra. Weight uncertainty in neural networks.
In International Conference on Machine Learning (ICML). 2015.
L. Chen and J. Bruna. Beyond the edge of stability via two-step gradient updates. In International Conference on Machine Learning (ICML), 2023.
X. Chen, K. Balasubramanian, P. Ghosal, and B. K. Agrawalla. From stability to chaos: Analyzing gradient descent dynamics in quadratic regression. Transactions on Machine Learning Research, 2024. URL https: http////ppe//eeeeeeeeeeeeee net wtwt//forrmiddd-iillllo..
J. Cohen, S. Kaur, Y. Li, J. Z. Kolter, and A. Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations (ICLR), 2021.
J. Cohen, A. Damian, A. Talwalkar, J. Z. Kolter, and J. D. Lee. Understanding optimization in deep learning with central flows. In International Conference on Learning Representations (ICLR), 2025.
J. M. Cohen, B. Ghorbani, S. Krishman, N. Agarwal, S. Medapati, M. Badura, D. Suo, D. Cardoze, Z. Nado, G. E. Dahl, al. Adaptive gradient methods at the edge of stability, arXiv:2207.14483 14484, 2022.
B. Cong, N. Daheim, Y. Shen, D. Cremers, R. Yokota, M. E. Khan, and T. Möllenhoff. Variational low-rank adaptation using IVON, arXiv preprint arXiv:2411.04422 2024.
A. Damian, Nichani, and D. Lee. Self-stabilization:: The implicit bias of gradient descent the edge of stability. In International Conference on Learning Representations (ICLR), 2023.
G. K. Dziugaite and D. M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Conference on Uncertainty in Artificial Intelligence (UAI), 2017.

---

M. Even, S. Pesme, S. Gunasekar, and N. Flammarion. (S)GD over diagonal linear networks: Implicit bias, large stepsizes and edge of stability, Advances in Neural Information Processing Systems (NeurIPS), 2024.
A. Ghosh, S. M. Kwon, R. Wang, S. Ravishankar, and Q. Qu. Learning dynamics of deep matrix factorization beyond the edge of stability. International Conference on Learning Representations (ICLR), 2025.
A. Graves. Practical variational inference for neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2011.
M. Gurbuzbalaban, U. Simsekli, and L. Zhu. The heavy-tail phenomenon in SGD. In International Conference on Machine Learning (ICML), 2021.
G. E. Hinton and D. Van Camp. Keeping the neural networks simple by minimizing the description length of the weights. In Conference on Learning Theory (COLT), 1993.
S. Hochreiter and J. Schmidhuber. Flat minima. Neural computation, 1997.
H. Ibayashi and M. Imaizumi. Why does sgd prefer flat minima?: Through the lens of dynamical systems. In When Machine Learning meets Dynamical Systems: Theory and Applications, 2023.
A. Immer, M. Bauer, V. Fortuin, G. Rätsch, and M. E. Khan. Scalable marginal likelihood estimation for model selection in deep learning. In International Conference on Machine Learning (ICML), 2021.
S. Jastrzebski, M. Szymezak, S. Fort, D. Arpit, Tabor, K. Cho and K. Geras* The break- even point on optimization trajectories of deep neural networks. In International Conference on Learning Representations (ICLR), 2020.
C. Jin, R. Ge, Netrapalli, S. M. Kakade, and M. I. Jordan. How escape saddle points efficiently. In International Conference on Machine Learning (ICML), 2017.
D. S. Kalra, T. He, and M. Barkeshli. Universal sharpness dynamics in neural network training: Fixed point amalysis, edge of stability, and route to chaos. In International Conference on Learning Representations (ICLR), 2025.
M. E. Khan and H. Rue. The Bayesian learning rule. J. Mach. Learn. Res. (JMLR), 2023.
M. E. Khan, D. Nielsen, V. Tangkaratt, W. Lin, Y. Gal, and A. Srivastava. Fast and scalable Bayesian deep learning by (ICML), 2018.
in adam. In International Conference on Machine Learning
I. Kreisler, M. S. Nacson, D. Soudry, and Y. Carmon. Gradient descent monotonically decreases the sharpness of gradient flow solutions in scalar networks and beyond. In International Conference on Machine Learning (ICML), 2023.
H. J. Kushner. Stochastic stability. In Stability of Stochastic Dynamical Systems. Springer, 1972.
S. Lee and C. Jang. new characterization of the edge of stability based on sharpness measure aware of batch gradient distribution. In International Conference on Learning Representations (ICLR), 2023.
Li, Z. He, Y. Li, Y. Wang, L. Shang, and X. Huang. Flat-LoRA: Low-rank adaption over a flat loss landscape, 2025.
W. Lin, M. Schmidt, and M. E. Khan. Handling the positive-definite constraint in the Bayesian learning rule. In International Conference on Machine Learning (ICML), 2020.
M. Long and P. L.. Bartlett. Sharpness-aware minimization and the edge of stability. Mach. Learn. Res. (JMLR), 2024.
S. Lotfi, M. A. Finzi, Y. Kuang, T. G. J. Rudner, M. Goldblum, and A. G. Wilson Non-vacuous generalization bounds for large language models, In International Conference on Machine Learning (ICML), 2024.
12

---

A. M. Lyapunov. The general problem of the stability of motion. International Journal of Control, 55(3)531-533, 1992.
K. Lyu, Z. Li, and S. Arora. Understanding the generalization benefit of normalization layers: Sharpness reduction. Advances in Neural Information Processing Systems (NeurIPS), 2022.
R. Mulayoff and T. Michaeli. Exact mean square linear stability analysis for sgd. Conference on Learning Theory (COLT), 2024.
Y. Nesterov al. Lectures on convex optimization, volume 137. Springer, 2018.
Y. Netzer, T. Wang, A. Coates, A. Bissacco, B. Wu, and A. Ng. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning 2011, 2011.
T. H. Nguyen, U. Simsekli, M. Gurbuzbalaban, and G. Richard. First exit time analysis of stochastic gradient descent under heavy-tailed gradient noise. Advances in Neural Information Processing Systems (NeurIPS), 2019.
A. Orvieto, H. Kersting, F. Proske. F. Bach, and A. Lucchi. Anticorrelated noise injection for improved generalization. In International Conference on Machine Learning (ICML), 2022.
A. Orvieto, A. Raj, H. Kersting, and Bach. Explicit regularization overparametrized models via noise injection. In International Conference on Artificial Intelligence and Statistics, 2023.
K. Osawa, S. Swaroop, M. Khan, A. Jain, R. Eschenhagen, R. E. Turner, and R. Yokota. Practical deep learning with Bayesian principles. Advances in Neural Information Processing Systems (NeurIPS), 2019.
A. Panigrahi, R. Somani, N. Goyal, and Netrapalli. Non-Gaussianity of stochastic gradient noise. orXiv:1910.0960 2019.
Phunyaphibarn, J. Lee, B. Wang, H. Zhang, and C. Yun. Gradient descent with Polyak's momentum finds flatter minima via large catapults. In High-dimensional Learning Dynamics 2024: The Emergence of Structure and Reasoning, 2024.
S. Reddi, M. Zaheer, Sra, B. Poczos, Bach, R. Salakhutdinov, and A. Smola. A generic approach for escaping saddle points. In International Conference Artificial Intelligence and Statistics (AISTATS), 2018.
M. Rudelson and R. Vershynin. Hanson- Wright incquality and sub-Gaussian concentration. arXiv:1306.287., 2013.
Y. Shen, N. Daheim, B. Cong, P. Nickl, G. M. Marconi, B. C. M. Raoul, R. Yokota, Gurevych, D. Cremers, M. E. Khan, and T. Möllenhoff. Variational learning is effective for large deep networks. In International Conference on Machine Learning (ICML), 2024.
U. Simsckli, L. Sagun, and M. Gurbuzbalaban. A tail-index analysis of stochastic gradient noise in deep neural networks. International Conference Machine Learning (ICML), 2019.
Smith, Elsen, and De. On the generalization benefit of noise in stochastic gradient descent. In International Conference on Machine Learning (ICML). 2020.
L. Smith and Q V. Le. Bayesian perspective on generalization and stochastic gradient descent. arXiv:1710.06451, 2017.
R. Vershynin. High-dimensional probability: An introduction with applications in data science. Cambridge University Press, 2018
X. Wang, S. Oh, and C.-H. Rhee. Eliminating sharp minima from SGD with truncated heavy-tailed noise. In International Conference on Learning Representations (ICLR), 2022a.
Y. Wang, M. Chen, T. Zhao, and M. Tao. Large learning rate tames homogeneity: Convergence and balancing effect. International Conference on Learning Representations (ICLR), 2022b.
13

---

Y. Wang, Z. Xu, T. Zhao, and M. Tao. Good regularity creates large learning rate implicit biases: edge of stability, balancing, and catapult. arXiv:2310. 17087, 2023a.
Z. Wang, N. Ding, T. Levinboim, X. Chen, and R. Soricut. Improving robust generalization by direct PAC-Bayesian bound minimization. In IEEE Conference on Computer Vision and Pattern Recognition 20236.
M. Wei and D. J. Schwab. How noise affects the Hessian spectrum in overparameterized neural networks. arXiv:1910.00195, 2019.
Wu, V. Braverman, and J. D. Lee. Implicit bias of gradient descent for logistic regression the edge of stability. Advances in Neural Information Processing Systems (NeurIPS), 2024.
L. Wu, C. Ma, et al. How sgd selects the global minima in wee-paaraeterize learning: A dynamical stability perspective. Advances Neural Information Processing Systems (NeurIPS), 2018.
L. Wu, M. Wang, and W. Su. The alignment property of SGD noise and how it helps select flat minima: A stability analysis. Advances in Neural Information Processing Systems (NeurIPS), 2022.
A. Zellner. Optimal information processing and Bayes's 's theorem. The American Statistician, 1988.
S. Zhai, T. Likhomanenko, E. Littwin, D. Busbridge, J. Ramapuram, Y. Zhang, J. Gu, and J. M. Susskind. Stabilizing transformer training by preventing attention entropy collapse. In International Conference on Machine Learning (ICML), 2023.
G. Zhang, L. Li, Nado, J. Martens, S. Sachdeva, G. Dahl, C. Shallue, and R. B. Grosse. Which algorithmic choices matter which batch sizes? insights from noisy quadratic model. Advances in Neural Information Processing Systems (NeurIPS), 2019.
H. R. Zhang. D. Li, and H. Ju. Noise stability optimization for finding flat minima: A hessian-based regularization approach. Transactions on Machine Learning Research, 2024a.
X. Zhang, A. Ghosh, G. Liu, and R. Wang. Improving generalization of complex models under unbounded loss using PAC-bayes bounds. Transactions on Machine Learning Research, 2024b.
W. Zhou, V. Veitch, M. Austern, R. Adams, and P. Orbanz. Non-vacuous generalization bounds at the ImageNet scale: PAC-Baye Bay-saa compression approach. In International Conference on Learning Representations (ICLR), 2019.
L. Zhu, C. Liu, Radhakrishnan, and M. Belkin. Quadratic models for understanding catapult dynamics of neural networks. In International Conference on Learning Representations (ICLR), 2024.
X. Zhu, Z. Wang, X. Wang, M. Zhou, and R. Ge. Understanding edge-of-stability training dynamics with minimalist example. In International Conference on Learning Representations (ICLR), 2023.
Z. Zhu, J. Wu, B. Yu, L. Wu, and J. Ma. The anisotropic noise stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In International Conference on Machine Learning (ICML), 2019.
14

---

Supplementary Material
The supplementary material contains the following appendix section.
Related Works
15
Proof of Theorem
16
Role of posterior samples quadratic
21
Loss Smoothing Beyond Local Quadratic Approximation
21
Discussion on stability and descent for Variational GD
25
Additional experiments across diverse datasets
25
A
Related Works
Variational Learning
Given a prior and a likelihood, variational inference approximates the posterior by optimizing the evidence lower bound within family of posterior distribution. The ELBO is optimized by natural gradient descent [Amari, 1998], in the hope that steepest descent induced by the KL divergence is a better metric to compare probability distributions, this gives to rise of class of algorithms called Natural Gradient Variational Learning. Several works have attempted to apply this to deep learning [Khan et al., 2018, Osawa al., 2019, Lin al., 2020]. Recently, Shen al. [2024] provide an improved version varlational Online Newton that largely scales and obtains state of the art accuracy and uncertainty at identical cost as Adam. A step of VL includes weight perturbation based on noise injection from the posterior distribution which has similarities to noise injection.
Regularization Effect of Noise
Injecting noise into the parameters within gradient descent has several desirable features such as escaping saddle points [Jin et al., 2017, 17, Reddi et al., 2018] and local minima Zhu et al. [2019], Nguyen al. [2019]. In fact has been widely observed empirically that the SGD noise has an important role to play to find flat minima. Noise with larger scale [Zhu et al., 2019, Nguyen et al., 2019, Zhang et al., 2019, Smith et al., 2020, Wei and Schwab, 2019] and heavy-tail [Simsekli et al., 2019, Panigrahi et al., 2019, Nguyen et al., 2019, Wang et al., 2022a] drives the optimization trajectories towards flatter minima. Orvieto et [2023] demonstrated through stochastic Taylor expansion that injecting Gaussian noise into parameters before gradient step implicitly regularizes by penalizing the curvature of the loss, while Orvieto al. [2022] showed that anticorrelated noise Perturbed Gradient Descent (PGD) specifically penalizes the trace of the Hessian. Zhang et al. [2024a] showed that injecting noise in opposite directions to the weight space leads to a better regularization of the trace of the Hessian.
Gradient Descent Edge of Stability
Cohen et al. [2021], building on the work of Jastrzebski et al. [2020] empirically showed that for full batch Gradient Descent (GD) with constant step-size (p), the operator norm of the Hessian (also termed as sharpness) settles in a neighborhood of 2/p. This threshold is termed as edge of stability (EoS) because gradient descent quadratic only converges if the sharpness is below 2/p. Strikingly in complex neural network landscapes, sharpness settling around the stability limit 2/p (instead of diverging) indicates that presence of self-stabilization mechanism Damian al. [2023] (which is absent in a quadratic) that regularizes the sharpness near 2/ p.p. The common occurrence of this phenomenon across various tasks, architectures and initialization has inspired substantial research on edge of stability. For example EoS has been studied across several non-convex optimization problems [Wang et al., 2022b, 2023a, Zhu et al., 2023, Chen and Bruna, 2023, Agarwala et al., 2023, Arora et al., 2022, Lyu et al., 2022, Ahn et al., 2024, Wu et al., 2024, Ghosh et al., 2025, Even et al., 2024, Chen et al., 2024, Kreisler et al., 2023, Kalra et al., 2025, Zhu et al., 2024] and across other descent based optimizers such as SGD [Lee and Jang, 2023, Andreyev and Beneventano, 2024], momentum (Phunyaphibarn et al., 2024], sharpness aware minimization (SAM) Agarwala and Dauphin [2023], Long and Bartlett [2024] and Adam/RMSProp [Cohen et al., 2022, 2025].
15

---

Rutmenira
((0)
m
102
Figure 10: Shows the mechanism through which VL finds flatter minima. GD (gray iterates) get stuck in a local sharp minima whereas VL (red iterates) escapes to flatter minima.
Stability of GD with Noise
Wu et al. [2018] analyze the linear dynamical stability of SGD, demonstrating that the batch size and the gradient covariance matrix impose an additional constraint, requiring the Hessian operator norm to be smaller than 2/p. Wu et al. [2022] extend this work by assuming aligmment between the gradient covariance matrix and the Fisher matrix, arguing that noise concentrates in the sharp directions of the landscape. Although neither study analyzes the EoS limit for SGD, they understanding the local stability of SGD near a global minimum. Lee and Jang [2023), Andreyev and Beneventano [2024] investigate the dynamics of SGD EoS and instead propose that different metric, called mini-batch aware sharpness, must be smaller than 2/p for stability. Since the actual sharpness is less than this mini-batch aware sharpness, sharpness itself must be smaller than 2/p. However, in our work, the perturbation in the weight-space and not on the gradient.
instrumental in
B Proof of Theorem 1
First, we prove sufficient condition for descent for variational GD in Theorem B. The following theorem states the condition on the eigenspectrum for which descent takes place in expectation. Theorem 3.1. Consider the VGD update (5) on the quadratic loss (2), then
[e(((m]+1)) - f(m.) <0 if X <<= 2
N Ci.t)
VF
(
for all
(6)
where C,
v,)?. (v,), and VF(.) denotes the Variational Factor given by
VF(2) = p. 13 sinh (1
aresinh
(3
(7)
Proof. We analyze the stability of the weight-perturbed gradient descent (GD) update applied to the quadratic loss
1 ((m) - Qm,
(13)
where Q E Iddxd is positive definite. The update rule is given by
1 N.
m m. lpg. where 9
Vl(m, EE), E e~~((0,)) E).
(14)
N
Since f quadratic, its gradient at perturbed point satisfies
ve(m; +) Vf(m,) + Qc.
(15)
16

---

Therefore, if e ~ (0, Σ), then
Vi(m, +e) N (Vé(m:), QΣQ),
(16)
by linear transformation of Gaussian variables.
As ÿ is the average of N₂ i..dd samples from this distribution, it follows that
((((m)) N, QIQ)
(17)
Now, the Taylor expansion (exact) of the quadratic loss around m, gives
f(m++1) = f(m.)) + Vé(m))) (m=+1 m.) + (m++1 m.)) m QQ(m+11 -
(18)
Substituting the update rule (14) into (18), get
0² 2 -9 Qg.
f(m++1) - E(m.)) = (V((m.))
(19)
We now observe that the change in loss, denoted by A = {(m++1) L {m,)), is random variable due to the stochasticity in the update. Using the expression from Equation (19), we can expand At into a deterministic (mean) and stochastic (fluctuation) component (R).
p² Al = ((m/+1) - ((m,) = - g+ +
2 g
Qg
(20)
1²
-p (g + () + (g+ 5)
Q + E)
(21)
p² g Qg
2²
E QE + 2
(-pgTe + p²g To + 2
1² (T Qf Elg Qε)
E[A
R
(22)
We write the perturbed gradient g g where ~ N (o, N. QIQ). Here g denotes the gradient evaluated at the posterior mean, that is, Vf(m.). This expresses the update the sum of deterministic component g and a random fluctuation E.
To ensure that the update leads to descent on average, we compute the expected change in loss and enforce the stability condition [[ < 0 with respect to the curvature.
Taking the expectation of the loss change derived in Equation (19), we obtain:
p²
A -pg T (I
P Qig + 2N,
Tr (ΣQ*)= <
(23)
Since the Hessian Q is symmetric, it admits cigendecomposition (Ae, vi) are the eigenvalue-eigenvvetor pairs. We expand the expression the eigenbasis Q. Note the following identities:
ΣL \vv.v where
g Di
. 'g =gg-- Qg===0 = Σi \vV. g)3.
TTEEQ)) < d a,
(Von Neuman's 's trace inequality).
Thus, we obtain an upper bound on the expected descent:
02 p(v g)2 + f ([v, (g(v)))² 2) &
B[ - -(mm)) < Σ
an
2
2N,
Σ v().v
17

---

The inequality is an equality for an isotropic Gaussian Σ. For anisotropic covariance matrix, the tightness of this inequality depends on the alignment of the Hessian and the covariance matrix.
To ensure descent expectation, we require ʃ(A,, v;) < for all i. We note that this is a sufficient condition for descent to take place. We further extend this Theorem to necessary condition in Theorem. Define
p² f(A.) = ppa; + A.a: +
2N, o.A.
2
g)² >00 0 This a cubic polynomial in x of the form where a; = (vTg)²
S(A) a bd c2
with:
pan <0,
Φ ai >
p²
c 2N, 01 > 0.
Since ('(A) = b+ 3c12 > 0 for all X > 0, the function is strictly increasing. Therefore, ensuring
f(A,)) < equivalent to requiring A; be smaller than the (unique) positive root of ((A)
By Cardano's 's method for solving cubics, the condition A = (1:) 3 real root. The root can be expressed as:
+
Taa 2c
> 0 implies a unique
2/
1/3
(1+
pGp2
1/3
p
1/3
A; =
V
+
+
1
1
1+
(24)
P
27
27
3
3
2 sinh
sinh-
(25)
2i
where z; N.a: di
N. (v.
Hence, the expected loss decreases if
3 (
3
0 <x <
2
sinh sinh -1
for all i.
Now that have established sufficient condition on the curvature matrix Q to ensure that the expected loss decreases, we next address the random variability of the actual loss change due to the stochasticity of the update. Specifically, we show that the random variable Al concentrates sharply around its expectation.
For simplicity, we assumed to Qto to be full rank. If Q has non-trivial null-space that is X; 0 for some d > 2 k, then we have v, g = v Qm, = \ m, = 0. This means there no component of the iterate in the null-space direction of Q and does not contribute to the descent objective.
This guarantees that, under the sufficient condition derived above, the actual loss decrease also holds with high probability, not just in expectation. The following lemma provides concentration bound for Al about its mean:
Lemma 3.2. In the same setting Theorem 3.1, when the expected loss at next iteration is smaller than the previous loss by some margin >> that is,
Ee (m < e(mi) - 5,
then -((())))((mm) < O occuus with probubility at least - 2 exp (-c min for constants 61 C2 > depending only on p. Q and E.
8N. (N./o2))
For GD, if vf mo = 0, = always holds. However, stochasticity ensures that * 0 even when an mp chosen such that it is orthogonal to some v.'s.
18

---

Proof. In Theorem 3.1, we showed that the descent step occurs in expectation, if the sharpness is less than the stability threshold. In this lemma, we show that the descent step occurs with high probability under the same condition. To show, this we use derive concentration inequality to show that the descent step concentrates around its expectation with high probability. Moreover, with larger MC samples N., the deviation is small. The proof occurs in the following steps:
1. Let e = g g denote the gradient noise which is a NoO 1 QEQ). The descent step has both linear and quadratic terms wrt E.
N.
2. To ensure the deviation of the descent A from its expectation E[A[] we analyze the concentration of both the linear term and the quadratic term. The concentration the linear term follows simply from the Gaussian tail. The concentration of the quadratic term is done using the Hanson Wright concentration inequality Rudelson and Vershynin [2013], Vershynin [2018].
3. The tail bounds from both the linear and the quadratic term is combined using the union bound which concludes the proof.
Step-1: Separating flcttuation and expectation term
The descent step Al as derived in Theorem 3.1 can be written as its expectation E|A{] and fluctuation R.
p² - k T AL- = =((m++1) 1(m,) = - g
--gg + +E) F (g TQ +6)
8²
1²
p²
ppgg'QQQ Qc & (d EEcc Qc))
-
2 +
c 2 g Qg E Qc + (-pg" c cc
TTQ
E[A
n
We show that the fluctuation R = At- -E[AA is small with high probability, i.e, it has a aub-Gaussian tail.
Step-2: Concentration bound on the fluctuation
We separate the fluctuation on linear and quadratic terms wrt e since applying concentration inequality on each term is different.
R = + Q, where L -p
---T--- E p'g" Qe and Q _ (E Qe - QeE).
Tail bound on I. using sub-Gaussin concentration: Since T. is linear function of the Gaussian vector 6, itself Gaussian. variance
of. N,
h (QEQ)h N.
where az az ::)QQQQ) h and h p02Qg pg. A standard Gaussian tail then implies that for any t t> 0,
t2 = 2 exp
Pr(1 2exp (- 207
2
Choosing t= 5 yields
2
Pr[[Z] 2 2cxp
2cxp (- 8al
(26)
Tail bound on Q using Hamnonn-rrigh Concentration inequality: For the quadratic term, note that
(eT Q EleT Qe)). The Hanson- Hanson-Wright Wright inequality states that sub-Gaussia vector (here
Gaussian) with sub-Gaussian norm bounded by K for any matrix A (in case A e Q), there
19

---

exists universal constant C1 >> 0 such that for any t >>< 0,
K AI K2a1AA12 3)
Pr Pr( (e(E"Ac Ae Acll > 1) - 2 exp C1 min
Since _N 1 N, QIQ), its sub-Gaussian norm satisfies K² I c2Xmaz &N QEQ)====-- c3. where
N,
we define 3 Amaz (QEQ) and C2 is another universal constant. Furthermore, we have K
N2
|A|
2 Q and | the inequalities
Q To substitute the sub-gausse gaussian norm K from the bound, we get the tail bound to be <3
12 K A
PN2 IQQ
and
I.N
AT
-1 mine Pr((Q1>1) < 2exp
N2
tN
))
3 Q |
23 IQ
Finally chosing 6g = max @ 2299 [Q]]] and f = 5. we get:
Pr( : 2xxp -1 min
(5)
(5)N.)
}
(27)
C3
C3
Step-3: Combining L and Q concentation using union bound
Since R L+ + Q. by the union bound and using equation (27) and (26) we get
Pr Pr((R] > 2) Pr (r ;) pP((Q))> ;
= Pr PP c) s 2xx (- 2No, sal
+ 2cxp -a min ( (9)2N3 (
(-a
-amin
c3
We combine the sum
single exponential tail by using the inequality exp(-XX exp(-Y) S 2 2exp(- exp(-min(Xx)
Now chosing d2
32 max((11, (al,, 03}
32 mmxafaa a,, 3 o QQ c288 2 [|Q (substituting oş) and d min 320, min(X, we we get
1 7} and us-
ing exp(-X] + exp( (Y) I
-d1 min (-do
{ N22 cN,
})
Pr R > e s 2xx
d2 d2
Assume the expected descent step is strictly negative by some margin 8 > i.e. E[AA <- < O.
Recall that R = Al- - E]. Under this assumption, if Al 2 0, then A - E|A 2 5. Hence
(A< 0} C {IR) 2 67.
We set e ö in the concentration bound
Pi( 2 c) s exp(- - di min
( N2 eN, }): dg d2
and obtain
Pr 2) 0) S expl exp(-d, -o di min
52 N'Ng
AN. da
d2
Therefore, with probability at least
2 exp di min
82 NN
}):
dg
dg
we have Al O. Thus, if the expected descent step is at most 8,, then Af is negative with high probability.
20

---

C Role of posterior samples on a quadratic
On a quadratic loss, the perturbed averaged gradient follows Gaussian distribution with variance inversely proportional to the number of posterior samples:
9 ~~ V((mm), 1 Q N,
(28)
For large N,. the dynamics of variational GD approach those of standard GD and exhibit similar stability characteristics. To examine the role of N.. we consider quadratic loss {(m) fm² and compare GD with variational GD across varying values of N,. When < 2/a, standard GD converges the minimum. In the limit N₂ oc, variational GD recovers this behavior. However, for finite N.. the gradient estimate becomes noisier, reducing the stability threshold as predicted by Theorem-l. Figurell shows the resulting trajectories and histograms of the iterates. As N, decreases, the iterates exhibit greater variability and cover wider range. reflecting increased instability. These results confirm the theoretical prediction that smaller N, increases the likelihood of the loss increasing in the next step, even under a stable learning rate.
(GD < 9
N₂ 1000
N, 200
N, 100
N₂ 50
N₂ 20
N₂
Figure 11: Histogram plot iterates on a quadratic for GD with noise injection from a posterior distribution with N, samples. As N, decreases, the iterates become more unstable.
We further visualize the stability threshold predicted by Theorem-1 1 12gure12. For each posterior sample size N., we compute the descent probability on quadratic loss and plot it alongside the corresponding stability threshold 2/p- VF. As shown the top row, descent occurs with high probability whenever the curvature is below the threshold-consiseeee bottom row, we binarize the descent probability by setting it to it exceeds 0.5, and 0 otherwise. The resulting transition boundary closely aligns with the predicted threshold, further validating our theoretical result.
with Theorem 1. In the
D Loss Smoothing Beyond Local Quadratic Approximation
In our paper, we consider local quadratic approximation of the loss to characterize the distribution of the perturbed gradient. By the property that Gaussian distribution invariant under linear transformation, showed that the perturbed gradient (or it's sample average) is also Gaussian. This observed by characterizing the distribution of the perturbed gradient Vé(m;; + e). Employing a Taylor expansion around the current mean parameter m, we obtain:
Vl(m + E) = Vé(m) + v=((m)e + ((())e, + O(Iee)), e~ (~,N E)
(29)
2
In general, the distribution of Vi(m + <) involves higher-order terms of Gaussian variables, making exact characterization challenging. However, if the perturbation covariance Σ is suffi- ciently small relative to the local curvature, defined by the Hessian H(m) = V²((m), specif-
ically, when ||V&8(m)]lop |E & ||H(m))l2, the gradient approximation effectively simpli- fies to linear approximation Vl(m E) V((m) H(m)e. This approximation is exact
21

---

for a quadratic. Under this condition, since e (O, the perturbed gradient is Gaussian Ve(m + e) ~ ((()(m)) H(m))H((m)) [
However, in perturbeturbtiion regimes (large covariance M where [ || the perturbed gradient is no more Gaussian (since it involves second-order terms on e). Furthermore, the expectation of the perturbed gradient has contribution from the neighborhood of the loss, which is referred to here as third-order curvature bias.
(H(m)))2)
B[ve(m + c)] R0 Ve(m) + 2r,,3 (233( +0(1=)13)
(30)
third-order curvature bias
Having large number of posterior samples, can help recover this modified expectation better with increasing number of samples. Here, the algorithm effectively follows the gradient of a smoothed version of the loss landscape, since samples are drawn from a wide neighborhood around m. Accurately estimating this biased but meaningful descent direction under heavy-tailed noise requires a sufficiently large N,, not just to reduce variance, but to faithfully recover the expected smoothed gradient.
In the next section and the subsequent theorem, we study this phenomenon in a non-quadratic function. Here show that, first for quadratic function, smoothing does not change the curvature of the underlying loss. But for a quartic function, smoothing by expectation changes the underlying curvature by the variance of the posterior, and furthermore, for smoothing by finite averaging, the curvature changes as function of both the variance and the number of samples used to approximate the expectation.
Let E(O) = afl + 60 + and let = (0 o²), then
lcono (m) = E, = (ab²² + b0 + +c) = am? + bm + c ag²
So w.r.t the reparameterized variable m, the curvature of the loss remains unchanged, only shift occurs, proportional to o². So, the stability dynamics on E(0) and (eonv(m) are the same. The loss diverges only when the learning rate is p > 2. However, this is not the case with general losses, especially for losses where g(4) (0) + 0, ie if it has non-zero fourth order derivative.
For example, let's take the example of quartic function E(O) (02-1) 1)2 where minima is at 0 +1 and (0(0) 120 4 8. Now, the smoothed loss Leonv (m) (E) E6 (02 1)2 me + m2 (6o2 2) + (3or 2a² + 1). The new loss has a global minima at m +v1 30
N₂ =1
N₂ 2
N₂ = 5
N₂ = 10
N 50
T
N, =1
N₂ 2
N₂
N, 10
N
50
Figure 12: (Top Row): Probability of descent evaluated over 10 trials on several simulations of combination of quadratic curvature and noise variance. (Bottom Row): Hard thresholded probability for descent. The transition occurs near the boundary given by the derived expression of 2/p VF from T1eorem-1.
22

---

(only o < v Comparing the curvature of both the losses at the global minima, we have:
(0) 126 4 8
Lomn (m) 12m² + 2(662 - = 12m2 4 12o = 246
So, the smoothing operator, changes the curvature at the global minima for quartic functions. note that this wasn't the case for quadratic functions, where the curvature was unchanged. This further changes the stability dynamics since the learning rate required for stable convergence is difference,
for original loss {()). it is E. whereas for the smoothed loss, is p 8 2-24 However, consider the effect of averaging of loss using finite mc samples N.. lang( (0) E~ N((000))
No E(O + Ez), where
N,
Lavg(0) = N.
Σ i=1
E(O + +.)
1
(0 61) - -1)2
N₂
N. N₂
N₂
1) 11) -1)
(
Nr
= gt 8
(6e, + de, -2)
40
1
Σ
cl
+
Σ (e
- 22 + +1)
N₂
N₂
So, the second derivative becomes:
2 N,
Capg (0) 120
+
Σ N
(GE + 4e;)
If we compare the second derivative of the smoothed out-loss Coon (m) and averaged loss over N, samples t" (9), observe that wrt samples, the expectation is same, i.e,
2 N. N₂
Er, (0)
(0) 120 4 Eeg
Σ
[ + de:)] = 120 4 12o²
(oooo (0)
However, the pointwise deviation of lang (0) and Loon (0) depends on the number of samples N, Taking the difference we have:
lag (0) (oon (0)
2
Σ (6 + +46) 607
N.
Define
Y = ( + 46)) - 6
Since Bie²] o² and [[[.]] O, each Y, has mean zero:
EYY, 0.
Hence,
2 N
Log (0)
E" Con (0)
Y,
N.
Note that e; ~ o2) is aub-Gaussian, and e, is sub-exponential. Thus Y, = =6,= + 4c; is linear combination sub-exponential and sub-Gaussian variable, which remains sub-exponential. Shifting by the constant - 6 o does not affect sub-exponential parameters. Therefore, each Y; is sub-exponential.
Let (v, b) be sub-exponential parameters for Y. By standard Bernstein-type concentration for sub-exponential random variables, there exists universal constant such that for all / 0:
N.
Yg
21
S 2 exp - N., min { 1 1 w 1)
23

---

Since
2 N
(0) E Loon (0)
Σ
Yi,
No i=1
we have
N Σ
N
(0)g(0) (0) - E (0)
=
N
Yi
=
N,
1
Y
1=1
Hence, for any 8>0 0,
N
(a (0)-Lcone (0) 2 5) = Pr
Σ
Y
2
exp - cN,, min { (8/2) e b
4/2
})
N.
i=1
Plugging in c and b = c202 which are the subexponential norms in terms of the variance, we get for some constant 61 and 02:
Pr(1aacgg
8) S expl NN, min 6a (one (0) (oon
Zoja 20yo }).
This inequality shows that the sinite-sample second derivative (arg(0) concentrates around the infinite- infinite-sample second derivative ( (6) at an exponential rate in N While they are identical in expectation, the above result quantifies their pointwise deviation with high probability. We formalize this observation for general analytical function which has continous derivatives.
Theorem D.1 (Concentration of Smoothed Curvature). Let f T - TR. be six-times contin- uously differentiable analytical function, with all derivatives up to order 6 bounded. Suppose o² sup, f(0) (x) < ssup, (x) and for ii.d samples {G with e; ~ NN0 o²), define
1 N.
favg (x)
Σ ((e+e)). fconv (x) = =EZ~N(0,nz) [f(r]]](]]] + 2)].
N.
Then there exist universal constants c, C12 02 O such that for any 8 > 0.
Pr( [( ((onv(a)) 8 < expl cN, min 82 4cc 2c2
Hence, | (x) - ('onv( (x)) concentrates around zero at an exponential rate in N
A general extension of this proof can be done for functions which has a finite fourth order derivative and small sixth order derivative such that x² sup,, (z) < sup, f (x). Starting with sufficiently smooth function f E Crk, we can expand f(x + 2) in a Taylor series around x. Let N/ o²). Then:
z²
2x = ((x) + = fx) +
(x) + 3!
(x) +
21
Since is Gaussian with zero mean, all the odd moments E[2], E| etc., vanish. Thus, when we take the expectation,
Sconv (x) = E E=-N(0.0)) =)] = f(x) + d²
(" (z) + + at (x) + 06 (z) +
2
8
48
In the above, use the fact that E[2?] 0². E[ 3 E| 15 08 etc., and only the even powers contribute. The curvature of the original function then becomes:
(x)= + σ²
2 M((x) + dt f((x)
Under the condition that δ² sup, (x) suup, f (I), we get
Poan (अ) (z) ~ (x) 2 (z)
(31)
The smoothing effect does change the second order curvature of the loss. Furthermore, approximating This approximation makes the noise due to the N, samples be subexponential and similar result holds.
24

---

E Discussion on stability and descent for Variational GD
On Necessity and Sufficiency of descent:
A further critical difference lies in the nature of the descent conditions for GD versus VGD. For GD on a quadratic loss, the standard descent lemma provides condition on the maximum eigenvalue Amaz < 2/p p that is both necessary and sufficient for ensuring stability and monotonic descent. A violation of this single condition guarantees divergence.
For VGD, the analysis is more nuanced. The true necessary and sufficient condition for the expected loss decrease is that the sum of all components its cigen-dccomposition must be negative, as shown below:
P 2 2N.
E A = PP P
(I-
P
2Qg + Tr (ΣQ*) = Σ <0
(32)
==1
However, analyzing this sum can be intractable. Instead, our work derives a practical sufficient condition by requiring each term in the sum to be negative independently, i.e., f(A,, v.) < 0 for all i. This is the condition presented in Theorem 1, which yields the stability limit X; < 2/0 VF(z;) for each eigenvalue.
This theoretical framework is validated by our extensive experiments on MLPs and ResNets across various learning-rate and variance. We consistently find that the leading Hessian eigenvalues, A hover around their respective stability thresholds. 2/p- VF(2.). This alignment provides strong empirical support for our sufficient condition, showing an active constraint that accurately describes the behavior of VGD in practice.
Mulayoff and Michaeli [2024] also derives stability condition (see their Theorem 5) which is necessary and sufficient condition for the stability of SGD, but their setting is fundamentally different from ours. Firstly, in their setting, the assumption is that the batch-size is drawn uniformly at random from the finite set of all possible data samples. Our work, in contrast, does not model noise from data sampling but rather from perturbations the model's weights, which we assume follow an explicit, continuous distribution (c.g., Gaussian). This leads to fundamentally different noise structures. In our VGD framework, the resulting gradient estimator shown to follow normal distribution whose covariance, N. QEQ, is shaped by the Hessian (Q), meaning noise is amplified in directions of high curvature. In the SGD paper, the gradient noise has discrete distribution with a covariance determined by the dataset's 's intrinsic variance.
F
Additional experiments across diverse datasets
While in the manuscript, we presented experiments on CIFAR-10, here we present several additional results to support our theorem and claims.
Gaussian Variational GD: In Figure 13, we compare GD and GD with weight perturbation trained across three different architectures conclusion that with Gaussian weight perturbation achieves smaller sharpness and better test accuracy than just GD. Similar trend is also observed for FashionMNIST dataset in Figure 15. In Figure 16, we show that in deep neural networks, sharpness also depends on the number of posterior samples, as smaller samples lead to smaller sharpness. In Figure 14, plot the Variational factor along with the sharpness in ResNet and show that they match closely.
SVHN dataset Netzer al. [2011]. We arrive at similar
Sharpness comparison for Cross Entropy Loss:
For cross-entropy (CE) loss, the sharpness dynamics differs from those observed with mean-squared error (MSE) loss. As training progresses and model predictions become increasingly confident, the term (1 in the Hessian approaches zero, driving the curvature and hence sharpness to vanish. This causes both GD and VL (or IVON) to converge to regions with negligible sharpness, although their transient behaviors differ. VL's sharper stability bound results in smaller peak sharpness compared to GD.
Figure 17 illustrates this effect for fully connected neural network trained on CIFAR-10 with CE loss. The left panel shows the training loss while the right panel presents the evolution of sharpness over training steps. Although both methods eventually exhibit vanishing curvature, IVON achieves
25

---

GD
GD with Gaussian weight perturbation
2/p
Train Accuracy
Train Accuracy
Train Accuracy
1.00
1.00
1.00
0.75
0.75
0.75
0.50
0.50
0.50
0.25
0.25
0.25
25000 Train steps
50000
25000 Train Train steps
50000
10000
20000
Train steps
Test Accuracy
Test Accuracy
Test Accuracy
0.6
0.75
0.4
0.4
0.50
0.25
25000 Train steps
50000
25000 steps
50000
10000 Train steps
20000
Sharpness
Sharpness
Sharpness
30
400
40
20-
200
30
10
25000 Train steps
50000
25000
50000
10000
20000
Train steps
Train steps
(a) ViT
(b) MLP
(c) ResNet-20
Figure 13: Similar to the trend in Fig. 4, weight perturbed GD consistently finds flatter minima on the SVHN dataset.
substantially lower peak sharpness than GD. This highlights the improved stability and flatter convergence landscape induced by variational learning even under the CE loss.
perturator from heavy-tail distribution: In the manuscript, we presented results on perturba- tion from a heavy-tailed Student-t distribution for MLP trained on CIFAR-10. Figure 18. we plot the training dynamics across ViT and ResNet-20 models. Here, we observe that heavier tails (smaller a) leads to smaller sharpness and better test accuracy.
Experiments in Vision Transformers: In Attention-basel architectures, such as Vision Transformers, has been widely observed that sharpness for GD often goes above the stability threshold 2/p. This phenomenon has been widely studied [Zhai et al., 2023] as attention entropy collapse that makes training unstable in Vision Transformers, with sharp spikes in both the training and accuracy. To mitigate such unstable behaviour, weight perturbation can be used, due to its sharpness reducing effect and thereby stabilizing training. For example, in Figure 20, we perform an ablation study of training ViTs with different perturbation covariance. Noise with larger variance consistently leads to more stable training and smaller sharpness in ViT. Similarly, for VON we observe that preconditioned sharpness is smaller than 2/p for ViT training 19.
Experiments in NLP tasks: To verify that Variational GD with isotropic Gaussian noise achieves lower sharpness compared to GD on an NLP task, we add classification head to frozen BERT-mini backbone and finetune the head on SST-2. For both GD and Variational GD, we use the same learning rate of 0.05. We report train loss, validation loss, and sharpness in Figure 21 and observe lower sharpness for Variational GD throughout.
26

---

Normalized- Normaliiee-Sharppess (a = 0.5)
Normalized Normalizze-Sharpness (a= 0.1)
Variational Factor (a 0.5)
Variational Factor (ar 0.1)
1.100
1,300
0.333
0.738
0.000
10
20
20
10
40
10
10
20
10
Epoch
(a) p 0.01
(b) P 0.02
(c) p 0.05
Figure 14: Normalized Sharpness (m22(m))l2/(2//))) hovers about the Variational Factor in ResNet.
GD
GD with Gaussian weight perturbation
2/p
Train Accuracy
Train Accuracy
Train Accuracy
1.0
1.0
1.0
0.9
0.9
0.9
0.75
0.7
0.60
25000 Train steps
50000
0.6
5000 Train steps
10000
0.6
10000 Train steps
20000
Test Accuracy
Test Accuracy
Test Accuracy
0.9
0.9
0.9
0.8
0.8
0.8
0.6
25000 Train steps
50000
0.6
5000 Train steps
10000
0.6
10000 Train steps
20000
Sharpness
Sharpness
Sharpness
12
20
200
10
10
25000
50000
5000
10000
10000
20000
Train steps
Train steps
Train steps
(a) ViT
(b) MLP
(c) ResNet-20
Figure 15: Similar to the trend in Fig. 4, variational learning finds flatter minima on the FashionM- NIST dataset.
Weight perturbation, N₁=1
Weight perturbation, N2 6,=6
2/p
Train Accuracy
Test Accuracy
Sharpness
0.4
40
0.3
0.2
0.1
5000
10000
5000
10000
5000
10000
Train Train steps
Train steps
Train steps
Figure 16: Using more samples per iteration leads to a large sharpness, as demonstrated in our ablation study where MLP network is trained on subset of the CIFAR-10 dataset containing 10000 images.
27

---

GD (p 0.05)
IVON (a? = 0.011)
IVON (o² = 0.05)
Trin Loos
Sharpness
-0
1030
2000
3000
4000
5000
500
1000
1500
2000
2500
Train Steps
Figure 17: Training dynamics of fully connected neural network trained with GD and IVON with fixed covariance. Although the sharpness always drops to zero fro CE loss, IVON achives smaller peak sharpness than GD.
a=3
a=6
o (Gaussian perturbation)
2/p
Train Accuracy
Test Accuracy
Sharpness
0.6
20
0.6
0.4
0.4
0.2
10000
20000
10000
20000
10000 Fain steps
20000
Train steps
Train steps
(a) ResNett-0
Train Accuracy
Test Accuracy
Sharpness
200
0.8
0.4
150
0.3
100
0.2
50
0.1g
25000 Train steps
50000
25000
50000
25000
50000
Train steps
aii steps
(b) VIT
Figure 18: We notice similar trends shown in Fig. 7 in training ViT and ResNet models. That is to say, smaller corresponds to deavier-tailess posterior which leads smaller sharpness. Note that as a approaches infinity, the Gaussian posterior is recovered.
Adam
Weight perturbation. 0.5.
Weight perturbation, T 0.1
2/p
Train Accuracy
Test Accuracy
Pre-conditioned Sharpness
1.0
0.4
10
0.3
10
25000
50000
25000
50000
25000
50000
Train steps
Train steps
Fain steps
Figure 19: We notice similar trends shown in Fig. training VIT models. That is to say, Smaller temperatures which shrinks the posterior reaches larger preconditioned sharpness.
28

---

Weight perturbation, of -
Weight perturbation,
Weight perturbation on, a'
2/p
Train Accuracy
Test Accuracy
Sharpness
500
0.8
400
200
100
50000
100000 0.3g
50000 Train steps
100000
50000
100000
Train
steps
Figure 20: For VIT, GD training is prone to loss spikes, and the sharpness often goes above the stability threshold 2/p p. On the other hand, with weight perturbation the training becomes more stable, and larger noise variance consistently leads to smaller sharpness.
Train Accuracy
Test Accuracy
Sharpness
0.7
40
0.6
35
0.6
30
0.5g
20000
40000
20000
40000
20000
40000
Train steps
Train steps
Train steps
Figure 21: SST-2 head finetuning frozen BERT mini backbone comparing Variational GD with isotropic Gaussian noise to vanilla GD, both with learning rate 0.05.
29

---

Normalized Normallzee-Sharppeess
Variational Factor
1.00
1,00
0.75
0.75
0.50
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(a) A
(b) A
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
100
300
400
500
100
400
500
Epoch
Epoch
(c) Àa
(d)
0.75
0.75
0.50
0.50
0.25
0.25
0.00
100
200
300
400
500
0.00
100
200
300
400
500
Epoch
(c) As
(f) As
1.00
1.00
0.75
0.75
0.50
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(g) X7
(h) As
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
100
200
300
400
500
100
300
400
500
Epoch
Epoch
Ag
(j) Auo
Figure 22: Eigenspectrum (A1 -1o) vs corresponding stability threshold 2/p- VF(z.). 2-layer ML. trained with VGD = 0.05 and o² 0.1.
30

---

Normalized NormallzeedShhrppeess
Variational Factor
1.00
1,00
0.75
0.75
0.50
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(a) A
(b) A
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
300
400
500
100
300
400
500
Epoch
Epoch
(c) Àa
(d)
0.75
0.75
0.50
0.50
0.25
0.25
0.00
100
200
300
400
500
0.00
100
200
300
400
500
Epoch
(c) As
(f) A
1.00
1.00
0.75
0.75
0.50
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(g) X7
(h) As
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
100
200
300
400
500
100
400
500
Epoch
Epoch
Ag
(i) Auo
Figure 23: Eigenspectrum (A1 -1o) vs corresponding stability threshold 2/p- VF(z.). 2-layer ML.P trained with VGD = 0.02 and o² = 0.5.
31

---

Normalized NormallzzedSShrppeess
Variational Factor
1.00
1,00
0.75
0.75
0.50
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(a) A
(b) A
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
100
300
400
500
100
300
400
500
Epoch
Epoch
(c) Àa
(d)
0.75
0.75
0.50
0.50
0.25
0.25
0.00
100
200
300
400
500
0.00
100
200
300
400
500
Epoch
(c) As
(f) As
1.00
1.00
0.75
0.75
0.50
0.25
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(g) X7
(h) As
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
100
200
300
400
500
100
300
400
500
Epoch
Epoch
Ag
(j) Auo
Figure 24: Eigenspectrum (A1 -1o) vs corresponding stability threshold 2/p- VF(z.). 2-layer ML.P trained with VGD p = 0.02 and o² 1.0.
32

---

Normalized NormallizedSShappeess
Variational Factor
1.00
1,00
0.75
0.75
0.50
0.50
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(a) A
(b) A
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
100
300
400
500
100
300
400
500
Epoch
Epoch
(c) Àa
(d)
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
100
200
300
400
500
0.00
100
200
300
400
500
Epoch
(c) As
(f) A
1.00
1.00
0.75
0.75
0.50
0.25
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(g) X7
(h) As
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
100
200
300
400
500
100
300
400
500
Epoch
Epoch
(i Ag
(j) Auo
Figure 25: Figenspectrum (A1 -1o) vs corresponding stability threshold 2/p- VF(z+). 2-layer ML. trained with VGD = 0.1 and x² 1.0.
33

---

Normalized Normalized-Sharpnesss
Variational Factor
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(a) A
(b) A
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
100
300
400
500
100
300
400
500
Epoch
Epoch
(c) Àa
(d)
1.00
1.00
0.75
0.75
0.50
0.25
0.25
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
(c) As
(f) As
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(g) X7
(h) As
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
100
200
300
400
500
100
300
400
500
Epoch
Epoch
(i Ag
(i) Auo
Figure 26: Eigenspectrum (A1=)10) vs corresponding stability threshold 2/p VF(2.). ResNet-20 trained with VGD = 0.1 and x² 0.1.
34

---

Normalized Normalized-Sharppesss
Variational Factor
1.5
1.00
0.75
0.50
0.5
0.25
0.0
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(a) A
(b) A
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.00
100
300
400
500
100
300
400
500
Epoch
Epoch
(c) Aa
(d)
1.00
1.00
0.75
0.75
0.50
0.25
0.25
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
(c) As
(f) A
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(g) A
(h) As
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
100
200
300
400
500
100
300
400
500
Epoch
Epoch
(i Ag
(i) Au
Figure 27: Eigenspectrum (A1=)1o) vs corresponding stability threshold 2/p VF(2.). ResNet-20 trained with VGD = 0.1 and x² 0.5 (low variance).
35

---

Normalized Normalizee-Sharppeess
Variational Factor
1.00
1.00
0.75
0.75
0.50
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(a) A
(b) A
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
100
300
400
500
100
300
400
500
Epoch
Epoch
(c) Àa
(d)
0.75
0.75
0.50
0.50
0.25
0.25
0.00
100
200
300
400
500
0.00
100
200
300
400
500
Epoch
(c) As
(f) A
1.00
1.00
0.75
0.75
0.50
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(g) X7
(h) As
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
100
200
300
400
500
100
300
400
500
Epoch
Epoch
(i Ag
(i) Auo
Figure 28: Eigenspectrum (A1=)10) vs corresponding stability threshold 2/p VF(2.). ResNet-20 trained with VGD p = 0.05 and o² = 0.1.
36

---

Normalized Normalizee-Sharppeess
Variational Factor
1.00
1.00
0.75
0.75
0.50
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(a) A
(b) A
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
100
300
400
500
100
300
400
500
Epoch
Epoch
(c) Àa
(d)
0.75
0.75
0.50
0.50
0.25
0.25
0.00
100
200
300
400
500
0.00
100
200
300
400
500
Epoch
(c) As
(f) A
1.00
1.00
0.75
0.75
0.50
0.00
0.00
100
200
300
400
500
0
100
200
300
400
500
Epoch
Epoch
(g) X7
(h) As
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
100
200
300
400
500
100
300
400
500
Epoch
Epoch
Ag
(i) Auo
Figure 29: Eigenspectrum (A1=)10) vs corresponding stability threshold 2/p VF(2.). ResNet-20 trained with VGD = 0.05 and o² = 0.5.
37

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Thorough experiments show that the observations made in the theorem per- fectly show up in experiments,
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
We state that our work is only limited to weight perturbed GD but not precon- ditioning.
Justification:
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model cell-speciiiitiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
38

---

Answer: [Yes]
Justification: Theorem-1, which states the main result of the paper has been validated extensively through experiments. For example in Figure-2, the theoritically calculated thereshold matches the experiment.
Guidelines:
The answer NA means that the paper does not include theoretical results,
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The code and the details of reproduction such as hyperparameter setting has been provided.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
39

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We provide anonymous code for reviewers to check. Once paper is published we will release the original code.
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: https/scccbclccgcuicceaeeaecpppa SCaeee//doeeSmmeeeeeeeee ssionPoiccy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: Yes, all these details have been stated explicitly.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Yes, averaged over various random seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
40

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Yes, states the computer resources.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: / //wuurrurrp.../ cc
Answer: [Yes]
Justification: Yes, it does.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Impact section added.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
41

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Yes, all code sources cited and credited.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
42

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We added Readme file document our code reviewers can check them.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [No]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
43

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [Yes]
Justification: Not used. only grammar checks and writing.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (https ftttps:/nneurpst c for what should or should not be described.
44


Reviews:
[{"review_id": "xhb5zPUJLz", "reviewer": "Reviewer_6LwC", "summary": "The 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.", "questions": "In 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$)?", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "## 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)", "was_revised": true, "final_justification": "I've read the other reviews and didn't feel like I missed any substantial flaws, which is why I remain my rating of \"strong accept\".\n\nAs far as I understand, the main remaining criticism by another review are partially due to some misunderstandings of VL, and partially point out that the assumptions made for the theoretical derivations in the paper do not perfectly match the experimental setup. In my opinion, the latter is the only way how theory of deep learning can make any progress: without strong simplifying assumptions, I doubt that any nontrivial insight can be derived theoretically. But by clearly stating the assumptions made for theoretical derivations, we can identify which aspect of the (inevitably much more complicated) experimental setups can explain observed phenomena. Further, the paper does compare the theoretical findings to experiments, and I think that it convincingly shows that empirical findings do indeed match the derived theory.", "rebuttal": "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.\n\n**Q1**: Fixing the reference issue in numbering equations and sections. \n\n**A1**: Thanks for catching this. In the updated manuscript, we will properly refer to whether it is equation or section. \n\n**Q2**: Can you plot (in addition to the $2/\\rho$ line) also a line that includes the \"VF\" factor?\n\n**A2**: Thanks for the suggestion! We will add it in the final version."}, {"review_id": "KguQgjrf5X", "reviewer": "Reviewer_YUMo", "summary": "The 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.", "questions": "As 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-", "limitations": "None", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Overall, 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.", "was_revised": true, "final_justification": "The authors addressed most of my concerns and, hence, I raised my score to 5.", "rebuttal": "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.\n\nWe 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:\n\n1) 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.\n\n2) 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).\n\n      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.\n\n3) We will add a discussion on extending MSE to cross-entropy etc. We will also add some numerical results.\n\n      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. \n\n\n| Iteration            |        GD       | VL (var=0.1) | VL (var=0.05) |\n|-------------------|--------------|---------------|-----------------|\n| 0-2000              |         45       |         26.       |         10           |\n| 2000-4000        |         40       |         22        |         10           |\n| End of training  |          8        |          3         |         1             |\n\n\n4) 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”.\n\n     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. \n\n| Iteration            |        BS=512     | BS=256   | BS=128 |\n|-------------------|--------------|---------------|-----------------|\n| 0-25k              |         37.6       |         22.3    |         15.4        |\n| 25k-40k          |         28.2       |         10.2    |         3.4      |\n\n\n[1] Alquier et al, “User-friendly introduction to PAC-Bayes bounds”.\n\n\n[2] Cohen et al, “Adaptive Gradient Methods at the Edge of Stability”. https://arxiv.org/abs/2207.14484\nMany thanks for pointing out typos. We will fix them!\n\nMany thanks for pointing out typos. We will fix them!"}, {"review_id": "3cyMMR7Ohx", "reviewer": "Reviewer_BkQ3", "summary": "Variational 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.", "questions": "See 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.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## 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.", "was_revised": true, "final_justification": "All my concers are resolved", "rebuttal": "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). \n\n**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. \n\n**A1**. Thanks for the suggestion. We will move Eq. 7 within the theorem.\n\n**W2**. Mathematical link between Theorem 1 and Hypothesis 1 is unclear.\n\n**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.\n\n**W3**: Hypothesis-1 assumes isotropic Gaussian noise, whereas Theorem-1 does not; the paper does not explain this discrepancy.\n\n**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.\n\n**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?\". \n\n**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. \n\n**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.\n\n**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.\n\n**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.\n\n**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.\n\n**W7**: NLP experiments that follow the theoretical Gaussian setting would strengthen the paper. \n\n**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.\nHere we show the results:\n\n**Variational GD with isotropic Gaussian noise**\n| Iteration | Train Loss | Val Loss | Sharpness |\n|-----------|------------|----------|-----------|\n| 10        | 0.6838     | 0.7149   | 29.2796   |\n| 20        | 0.6760     | 0.6715   | 30.7939   |\n| 50        | 0.6513     | 0.6803   | 29.9545   |\n| 100       | 0.6266     | 0.6598   | 34.0424   |\n| 200       | 0.6140     | 0.6159   | 33.9855   |\n| 500       | 0.5675     | 0.6119   | 35.6071   |\n| 1000      | 0.5657     | 0.5926   | 35.8911   |\n| 2000      | 0.5334     | 0.5806   | 34.4932   |\n| 5000      | 0.5232     | 0.5754   | 31.5243   |\n| 10000     | 0.5158     | 0.5644   | 31.7331   |\n| 20000     | 0.4924     | 0.5461   | 33.1990   |\n| 50000     | 0.4596     | 0.5320   | 32.1427   |\n\n**GD**\n| Iteration | Train Loss | Val Loss | Sharpness |\n|-----------|------------|----------|-----------|\n| 10        | 0.6834     | 0.6946   | 33.0345   |\n| 20        | 0.6735     | 0.6866   | 33.3173   |\n| 50        | 0.6502     | 0.6669   | 36.1241   |\n| 100       | 0.6183     | 0.6378   | 44.8439   |\n| 200       | 0.6192     | 0.6196   | 40.6111   |\n| 500       | 0.5766     | 0.5983   | 40.2430   |\n| 1000      | 0.5548     | 0.5882   | 40.1827   |\n| 2000      | 0.5414     | 0.5821   | 40.1149   |\n| 5000      | 0.5284     | 0.5741   | 40.0432   |\n| 10000     | 0.5138     | 0.5605   | 40.1765   |\n| 20000     | 0.4922     | 0.5477   | 40.3987   |\n| 50000     | 0.4570     | 0.5359   | 40.7298   |\n\n\n[1] Cohen, Jeremy M., et al. \"Gradient descent on neural networks typically occurs at the edge of stability.\", ICLR-21. \n\n[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. \n\n[3] Long, Philip M., and Peter L. Bartlett. \"Sharpness-aware minimization and the edge of stability.\" JMLR-24\n\n[4] Damian et al, \"Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability.\" ICLR-23. \n\n[5] Devlin, Jacob, et al. \"BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.\" NAACL-19\n\n[6] Bhargava, Prajjwal, Aleksandr Drozd, and Anna Rogers. \"Generalization in NLI: Ways (Not) To Go Beyond Simple Heuristics.\" arXiv\n\n[7] Socher, Richard, et al. \"Recursive deep models for semantic compositionality over a sentiment treebank.\" EMNLP-13\n\n\n**Q1**. How does VL being less stable enables it to find flatter minima?\n\n**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.\n\nConsider 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.\n\n**Q2:** Use of isotropic assumption. \n\n**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:\n\n*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.\n\n*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.\n\nOur 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.\n\nWe 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.\n\n**Q3** is addressed with **W4**. \n\n**Q4** Numbering issue. \n\n**A**: Thank you for spotting the error in section number. We will fix it in our update.\n\nIf the reviewer's concerns are addressed, we would be happy to see an increase in score."}]

Author rebuttal:

Meta-review: This paper analyses the training dynamics and implicit regularisation of variational inference for deep neural networks using Edge of Stability theory, showing that variational inference leads to flatter solutions than gradient descent. Reviewers found the paper exceptionally well written, with novel and sound theoretical contributions and well-designed experiments. While concerns were raised about theorem clarity, the connections between theory and experiments, and the scope of evaluation, the authors addressed these during rebuttal, and I recommend acceptance.
Decision: Accept (spotlight)