---
title: Flatness-Aware Stochastic Gradient Langevin Dynamics
url: https://www.emergentmind.com/topics/flatness-aware-stochastic-gradient-langevin-dynamics-fsgld
type: topic
---

# Flatness-Aware Stochastic Gradient Langevin Dynamics

Searching arXiv for the main paper and related context papers.
Flatness-Aware Stochastic Gradient Langevin Dynamics (fSGLD) is a stochastic optimization method for high-dimensional nonconvex learning problems that combines Langevin dynamics with random weight perturbation (RWP) in order to bias training toward flat minima. In the formulation introduced in "Flatness-Aware Stochastic Gradient Langevin Dynamics" [2510.02174], the method evaluates stochastic gradients at Gaussian-perturbed parameters and thereby optimizes a randomized-smoothing surrogate whose leading second-order term is the Hessian trace. Under a principled coupling between the inverse temperature $\beta$ and the perturbation scale $\sigma$, the invariant measure of the resulting dynamics stays close to a stationary measure concentrated on the global minimizers of a Hessian-trace-regularized loss, while retaining the computational cost of SGD and about half the cost of SAM [2510.02174].

## 1. Definition and motivation

The motivating premise of fSGLD is that generalization in deep learning is closely tied to the pursuit of flat minima in the loss landscape. Classical SGLD adds isotropic Gaussian noise to SGD and induces a Gibbs measure $\pi_\beta(\theta)\propto \exp(-\beta u(\theta))$, but it has no intrinsic mechanism to distinguish flat from sharp minimizers when several solutions attain comparable objective values [2510.02174]. fSGLD addresses that deficiency by replacing gradients at $\theta$ with gradients at $\theta+\epsilon$, where $\epsilon\sim \mathcal{N}(0,\sigma^2 I)$, and by preserving the Langevin noise that underlies SGLD.

The role of RWP is central. Evaluating the loss at perturbed weights corresponds to optimizing a randomized-smoothing surrogate, and Gaussian smoothing exposes curvature explicitly because its leading correction to the original loss is proportional to the Hessian trace. This makes fSGLD an optimizer with an implicit curvature penalty, but without explicit Hessian computation [2510.02174].

A common misconception is that random perturbation alone automatically yields a clean flatness regularizer. The paper argues that this is not generally true in high dimensions, because randomized smoothing also contains higher-order remainder terms that can be non-negligible and can corrupt the intended flatness bias. The distinctive theoretical contribution of fSGLD is therefore not only the use of RWP, but the coupling of $\beta$ and $\sigma$ that isolates the Hessian-trace regularization effect in the invariant measure [2510.02174].

This perspective is consistent with adjacent lines of work. One related direction shows that anisotropic SGD noise induces an effective landscape-dependent regularization favoring flat solutions [2206.01246]. Another, Natural Langevin Dynamics, uses inverse-Fisher preconditioning so that sensitive directions receive smaller steps and less noise, yielding a curvature-aware SGLD variant based on Fisher geometry rather than randomized smoothing [1712.01076]. fSGLD differs from both by tying its flatness bias directly to Gaussian smoothing and to a Hessian-trace-regularized objective.

## 2. Mathematical formulation

The starting point is the nonconvex objective
$$
u(\theta):=\mathbb{E}[U(\theta,X)],
$$
with gradient $h(\theta):=\nabla u(\theta)$. The flatness-aware regularized objective is
$$
v(\theta):=u(\theta)+\frac{\sigma^2}{2}\operatorname{tr}(H(\theta)),
$$
where $H(\theta)$ is the Hessian of $u$ [2510.02174].

Introducing $\epsilon\sim \mathcal{N}(0,\sigma^2 I_d)$, independent of $X$, yields the randomized-smoothing surrogate
$$
g_\epsilon(\theta):=\mathbb{E}[u(\theta+\epsilon)].
$$
A Taylor expansion gives
$$
u(\theta+\epsilon)
=
u(\theta)+\nabla u(\theta)^\top \epsilon+\frac12 \epsilon^\top H(\theta)\epsilon+\mathcal{R}(\theta,\epsilon),
$$
and averaging over $\epsilon$ produces
$$
g_\epsilon(\theta)
=
u(\theta)+\frac{\sigma^2}{2}\operatorname{tr}(H(\theta))+\mathbb{E}[\mathcal{R}(\theta,\epsilon)]
=
v(\theta)+\mathbb{E}[\mathcal{R}(\theta,\epsilon)].
$$
The remainder has an explicit fourth-order form, and $\mathbb{E}[\mathcal{R}(\theta,\epsilon)]=O(\sigma^4)$ under mild smoothness [2510.02174]. This is the precise sense in which randomized smoothing approximates Hessian-trace regularization while also introducing higher-order corrections.

