Flatness-Aware Stochastic Gradient Langevin Dynamics

TL;DR

Flatness-Aware Stochastic Gradient Langevin Dynamics is introduced, designed to efficiently and provably seek flat minima in high-dimensional nonconvex optimization problems and achieves superior or comparable generalization and robustness to baseline algorithms while maintaining the computational cost of SGD, about half that of SAM.

Abstract

Generalization in deep learning is closely tied to the pursuit of flat minima in the loss landscape, yet classical Stochastic Gradient Langevin Dynamics (SGLD) offers no mechanism to bias its dynamics toward such low-curvature solutions. This work introduces Flatness-Aware Stochastic Gradient Langevin Dynamics (fSGLD), designed to efficiently and provably seek flat minima in high-dimensional nonconvex optimization problems. At each iteration, fSGLD uses the stochastic gradient evaluated at parameters perturbed by isotropic Gaussian noise, commonly referred to as Random Weight Perturbation (RWP), thereby optimizing a randomized-smoothing objective that implicitly captures curvature information. Leveraging these properties, we prove that the invariant measure of fSGLD stays close to a stationary measure concentrated on the global minimizers of a loss function regularized by the Hessian trace whenever the inverse temperature and the scale of random weight perturbation are properly coupled. This result provides a rigorous theoretical explanation for the benefits of random weight perturbation. In particular, we establish non-asymptotic convergence guarantees in Wasserstein distance with the best known rate and derive an excess-risk bound for the Hessian-trace regularized objective. Extensive experiments on noisy-label and large-scale vision tasks, in both training-from-scratch and fine-tuning settings, demonstrate that fSGLD achieves superior or comparable generalization and robustness to baseline algorithms while maintaining the computational cost of SGD, about half that of SAM. Hessian-spectrum analysis further confirms that fSGLD converges to significantly flatter minima.

Figures (4)

• Figure 1: A schematic overview of the theoretical framework of fSGLD. The process begins with the original objective$u(\theta)$ and its associated Gibbs measure $\pi_\beta^{\text{SGLD}}$ (left). Randomized smoothing transforms this into a tractable surrogate objective, $g_\epsilon(\theta)$, which is the basis for the fSGLD algorithm and its invariant measure, $\pi^{\text{fSGLD}}_\beta$ (center). This highlights a key distinction: while the Gibbs measure of standard SGLD, $\pi_\beta^{\text{SGLD}}$, is indifferent to the flatness of the minima, the fSGLD framework is designed such that its invariant measure, $\pi^{\text{fSGLD}}_\beta$, targets the distribution over the flattest minima. Our ultimate goal is to target the Hessian-trace regularized objective$v(\theta)$ and its corresponding measure $\pi^{\star}_{\beta, \sigma}$, which concentrates on the desired global flat minima (right).
• Figure 2: Sensitive analysis of noise standard deviation $\sigma$ on CIFAR-10N with ResNet-34.
• Figure 3: Performance with respect to the number of random perturbations $n$ used in fSGLD.
• Figure 3: The distribution of the leading eigenvalues and Hessian trace of ResNet-34 trained on CIFAR-10N with SGD, SAM, and fSGLD.

Theorems & Definitions (48)

• Proposition 3.1
• Theorem 3.2
• Corollary 3.3
• Remark 3.4
• Theorem 3.5
• Remark 3.6
• Remark B.1
• Remark B.2
• proof : Proof of Remark \ref{['dissipativity_fSGLD_remark']}
• Remark B.3