Flatness-Aware Stochastic Gradient Langevin Dynamics

Stefano Bruno, Youngsik Hwang, Jaehyeon An, Sotirios Sabanis, Dongyoung Lim
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:10017-10050, 2026.

Abstract

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-bruno26a, title = {Flatness-Aware Stochastic Gradient {L}angevin Dynamics}, author = {Bruno, Stefano and Hwang, Youngsik and An, Jaehyeon and Sabanis, Sotirios and Lim, Dongyoung}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {10017--10050}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/bruno26a/bruno26a.pdf}, url = {https://proceedings.mlr.press/v306/bruno26a.html}, abstract = {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.} }
Endnote
%0 Conference Paper %T Flatness-Aware Stochastic Gradient Langevin Dynamics %A Stefano Bruno %A Youngsik Hwang %A Jaehyeon An %A Sotirios Sabanis %A Dongyoung Lim %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-bruno26a %I PMLR %P 10017--10050 %U https://proceedings.mlr.press/v306/bruno26a.html %V 306 %X 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.
APA
Bruno, S., Hwang, Y., An, J., Sabanis, S. & Lim, D.. (2026). Flatness-Aware Stochastic Gradient Langevin Dynamics. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:10017-10050 Available from https://proceedings.mlr.press/v306/bruno26a.html.

Related Material