In continuous time, the smoothed Langevin dynamics is
$$
dZ_t^{\mathrm{fSGLD}}=-\nabla g_\epsilon(Z_t^{\mathrm{fSGLD}})\,dt+\sqrt{2\beta^{-1}}\,dB_t,
$$
with invariant measure
$$
\pi_\beta^{\mathrm{fSGLD}}(\theta)\propto \exp(-\beta g_\epsilon(\theta)).
$$
The corresponding ideal Gibbs measure for the Hessian-trace-regularized objective is
$$
\pi^\star_{\beta,\sigma}(d\theta)\propto \exp(-\beta v(\theta))\,d\theta
$$
[2510.02174].

The discrete algorithm used in practice combines RWP and Langevin noise:
$$
\theta_{k+1}^{\mathrm{fSGLD}}
=
\theta_k^{\mathrm{fSGLD}}
-
\lambda \nabla_\theta U(\theta_k^{\mathrm{fSGLD}}+\epsilon_{k+1},X_{k+1})
+
\sqrt{2\lambda\beta^{-1}}\,\xi_{k+1},
$$
where $(X_k)$ are i.i.d. data, $(\epsilon_k)$ are i.i.d. $\mathcal{N}(0,\sigma^2 I_d)$, and $(\xi_k)$ are i.i.d. $\mathcal{N}(0,I_d)$, all mutually independent [2510.02174]. The perturbed gradient is unbiased for $\nabla g_\epsilon(\theta)$:
$$
\nabla g_\epsilon(\theta)=\mathbb{E}\big[\mathbb{E}_X[\nabla_\theta U(\theta+\epsilon,X)]\big].
$$

A practical variant averages multiple perturbations per step, but empirically $n=1$ already works well [2510.02174].

## 3. Invariant measures and convergence theory

The core theoretical result is an invariant-measure alignment theorem under the coupling
$$
\sigma=\beta^{-\frac{1+\eta}{4}},\qquad \eta>0.
$$
Under this choice, $\mathbb{E}[\mathcal{R}]=O(\sigma^4)=O(\beta^{-(1+\eta)})$, and the invariant measure of fSGLD stays close to the Gibbs measure associated with the Hessian-trace-regularized objective. Specifically,
$$
\lim_{\beta\to\infty}\mathrm{KL}(\pi_\beta^{\mathrm{fSGLD}}\|\pi^\star_{\beta,\sigma})=0,
\qquad
\lim_{\beta\to\infty}W_2(\pi_\beta^{\mathrm{fSGLD}},\pi^\star_{\beta,\sigma})=0
$$
[2510.02174].

The relation between the two measures is made explicit by
$$
\pi_\beta^{\mathrm{fSGLD}}(d\theta)
=
Z_\beta^{-1}Z_{\beta,\sigma}\,
\exp\!\big(-\beta\,\mathbb{E}[\mathcal{R}(\theta,\epsilon)]\big)\,
\pi^\star_{\beta,\sigma}(d\theta),
$$
where
$$
Z_\beta:=\int_{\mathbb{R}^d}e^{-\beta g_\epsilon(\theta)}\,d\theta,
\qquad
Z_{\beta,\sigma}:=\int_{\mathbb{R}^d}e^{-\beta v(\theta)}\,d\theta.
$$
This identity clarifies why the $\beta$-$\sigma$ coupling matters: it suppresses the effect of the smoothing remainder in the Gibbs comparison and makes the flatness regularization dominant [2510.02174].

The paper also establishes non-asymptotic convergence of the $k$-th iterate to $\pi^\star_{\beta,\sigma}$ in Wasserstein distance under unbiased gradients, independence, Lipschitzness, and dissipativity. The $W_1$ bound is
$$
W_1(\mathcal{L}(\theta_k^{\mathrm{fSGLD}}),\pi^\star_{\beta,\sigma})
\le
D_1 e^{-\dot c\lambda k/2}(1+\mathbb{E}[|\theta_0|^4])
+
(D_2+D_3)\sqrt{\lambda}
+
\underline D,
$$
with $\underline D=\mathcal{O}(\beta^{-\eta})$, and the $W_2$ bound is
$$
W_2(\mathcal{L}(\theta_k^{\mathrm{fSGLD}}),\pi^\star_{\beta,\sigma})
\le
D_4 e^{-\dot c\lambda k/4}(\mathbb{E}[|\theta_0|^4]+1)
+
(D_5+D_6)\lambda^{1/4}
+
\underline D
$$
[2510.02174]. The $\sqrt{\lambda}$ and $\lambda^{1/4}$ discretization terms match the best-known non-asymptotic rates for SGLD under comparable assumptions, but the target measure is now associated with global minimizers of $v$, not of the original $u$.

