Adaptive Stepsizing for Stochastic Gradient Langevin Dynamics in Bayesian Neural Networks

TL;DR

This paper introduces SA-SGLD, an adaptive stepsize scheme for stochastic gradient Langevin dynamics in Bayesian neural networks, improving sampling accuracy and stability without bias.

Contribution

It develops an adaptive stepsize method based on time rescaling, enhancing sampling efficiency in Bayesian neural networks without divergence correction.

Findings

SA-SGLD outperforms SGLD in high-curvature 2D examples

It achieves more accurate posterior sampling in BNNs with sharp priors

The method improves stability and mixing without bias

Abstract

Bayesian neural networks (BNNs) require scalable sampling algorithms to approximate posterior distributions over parameters. Existing stochastic gradient Markov Chain Monte Carlo (SGMCMC) methods are highly sensitive to the choice of stepsize and adaptive variants such as pSGLD typically fail to sample the correct invariant measure without addition of a costly divergence correction term. In this work, we build on the recently proposed `SamAdams' framework for timestep adaptation (Leimkuhler, Lohmann, and Whalley 2025), introducing an adaptive scheme: SA-SGLD, which employs time rescaling to modulate the stepsize according to a monitored quantity (typically the local gradient norm). SA-SGLD can automatically shrink stepsizes in regions of high curvature and expand them in flatter regions, improving both stability and mixing without introducing bias. We show that our method can achieve more accurate posterior sampling than SGLD on high-curvature 2D toy examples and in image classification with BNNs using sharp priors.