[Paper Review] Selecting the Metric in Hamiltonian Monte Carlo

Ben Bales, Arya Pourzanjani+2 more|arXiv (Cornell University)|May 28, 2019

Markov Chains and Monte Carlo Methods21 references4 citations

TL;DR

This paper introduces a selection criterion to automatically choose an optimal Euclidean metric for Hamiltonian Monte Carlo (HMC) based on posterior geometry, and proposes a low-rank Hessian-based metric adaptation that requires significantly fewer warmup draws than traditional sample covariance methods. The method improves sampling efficiency on correlated posteriors, especially with limited warmup data, and is computationally feasible via the Lanczos algorithm.

ABSTRACT

We present a selection criterion for the Euclidean metric adapted during warmup in a Hamiltonian Monte Carlo sampler that makes it possible for a sampler to automatically pick the metric based on the model and the availability of warmup draws. Additionally, we present a new adaptation inspired by the selection criterion that requires significantly fewer warmup draws to be effective. The effectiveness of the selection criterion and adaptation are demonstrated on a number of applied problems. An implementation for the Stan probabilistic programming language is provided.

Motivation & Objective

• To address the challenge of selecting an effective Euclidean metric in HMC that adapts to posterior geometry without requiring high-dimensional derivatives or complex solvers.
• To reduce the number of warmup draws needed for effective metric adaptation, particularly in high-correlation settings.
• To develop a practical, scalable method for metric selection and adaptation that integrates seamlessly into probabilistic programming frameworks like Stan.
• To demonstrate that low-rank Hessian approximations can outperform standard diagonal and full covariance metrics in terms of effective sample size per second (ESS/s), especially with short warmup phases.

Proposed method

• The selection criterion evaluates candidate metrics using the eigenvalues of the Hessian of the log-posterior at the MAP estimate, derived from linearized Hamiltonian dynamics.
• The criterion computes a stability proxy based on the eigenvalues of the matrix $ L^T \nabla^2_{qq}H(q_0)L $, where $ L $ is the Cholesky factor of the metric inverse.
• A new low-rank metric adaptation is proposed using rank-1 to rank-8 Hessian approximations, computed efficiently via the Lanczos algorithm.
• The method incorporates an inverse-Wishart update to stabilize the Hessian-based metric estimate during adaptation.
• The selection criterion is used to dynamically switch between metric types (diagonal, dense, low-rank) during warmup, choosing the one with the lowest criterion value.
• The approach is implemented in Stan and evaluated on benchmark models using ESS/s and $\hat{R}$ as performance metrics.

Experimental results

Research questions

• RQ1 Can a selection criterion be derived that reliably predicts the effectiveness of a metric in HMC without requiring full posterior sampling?
• RQ2 Can a low-rank Hessian-based metric adaptation achieve high sampling efficiency with fewer warmup draws than sample covariance or diagonal approximations?
• RQ3 Does the proposed metric adaptation outperform standard Stan defaults (diagonal and dense) in terms of effective sample size per second on correlated posteriors?
• RQ4 How does the inclusion of an inverse-Wishart update affect the stability and performance of the Hessian-based metric adaptation?
• RQ5 Can the selection criterion enable automatic, dynamic metric switching during warmup to improve HMC efficiency?

Key findings

• The selection criterion successfully identifies high-performing metrics, with lower values indicating better HMC performance, and is effective across diverse models including Radon and Prophet.
• The switching adaptation, which selects among diagonal, dense, and low-rank Hessian metrics based on the selection criterion, achieves ESS/s values comparable to or better than fixed metrics.
• On the Radon model, the switching adaptation achieved a minimum ESS/s of 370 with only 75 warmup draws, outperforming the diagonal metric under the same conditions.
• The low-rank Hessian adaptation (rank 1–8) required significantly fewer warmup draws than full sample covariance to achieve effective sampling, especially in high-correlation settings.
• The inverse-Wishart update improved stability of the Hessian-based metric but did not guarantee performance gains in all cases, as shown in Figure 4.
• The Lanczos algorithm enabled efficient computation of required eigenvalues and eigenvectors, making the method scalable to high-dimensional models.