Adaptive Stepsizing for Stochastic Gradient Langevin Dynamics in Bayesian Neural Networks
This work addresses the challenge of scalable and stable posterior sampling in Bayesian neural networks, though it is incremental as it builds on an existing framework.
The paper tackled the problem of stepsize sensitivity in stochastic gradient Langevin dynamics for Bayesian neural networks by introducing an adaptive scheme that automatically adjusts stepsizes based on local gradient norms, achieving more accurate posterior sampling in high-curvature toy examples and image classification tasks.
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.