An excess-risk bound is also derived:
$$
\mathbb{E}[g_\epsilon(\theta_k^{\mathrm{fSGLD}})]-\inf_{\theta\in\mathbb{R}^d} v(\theta)
\le
D_1^\# e^{-\dot c\lambda k/4}
+
D_2^\#\lambda^{1/4}
+
D_3^\#,
$$
where
$$
D_3^\#=\mathcal{O}\!\left(\frac{d}{\beta}\log\!\bigl(\tfrac{\beta}{d}+1\bigr)+\beta^{-(1+\eta)}\right).
$$
This gives an optimizer-level guarantee that the method reaches near-global flat minima of the Hessian-trace-regularized objective [2510.02174].

The proof strategy combines a Taylor expansion of the smoothing objective, KL-to-$W_2$ control through a weighted transportation inequality of Bolley–Villani type, and contractivity plus moment bounds under dissipativity. The use of Eberle-like contractivity for $W_1$ and a lift to $W_2$ explains the exponential decay terms, while stability of Euler–Maruyama discretization with stochastic gradients yields the $\sqrt{\lambda}$ and $\lambda^{1/4}$ bias terms [2510.02174].

## 4. Relation to SGD, SGLD, SAM, SWA, and other Langevin variants

Relative to SGD, fSGLD changes the objective implicitly rather than merely altering the trajectory noise. SGD minimizes $u$ using noisy gradients and has no explicit bias toward flatness. Classical SGLD augments SGD with isotropic Gaussian noise and converges toward $\pi_\beta(\theta)\propto \exp(-\beta u(\theta))$, but that measure is indifferent to curvature when several minima have similar loss values [2510.02174].

fSGLD instead optimizes the smoothed objective
$$
g_\epsilon(\theta)=\mathbb{E}[u(\theta+\epsilon)],
$$
which approximates
$$
v(\theta)=u(\theta)+\frac{\sigma^2}{2}\operatorname{tr}(H(\theta)).
$$
Its update rule captures curvature through Gaussian smoothing while avoiding direct Hessian computation. This distinguishes the method from heuristic sharpness-averse procedures by making the implicit regularizer mathematically explicit [2510.02174].

The comparison with SAM and ASAM is practical as well as conceptual. SAM replaces the gradient with one computed from a worst-case local perturbation in a min–max problem and is effective at penalizing sharpness, but it requires approximately two gradient evaluations per step and nearly doubles the cost of SGD. fSGLD uses symmetric Gaussian perturbations and a single gradient evaluation at $\theta+\epsilon$ per step, so its cost remains comparable to SGD. Empirically it matches or surpasses SAM with roughly half the cost [2510.02174].

SWA occupies a different position. Averaging weights along the optimization trajectory can yield flatter solutions, but SWA lacks both an explicit curvature surrogate and an invariant-measure analysis. fSGLD, by contrast, supplies a randomized-smoothing surrogate tied to the Hessian trace and a convergence theory that links its invariant distribution to a Hessian-trace-regularized Gibbs measure [2510.02174].

A separate curvature-aware Langevin family is represented by Natural Langevin Dynamics, which uses inverse-Fisher preconditioning for both drift and noise so that sensitive directions receive smaller steps and less noise [1712.01076]. This suggests a different operational notion of flatness awareness: fSGLD derives it from randomized smoothing and Hessian-trace regularization, whereas Natural Langevin Dynamics derives it from Fisher geometry and posterior covariance.

## 5. Empirical behavior

The empirical study in [2510.02174] covers noisy-label and large-scale vision tasks in both training-from-scratch and fine-tuning settings. The noisy-label benchmarks are CIFAR-10N and CIFAR-100N, both with real human annotation noise, and WebVision using the first 50 classes from the Google subset. The reported metrics are CIFAR accuracy, WV-1/Top-1, and WV-5/Top-5. The models are ResNet-34 and ResNet-50 for training from scratch, and a pretrained ViT-B/16 for fine-tuning.

Training from scratch uses 150 epochs, batch size 128, and learning-rate decays by a factor of $0.1$ at epochs 50 and 100. Fine-tuning uses 75 epochs, batch size 128, and a learning-rate decay by a factor of $0.1$ at epoch 50. Hyperparameters are tuned with Optuna using 20 trials per optimizer. For fSGLD, $\sigma$ is tuned and $\beta$ is set from the theory as $\beta=\sigma^{-4/(1+\eta)}$ with $\eta=0.01$; a comparison variant fixes $\beta=10^{14}$ [2510.02174].

