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.
翻译:贝叶斯神经网络(BNNs)需要可扩展的采样算法来近似参数的后验分布。现有随机梯度马尔可夫链蒙特卡洛(SGMCMC)方法对步长选择高度敏感,且诸如pSGLD等自适应变体通常无法在未添加昂贵散度校正项的情况下正确采样不变测度。本研究基于近期提出的时间步长自适应"SamAdams"框架(Leimkuhler, Lohmann, and Whalley 2025),引入一种自适应方案:SA-SGLD,该方法通过时间重缩放根据监测量(通常为局部梯度范数)调节步长。SA-SGLD能自动在高曲率区域收缩步长,在平坦区域扩展步长,从而在不引入偏差的情况下提升稳定性与混合效率。我们证明,该方法在高曲率二维玩具示例及采用尖锐先验的BNN图像分类任务中,可实现比SGLD更精确的后验采样。