Self-tuning Hamiltonian Monte Carlo for accelerated sampling

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 that 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 commonly used 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 speedup in the 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.