Self-Tuning Hamiltonian Monte Carlo for Accelerated Sampling

TL;DR

This paper introduces an adaptive framework for tuning Hamiltonian Monte Carlo parameters automatically, using a differentiable loss function that correlates with autocorrelation times, leading to faster sampling in molecular simulations.

Contribution

The authors develop a novel, gradient-based, adaptive tuning method for HMC parameters that improves sampling efficiency and can learn a distribution over integration steps.

Findings

Achieved over 100-fold speed-up in parameter optimization for alanine dipeptide.

Extended the integrator to include atom-dependent timesteps, reducing autocorrelation times by 25%.

Demonstrated the method's effectiveness on harmonic oscillator and alanine dipeptide systems.

Abstract

The performance of Hamiltonian Monte Carlo simulations crucially depends on both the integration timestep and the number of integration steps. We present an adaptive general-purpose framework to automatically tune such parameters, based on a local loss function which promotes the fast exploration of phase-space. We show that a good correspondence between loss and autocorrelation time can be established, allowing for gradient-based optimization using a fully-differentiable set-up. The loss is constructed in such a way that it also allows for gradient-driven learning of a distribution over the number of integration steps. Our approach is demonstrated for the one-dimensional harmonic oscillator and alanine dipeptide, a small protein common as a test case for simulation methods. Through the application to the harmonic oscillator, we highlight the importance of not using a fixed timestep to avoid a rugged loss surface with many local minima, otherwise trapping the optimization. In the case of alanine dipeptide, by tuning the only free parameter of our loss definition, we find a good correspondence between it and the autocorrelation times, resulting in a $>100$ fold speed up in optimization of simulation parameters compared to a grid-search. For this system, we also extend the integrator to allow for atom-dependent timesteps, providing a further reduction of $25\%$ in autocorrelation times.