[论文解读] AMAGOLD: Amortized Metropolis Adjustment for Efficient Stochastic Gradient MCMC

Ruqi Zhang, A. Feder Cooper+1 more|arXiv (Cornell University)|Feb 29, 2020

Markov Chains and Monte Carlo Methods被引用 8

一句话总结

AMAGOLD 是一种新颖的二阶随机梯度 MCMC 算法，通过不频繁地应用 Metropolis-Hastings 修正来消除 SGHMC 中的偏差，同时分摊其计算成本。它在固定步长下实现渐近精确采样，且收敛速率在常数因子范围内与全批量基线方法相当，在合成数据和真实世界贝叶斯推断任务中，其鲁棒性和准确性均优于 SGHMC。

ABSTRACT

Stochastic gradient Hamiltonian Monte Carlo (SGHMC) is an efficient method for sampling from continuous distributions. It is a faster alternative to HMC: instead of using the whole dataset at each iteration, SGHMC uses only a subsample. This improves performance, but introduces bias that can cause SGHMC to converge to the wrong distribution. One can prevent this using a step size that decays to zero, but such a step size schedule can drastically slow down convergence. To address this tension, we propose a novel second-order SG-MCMC algorithm---AMAGOLD---that infrequently uses Metropolis-Hastings (M-H) corrections to remove bias. The infrequency of corrections amortizes their cost. We prove AMAGOLD converges to the target distribution with a fixed, rather than a diminishing, step size, and that its convergence rate is at most a constant factor slower than a full-batch baseline. We empirically demonstrate AMAGOLD's effectiveness on synthetic distributions, Bayesian logistic regression, and Bayesian neural networks.

研究动机与目标

• 为解决 SGHMC 中由随机梯度和固定步长引入的偏差问题，这些偏差会阻碍收敛到真实后验分布。
• 在保持随机梯度 MCMC 计算效率的同时，确保渐近精确性。
• 通过仅每 T 步应用一次 Metropolis-Hastings 修正，降低其计算负担，从而分摊成本。
• 在不依赖递减步长或精确噪声方差估计的前提下，提供收敛速率和偏差消除的理论保证。
• 在合成分布、贝叶斯逻辑回归和贝叶斯神经网络上，实证验证 AMAGOLD 的鲁棒性和性能。

提出的方法

• AMAGOLD 在 SGHMC 更新序列中每 T 步应用一次 Metropolis-Hastings 修正，而非在每次迭代中应用，以校正随机梯度带来的偏差。
• 它采用带有动量变量和能量累加器的可逆或不可逆 SGHMC 变体，用于追踪更新过程中的能量变化。
• 使用完整的能量差值和累积动量项计算 Metropolis-Hastings 接受概率，确保细致平衡。
• 通过在 T 步内重用同一修正，分摊 M-H 修正的成本，降低每轮迭代的开销。
• 推导出基于谱隙的收敛速率上界，表明 AMAGOLD 的收敛速率在全批量基线的常数因子范围内。
• 该方法无需步长衰减或精确的噪声方差估计，提升了实际可用性。

实验结果

研究问题

• RQ1 是否可以在 SGHMC 中不频繁地应用 Metropolis-Hastings 修正，以在不产生禁止性计算成本的前提下保持渐近精确性？
• RQ2 在 T 步内分摊 M-H 修正是否能保持全批量 MCMC 的收敛特性，同时维持计算效率？
• RQ3 AMAGOLD 是否能在存在随机梯度的情况下，实现与全批量 HMC 相当的收敛速率（常数因子内）？
• RQ4 AMAGOLD 在超参数选择（尤其是步长）方面的鲁棒性是否优于 SGHMC？
• RQ5 AMAGOLD 在真实世界贝叶斯推断任务（如贝叶斯逻辑回归和贝叶斯神经网络）中是否能保持准确性？

主要发现

• AMAGOLD 在固定步长下收敛到真实后验分布，消除了 SGHMC 中为减少偏差而常用递减步长的需求。
• AMAGOLD 的收敛速率被限制在全批量基线的常数因子范围内，展示了理论上的高效性。
• AMAGOLD 在步长选择上的鲁棒性显著优于 SGHMC，在广泛步长范围内均表现出稳定性能。
• 在合成分布上，AMAGOLD 的 KL 散度低于 SGHMC，尤其在步长非最优时表现更优。
• 在贝叶斯逻辑回归和贝叶斯神经网络中，AMAGOLD 在采样精度和泛化能力上优于 SGHMC，同时在运行时间上与全批量基线方法保持相当。
• AMAGOLD 中的 M-H 接受概率在广泛步长范围内保持足够高，从而在避免过度拒绝的同时实现有效探索。