Flatness-Aware Stochastic Gradient Langevin Dynamics

Flatness of the loss landscape has been widely studied as an important perspective for understanding the behavior and generalization of deep learning algorithms. Motivated by this view, we propose Flatness-Aware Stochastic Gradient Langevin Dynamics (fSGLD), a first-order optimization method that biases learning its dynamics toward flat basins while retaining the computational and memory efficiency of SGD and SGLD. We provide a non-asymptotic theoretical analysis showing that fSGLD converges to a flatness-biased Gibbs distribution under a theoretically prescribed coupling between the noise scale $σ$ and the inverse temperature $β$, together with explicit excess risk guarantees. We empirically evaluate fSGLD across standard optimizer benchmarks, Bayesian image classification, uncertainty quantification, and out-of-distribution detection, demonstrating consistently strong performance and reliable uncertainty estimates. Additional experiments confirm the effectiveness of the theoretically prescribed $β$-$σ$ coupling compared to decoupled choices.

💡 Research Summary

This paper introduces Flatness-Aware Stochastic Gradient Langevin Dynamics (fSGLD), a novel first-order optimization algorithm designed to bridge the gap between two important lines of research: methods that bias learning towards flat minima for better generalization, and sampling-based methods like Stochastic Gradient Langevin Dynamics (SGLD) that offer global exploration capabilities.

The core motivation stems from the well-known “flat minima hypothesis,” which suggests that parameters located in flat, low-curvature regions of the loss landscape lead to models that generalize better. While methods like Sharpness-Aware Minimization (SAM) explicitly seek flat regions, they are inherently local and computationally expensive, requiring double gradient computations. On the other hand, SGLD provides a principled global exploration mechanism but is agnostic to the geometry (flatness) of the loss landscape, as it solely targets the Gibbs distribution defined by the original loss function.

fSGLD elegantly combines these advantages through a simple yet powerful modification to the standard SGLD update rule. At each iteration, instead of using the stochastic gradient at the current parameters θ, fSGLD perturbs the parameters with Gaussian noise ε and computes the gradient at this perturbed point: θ_{k+1} = θ_k - λ∇U(θ_k + ε_{k+1}, X_{k+1}) + √(2λβ^{-1}) ξ_{k+1}. The perturbation scale σ is not a free parameter but is theoretically coupled with the inverse temperature β of the Langevin dynamics as σ := β^{-(1+η)/4} for a small fixed η.

The authors provide a profound insight: this perturbed gradient is an unbiased estimator of the gradient of a “randomized-smoothing surrogate” objective, g_ε(θ) = E