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 $\sigma$ and the inverse temperature $\beta$, 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 $\beta$–$\sigma$ coupling compared to decoupled choices.

Training a neural network is like searching a vast landscape for a low point, but not all low points are equally useful. Some solutions sit in narrow valleys, where a small change can make the model worse, while others sit in wide, flatter valleys that often lead to better performance on new data. Existing methods can encourage models to find these flatter solutions, but they often require extra computations, extra memory, or only look in a small neighborhood of the current model. We introduce Flatness-Aware Stochastic Gradient Langevin Dynamics, a training method that uses carefully calibrated random noise to help the model explore the landscape and favor flatter, more stable solutions. The key idea is simple: instead of measuring flatness directly, which is expensive, we add small random perturbations during training so that the algorithm can sense whether a solution remains good nearby. Our theory shows how to choose the amount of noise so that this effect reliably targets flat solutions. In experiments on image classification, uncertainty estimation, and detecting unfamiliar data, our method performs competitively or better than strong baselines while using about the same memory and computation as standard training methods. This makes flatness-aware training more practical for building neural networks that generalize well and provide more reliable uncertainty estimates.