[论文解读] Introducing an Explicit Symplectic Integration Scheme for Riemannian Manifold Hamiltonian Monte Carlo

Adam D. Cobb, Atılım Güneş Baydin+2 more|arXiv (Cornell University)|Oct 14, 2019

Model Reduction and Neural Networks参考文献 33被引用 18

一句话总结

本文提出了一种显式辛积分器用于黎曼流形哈密顿蒙特卡洛（RMHMC），减少了在跃迁步长中进行迭代定点求解的需求。通过最小化高阶导数的计算，该方法在保持采样精度的同时，使运行时间几乎缩短为原来的一半，速度接近隐式RMHMC的两倍，且在贝叶斯逻辑回归和层次模型上通过开源包 hamiltorch 得到了验证。

ABSTRACT

We introduce a recent symplectic integration scheme derived for solving physically motivated systems with non-separable Hamiltonians. We show its relevance to Riemannian manifold Hamiltonian Monte Carlo (RMHMC) and provide an alternative to the currently used generalised leapfrog symplectic integrator, which relies on solving multiple fixed point iterations to convergence. Via this approach, we are able to reduce the number of higher-order derivative calculations per leapfrog step. We explore the implications of this integrator and demonstrate its efficacy in reducing the computational burden of RMHMC. Our code is provided in a new open-source Python package, hamiltorch.

研究动机与目标

• 为解决RMHMC计算成本过高的问题，其根源在于对非可分哈密顿量求解定点迭代。
• 减少RMHMC中每步跃迁步长的高阶导数计算次数。
• 开发一种更高效的显式辛积分器，避免现有隐式方案中使用的迭代求解器。
• 提升RMHMC在复杂模型（如贝叶斯神经网络）中的可扩展性和实用性。
• 通过统一的API接口，提供用户友好的、GPU加速的实现，借助hamiltorch Python包实现更广泛的采用。

提出的方法

• 提出一种专为RMHMC中非可分哈密顿量设计的显式辛积分方案。
• 用无需多次定点迭代的直接显式更新规则，替代需要多次定点迭代的隐式广义跃迁积分器。
• 利用PyTorch的自动微分高效计算所需的黎曼几何量（度量张量、克里斯托费尔符号）。
• 设计积分器以保持辛结构，同时最小化高阶导数带来的计算开销。
• 在hamiltorch包中实现该方法，支持通过最小代码修改实现基于GPU的采样。
• 通过统一的API支持在HMC、RMHMC和显式RMHMC之间无缝切换。

实验结果

研究问题

• RQ1 显式辛积分器是否能在不牺牲采样精度的前提下，降低RMHMC的计算负担？
• RQ2 在运行时间与收敛性方面，显式RMHMC与隐式RMHMC相比表现如何？
• RQ3 新积分器在每步跃迁中将高阶导数计算次数减少了多少？
• RQ4 在具有挑战性的后验几何结构（如漏斗分布）上，显式RMHMC是否能实现与隐式RMHMC相当的后验近似质量？
• RQ5 在真实世界模型（如贝叶斯逻辑回归和层次模型）中，该方法在实际应用中的可扩展性如何？

主要发现

• 与隐式RMHMC相比，显式RMHMC实现了接近2倍的运行时间加速，收敛时间减少了约三分之二。
• 尽管速度显著提升，显式RMHMC仍保持了相近的采样性能，其KL散度为0.142，与隐式RMHMC的0.130相近。
• 有效样本量方面，两种RMHMC变体均优于HMC或NUTS，表现为自相关函数衰减更快，DKL(p(v)∥q(v))随时间更快趋于平稳。
• 该方法通过避免迭代定点求解，直接降低了每步跃迁的高阶导数计算次数，从而显著降低计算成本。
• 在含431,080个参数的贝叶斯神经网络上，hamiltorch在GPU上实现了每秒116.33个样本，而CPU上仅为13.92个，展现出强大的硬件加速能力。
• 显式RMHMC估计的边缘后验分布偏差更低，且在高曲率区域（如漏斗的狭窄颈部）与真实后验分布的吻合度更高。