Hessian-informed, Coordinate Friendly Hamiltonian Monte Carlo in Linear Time

2026-06-04

AI Summary

This work addresses the scalability limitations of Riemannian Hamiltonian Monte Carlo (RHMC) in high dimensions, where fixed-point iterations require full Hessian computations with O(d²)–O(d³) complexity. Focusing on target distributions exhibiting “coordinate-friendly” structure—such as generalized linear models and dense or sparse graphical models—the authors propose a linear-time RHMC variant that leverages a diagonal Hessian approximation. By automatically detecting coordinate-friendly structure, restructuring the computational graph, and integrating automatic differentiation, the method reduces per-iteration complexity to O(d). This approach enables, for the first time, black-box, automated linear-time RHMC sampling while preserving the benefits of position-dependent preconditioning. Empirical results demonstrate that the proposed sampler yields higher-quality samples per unit of computation compared to existing HMC and NUTS algorithms, including those employing various preconditioning strategies.

Abstract

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.

Problem

• Hamiltonian Monte Carlo
• computational complexity
• high-dimensional inference
• Riemannian manifold
• preconditioning

Innovation

• Hamiltonian Monte Carlo
• Riemannian geometry
• diagonal Hessian preconditioning
• coordinate-friendly structure
• linear-time complexity