ATLAS: Adapting Trajectory Lengths and Step-Size for Hamiltonian Monte Carlo

TL;DR

ATLAS introduces an adaptive Hamiltonian Monte Carlo method that dynamically adjusts step size and trajectory length based on local geometry, improving sampling accuracy for complex distributions while maintaining computational efficiency.

Contribution

The paper presents ATLAS, a novel adaptive HMC algorithm that locally tunes step size and trajectory length, enhancing sampling of complex geometries and robustness to hyperparameter tuning.

Findings

ATLAS accurately samples complex distributions with varying curvature.

It is computationally competitive with NUTS on simpler problems.

ATLAS demonstrates greater robustness to hyperparameter tuning.

Abstract

Hamiltonian Monte-Carlo (HMC) and its auto-tuned variant, the No U-Turn Sampler (NUTS) can struggle to accurately sample distributions with complex geometries, e.g., varying curvature, due to their constant step size for leapfrog integration and fixed mass matrix. In this work, we develop a strategy to locally adapt the step size parameter of HMC at every iteration by evaluating a low-rank approximation of the local Hessian and estimating its largest eigenvalue. We combine it with a strategy to similarly adapt the trajectory length by monitoring the no U-turn condition, resulting in an adaptive sampler, ATLAS: adapting trajectory length and step-size. We further use a delayed rejection framework for making multiple proposals that improves the computational efficiency of ATLAS, and develop an approach for automatically tuning its hyperparameters during warmup. We compare ATLAS with state-of-the-art samplers like NUTS on a suite of synthetic and real world examples, and show that i) unlike NUTS, ATLAS is able to accurately sample difficult distributions with complex geometries, ii) it is computationally competitive to NUTS for simpler distributions, and iii) it is more robust to the tuning of hyperparamters.