Generalizing Hamiltonian Monte Carlo with Neural Networks

TL;DR

This paper introduces a neural network-based extension of Hamiltonian Monte Carlo that significantly improves sampling efficiency and mixing speed across various challenging distributions and real-world tasks.

Contribution

It proposes a novel neural network parameterization for MCMC kernels that generalizes HMC and maximizes mixing speed, with demonstrated empirical improvements.

Findings

Achieved up to 106x increase in effective sample size.

Enabled mixing where standard HMC fails.

Showed benefits on real-world latent-variable modeling.

Abstract

We present a general-purpose method to train Markov chain Monte Carlo kernels, parameterized by deep neural networks, that converge and mix quickly to their target distribution. Our method generalizes Hamiltonian Monte Carlo and is trained to maximize expected squared jumped distance, a proxy for mixing speed. We demonstrate large empirical gains on a collection of simple but challenging distributions, for instance achieving a 106x improvement in effective sample size in one case, and mixing when standard HMC makes no measurable progress in a second. Finally, we show quantitative and qualitative gains on a real-world task: latent-variable generative modeling. We release an open source TensorFlow implementation of the algorithm.