| Setting | fSGLD | Comparator highlights |
|---|---:|---:|
| ResNet-34, CIFAR-100N | 62.02±0.29 | SAM 59.18±0.33; SGLD 57.33±0.36; SGD 58.47±0.20 |
| ResNet-50, CIFAR-100N | 62.08±0.45 | SAM 59.01±0.60; SGLD 56.90±0.65 |
| CIFAR-10N | 91.72±0.20 | SAM 91.53±0.22; SGLD 88.77±0.51 |
| WebVision WV-1 | 73.55±0.27 | SAM 73.49±0.36 |
| WebVision WV-5 | 89.86±0.12 | SAM 90.32±0.31 |
| ViT-B/16, CIFAR-100N | 75.18 | SAM 74.66 |
| ViT-B/16, CIFAR-10N | 96.72 | SAM 96.75 |

These results support two claims made in the paper: first, that the $\beta$-$\sigma$ coupled version is generally the strongest variant; and second, that the method is competitive with or better than SAM while preserving a much lower cost profile [2510.02174].

The runtime measurements quantify that profile. On ResNet-34 with CIFAR-10N, wall-clock per iteration is 23.7 s for fSGLD, 22.0 s for SGD, and 41.3 s for SAM; on ResNet-50, the corresponding numbers are 34.1 s, 31.9 s, and 60.7 s. For ViT-B/16, per-epoch runtime is 345.8 s for fSGLD and 656.7 s for SAM [2510.02174].

Ablation results further refine the picture. Increasing the number of perturbations from $n=1$ to $n=2$-$5$ yields only minor gains while increasing runtime nearly linearly, so the single-perturbation version is already effective. Performance is stable for small to moderate $\sigma$, whereas overly large $\sigma$ degrades training stability. The theoretically prescribed $\beta$-$\sigma$ coupling consistently outperforms the heuristic fixed-$\beta$ variant in most settings [2510.02174].

Hessian-spectrum analysis provides direct evidence for the claimed mechanism. Using Lanczos for the top-50 eigenvalues and Hutchinson’s estimator for the trace, the paper reports that fSGLD converges to significantly flatter minima than SGD, with lower $\lambda_{\max}$ and lower trace, and to minima comparable to or flatter than those found by SAM, with similar $\lambda_{\max}$ and lower trace [2510.02174].

## 6. Practical use, limitations, and open directions

The practical prescription given in [2510.02174] is deliberately simple. The perturbation scale $\sigma$ should be chosen in a small to moderate regime, and tuning over $\sigma\in 10^{[-3,-2]}$ is reported as effective in the experiments. The inverse temperature should be coupled to $\sigma$ by
$$
\beta=\sigma^{-4/(1+\eta)},
$$
with $\eta\approx 0.01$ so that Langevin noise does not overwhelm the gradient and the higher-order smoothing remainder is suppressed. The step size $\lambda$ can follow standard SGD schedules; smaller $\lambda$ reduces the discretization bias terms but slows optimization. Batch size and weight decay can follow standard training settings, and the method is compatible with typical training pipelines [2510.02174].

Operationally, each iteration samples $\epsilon_k\sim\mathcal{N}(0,\sigma^2 I)$ and $\xi_k\sim\mathcal{N}(0,I)$ independently, forms the perturbed parameter $\tilde\theta_k=\theta_k+\epsilon_k$, computes the stochastic gradient on a minibatch at $\tilde\theta_k$, and applies the Langevin update. Averaging multiple perturbed gradients is optional and primarily serves variance reduction at higher computational cost [2510.02174].

Several failure modes are explicitly identified. If $\sigma$ is too large, perturbations destabilize training; if it is too small, the flatness bias weakens. If $\beta$ is fixed independently of $\sigma$ and made very large, the method approaches near-SGD on the perturbed loss and loses the invariant-measure alignment to $v$. Excessive Langevin noise, corresponding to large $\beta^{-1}$, increases exploration but can harm optimization [2510.02174].

The method is presented as especially useful in noisy-label regimes and in overparameterized deep networks where flat minima correlate strongly with generalization and robustness. Fine-tuning large-scale vision models is another setting in which it often matches or exceeds SAM at half the cost [2510.02174].

The theoretical scope is not universal. The analysis assumes unbiased stochastic gradients, Lipschitz continuity, and dissipativity, and extending it to semiconvex, one-sided Lipschitz gradients is identified as an open direction. Applying the method to diffusion-based generative models is described as promising, but whether a bias toward flatter regions improves sample diversity or quality remains unresolved [2510.02174]. A plausible implication is that future work will focus less on the existence of a flatness bias, which is already formalized, and more on how different curvature surrogates and geometry-aware noise models interact with large-model training dynamics.

Source: https://www.emergentmind.com/topics/flatness-aware-stochastic-gradient-langevin-dynamics-fsgld