Towards Enhanced Proposals for PINN-Based Neural Sampler Training

Abstract

Sampling from distributions whose density is known up to a normalizing constant is an important problem with a wide range of applications including Bayesian posterior inference, statistical physics, and structural biology. 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 vector field 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. In this work we extend existing work and characterize new families of proposals based on controlled Langevin dynamics. In particular, 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 additionally explore proposals based on Langevin dynamics involving non-Newtonian kinetic energies. We empirically validate improved sampler performance driven by extended exploration.