Riemannian Hamiltonian Monte Carlo (RHMC) is a promising MCMC methodology thanks to its ability to accommodate position-dependent preconditioning and multi-step proposals. While RHMC performs well in low dimensions, it becomes infeasible in high dimensions due to its $O(d^3)$ cost per fixed-point iteration, where $d$ is the dimension of the target density. Even when the position-dependent preconditioner is based on the diagonal of the Hessian, the cost is still $O(d^2)$ per fixed-point iteration. In this paper, we propose a computational method to reduce the computational complexity of RHMC fixed-point iterations with diagonal preconditioners from $O(d^2)$ to $O(d)$ for targets with ``coordinate friendly'' structures. This distribution class includes generalized linear models as well as other dense and sparse graphical models. The method is expressed as manipulating the compute graph and can therefore be automated to work on black box targets. Finally, we show empirically that our implementation of RHMC results in better sample quality per unit of compute time for various target distributions compared to state-of-the-art HMC NUTS algorithms with both position-independent and position-dependent preconditioners.

翻译：黎曼流形哈密顿蒙特卡罗（RHMC）是一种很有前景的马尔可夫链蒙特卡罗方法，因为它能够适应位置相关的预条件和多步提案。尽管RHMC在低维空间中表现良好，但在高维空间中，由于每次不动点迭代的计算成本为$O(d^3)$（其中$d$是目标密度的维度），变得不可行。即使位置相关预条件基于Hessian矩阵的对角线，每次不动点迭代的成本仍然为$O(d^2)$。本文提出了一种计算方法，将采用对角预条件器的RHMC不动点迭代的计算复杂度从$O(d^2)$降低至$O(d)$，适用于具有“坐标友好”结构的分布目标。该分布类别包括广义线性模型以及其他稠密和稀疏图模型。该方法表现为对计算图的操纵，因此可以自动化地应用于黑箱目标。最后，我们通过实验表明，与最先进的采用位置无关和位置相关预条件器的HMC NUTS算法相比，我们的RHMC实现在各种目标分布上能在单位计算时间内获得更好的样本质量。