Flatness-Aware Stochastic Gradient Langevin Dynamics

• The paper introduces fSGLD, which leverages Gaussian perturbations to evaluate gradients and implicitly regularize via the Hessian trace.
• fSGLD couples the inverse temperature and perturbation scale, ensuring the invariant measure aligns with a Hessian-trace-regularized objective for flat minima.
• Empirical results on noisy-label and vision tasks demonstrate that fSGLD achieves competitive accuracy with lower computational cost compared to methods like SAM.

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" (Bruno et al., 2 Oct 2025), 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 (Bruno et al., 2 Oct 2025).

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 (Bruno et al., 2 Oct 2025). 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 (Bruno et al., 2 Oct 2025).

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 (Bruno et al., 2 Oct 2025).

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 (Yang et al., 2022). 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 (Marceau-Caron et al., 2017). 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

$\sigma$0

where $\sigma$1 is the Hessian of $\sigma$2 (Bruno et al., 2 Oct 2025).

Introducing $\sigma$3, independent of $\sigma$4, yields the randomized-smoothing surrogate

$\sigma$5

A Taylor expansion gives

$\sigma$6

and averaging over $\sigma$7 produces

$\sigma$8

The remainder has an explicit fourth-order form, and $\sigma$9 under mild smoothness (Bruno et al., 2 Oct 2025). 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

$\pi_\beta(\theta)\propto \exp(-\beta u(\theta))$0

with invariant measure

$\pi_\beta(\theta)\propto \exp(-\beta u(\theta))$1

The corresponding ideal Gibbs measure for the Hessian-trace-regularized objective is

$\pi_\beta(\theta)\propto \exp(-\beta u(\theta))$2

(Bruno et al., 2 Oct 2025).

The discrete algorithm used in practice combines RWP and Langevin noise:

$\pi_\beta(\theta)\propto \exp(-\beta u(\theta))$3

where $\pi_\beta(\theta)\propto \exp(-\beta u(\theta))$4 are i.i.d. data, $\pi_\beta(\theta)\propto \exp(-\beta u(\theta))$5 are i.i.d. $\pi_\beta(\theta)\propto \exp(-\beta u(\theta))$6, and $\pi_\beta(\theta)\propto \exp(-\beta u(\theta))$7 are i.i.d. $\pi_\beta(\theta)\propto \exp(-\beta u(\theta))$8, all mutually independent (Bruno et al., 2 Oct 2025). The perturbed gradient is unbiased for $\pi_\beta(\theta)\propto \exp(-\beta u(\theta))$9:

$\theta$0

A practical variant averages multiple perturbations per step, but empirically $\theta$1 already works well (Bruno et al., 2 Oct 2025).

3. Invariant measures and convergence theory

The core theoretical result is an invariant-measure alignment theorem under the coupling

$\theta$2

Under this choice, $\theta$3, and the invariant measure of fSGLD stays close to the Gibbs measure associated with the Hessian-trace-regularized objective. Specifically,

$\theta$4

(Bruno et al., 2 Oct 2025).

The relation between the two measures is made explicit by

$\theta$5

where

$\theta$6

This identity clarifies why the $\theta$7-$\theta$8 coupling matters: it suppresses the effect of the smoothing remainder in the Gibbs comparison and makes the flatness regularization dominant (Bruno et al., 2 Oct 2025).

The paper also establishes non-asymptotic convergence of the $\theta$9-th iterate to $\theta+\epsilon$0 in Wasserstein distance under unbiased gradients, independence, Lipschitzness, and dissipativity. The $\theta+\epsilon$1 bound is

$\theta+\epsilon$2

with $\theta+\epsilon$3, and the $\theta+\epsilon$4 bound is

$\theta+\epsilon$5

(Bruno et al., 2 Oct 2025). The $\theta+\epsilon$6 and $\theta+\epsilon$7 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 $\theta+\epsilon$8, not of the original $\theta+\epsilon$9.

An excess-risk bound is also derived:

$\epsilon\sim \mathcal{N}(0,\sigma^2 I)$0

where

$\epsilon\sim \mathcal{N}(0,\sigma^2 I)$1

This gives an optimizer-level guarantee that the method reaches near-global flat minima of the Hessian-trace-regularized objective (Bruno et al., 2 Oct 2025).

