Continuously Tempered Diffusion Samplers

Ezra Erives, Bowen Jing, Peter Holderrieth, Tommi Jaakkola

Keywords: neural sampler, Langevin dynamics, PINN, physics-informed neural network, variational inference, diffusion models

TL;DR: A fusion of MD-inspired continuous tempering techniques and annealed neural sampler training.

Abstract: Annealing-based neural samplers seek to amortize sampling from unnormalized distributions by training neural networks to transport a family of densities interpolating from source to target. A crucial design choice in the training phase of such samplers is the proposal distribution by which locations are generated at which to evaluate the loss. Previous work has obtained such a proposal distribution by combining a partially learned transport with annealed Langevin dynamics. However, isolated modes and other pathological properties of the annealing path imply that such proposals achieve insufficient exploration and thereby lower performance post training. To remedy this, we propose continuously tempered diffusion samplers, which leverage exploration techniques developed in the context of molecular dynamics to improve proposal distributions. Specifically, a family of distributions across different temperatures is introduced to lower energy barriers at higher temperatures and drive exploration at the lower temperature of interest. We empirically validate improved sampler performance driven by extended exploration.

Submission Number: 81