Self-Tuning Hamiltonian Monte Carlo for Accelerated Sampling

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.

References in corpus (12)

• Adam: A Method for Stochastic Optimization
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• MCMC using Hamiltonian dynamics
• The No-U-Turn Sampler: Adaptively Setting Path Lengths in Hamiltonian Monte Carlo
• Normalizing Flows for Probabilistic Modeling and Inference
• A Gentle Introduction to Deep Learning for Graphs
• TorchMD: A deep learning framework for molecular simulations
• Optimal scalings for local Metropolis--Hastings chains on nonproduct targets in high dimensions
• Monte Carlo Simulation of Dense Polymer Melts Using Event Chain Algorithms
• Generalizing Hamiltonian Monte Carlo with Neural Networks
• Policy Guided Monte Carlo: Reinforcement Learning Markov Chain Dynamics
• Focusing on Difficult Directions for Learning HMC Trajectory Lengths

Cited by in corpus (4)

• Uncertainty-biased molecular dynamics for learning uniformly accurate interatomic potentials
• Roadmap on machine learning glassy dynamics
• Policy-guided Monte Carlo on general state spaces: Application to glass-forming mixtures
• Fast, Modular, and Differentiable Framework for Machine Learning-Enhanced Molecular Simulations