The proof strategy combines a Taylor expansion of the smoothing objective, KL-to-$\epsilon\sim \mathcal{N}(0,\sigma^2 I)$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 $\epsilon\sim \mathcal{N}(0,\sigma^2 I)$3 and a lift to $\epsilon\sim \mathcal{N}(0,\sigma^2 I)$4 explains the exponential decay terms, while stability of Euler–Maruyama discretization with stochastic gradients yields the $\epsilon\sim \mathcal{N}(0,\sigma^2 I)$5 and $\epsilon\sim \mathcal{N}(0,\sigma^2 I)$6 bias terms (Bruno et al., 2 Oct 2025).

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 $\epsilon\sim \mathcal{N}(0,\sigma^2 I)$7 using noisy gradients and has no explicit bias toward flatness. Classical SGLD augments SGD with isotropic Gaussian noise and converges toward $\epsilon\sim \mathcal{N}(0,\sigma^2 I)$8, but that measure is indifferent to curvature when several minima have similar loss values (Bruno et al., 2 Oct 2025).

fSGLD instead optimizes the smoothed objective

$\epsilon\sim \mathcal{N}(0,\sigma^2 I)$9

which approximates

$\beta$0

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 (Bruno et al., 2 Oct 2025).

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 $\beta$1 per step, so its cost remains comparable to SGD. Empirically it matches or surpasses SAM with roughly half the cost (Bruno et al., 2 Oct 2025).

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 (Bruno et al., 2 Oct 2025).

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 (Marceau-Caron et al., 2017). 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 (Bruno et al., 2 Oct 2025) 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 $\beta$2 at epochs 50 and 100. Fine-tuning uses 75 epochs, batch size 128, and a learning-rate decay by a factor of $\beta$3 at epoch 50. Hyperparameters are tuned with Optuna using 20 trials per optimizer. For fSGLD, $\beta$4 is tuned and $\beta$5 is set from the theory as $\beta$6 with $\beta$7; a comparison variant fixes $\beta$8 (Bruno et al., 2 Oct 2025).

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$9-$\sigma$0 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 (Bruno et al., 2 Oct 2025).

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 (Bruno et al., 2 Oct 2025).

Ablation results further refine the picture. Increasing the number of perturbations from $\sigma$1 to $\sigma$2-$\sigma$3 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$4, whereas overly large $\sigma$5 degrades training stability. The theoretically prescribed $\sigma$6-$\sigma$7 coupling consistently outperforms the heuristic fixed-$\sigma$8 variant in most settings (Bruno et al., 2 Oct 2025).

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 $\sigma$9 and lower trace, and to minima comparable to or flatter than those found by SAM, with similar $u(\theta):=\mathbb{E}[U(\theta,X)],$0 and lower trace (Bruno et al., 2 Oct 2025).

6. Practical use, limitations, and open directions

The practical prescription given in (Bruno et al., 2 Oct 2025) is deliberately simple. The perturbation scale $u(\theta):=\mathbb{E}[U(\theta,X)],$1 should be chosen in a small to moderate regime, and tuning over $u(\theta):=\mathbb{E}[U(\theta,X)],$2 is reported as effective in the experiments. The inverse temperature should be coupled to $u(\theta):=\mathbb{E}[U(\theta,X)],$3 by

$u(\theta):=\mathbb{E}[U(\theta,X)],$4

with $u(\theta):=\mathbb{E}[U(\theta,X)],$5 so that Langevin noise does not overwhelm the gradient and the higher-order smoothing remainder is suppressed. The step size $u(\theta):=\mathbb{E}[U(\theta,X)],$6 can follow standard SGD schedules; smaller $u(\theta):=\mathbb{E}[U(\theta,X)],$7 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 (Bruno et al., 2 Oct 2025).

Operationally, each iteration samples $u(\theta):=\mathbb{E}[U(\theta,X)],$8 and $u(\theta):=\mathbb{E}[U(\theta,X)],$9 independently, forms the perturbed parameter $h(\theta):=\nabla u(\theta)$0, computes the stochastic gradient on a minibatch at $h(\theta):=\nabla u(\theta)$1, and applies the Langevin update. Averaging multiple perturbed gradients is optional and primarily serves variance reduction at higher computational cost (Bruno et al., 2 Oct 2025).

Several failure modes are explicitly identified. If $h(\theta):=\nabla u(\theta)$2 is too large, perturbations destabilize training; if it is too small, the flatness bias weakens. If $h(\theta):=\nabla u(\theta)$3 is fixed independently of $h(\theta):=\nabla u(\theta)$4 and made very large, the method approaches near-SGD on the perturbed loss and loses the invariant-measure alignment to $h(\theta):=\nabla u(\theta)$5. Excessive Langevin noise, corresponding to large $h(\theta):=\nabla u(\theta)$6, increases exploration but can harm optimization (Bruno et al., 2 Oct 2025).

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 (Bruno et al., 2 Oct 2025).

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 (Bruno et al., 2 Oct 2025). 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.