Iterated Denoising Energy Matching for Sampling from Boltzmann Densities

Iterated Denoising Energy Matching for Sampling from Boltzmann Densities

2024 | Tara Akhound-Sadegh, Jarrid Rector-Brooks, Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera, Siamak Ravanbakhsh, Gauthier Gidel, Yoshua Bengio, Nikolay Malkin, Alexander Tong
This paper introduces Iterated Denoising Energy Matching (iDEM), an iterative algorithm designed to efficiently generate statistically independent samples from unnormalized probability distributions, particularly equilibrium samples of many-body systems. iDEM leverages a novel stochastic score matching objective that uses only the energy function and its gradient, without requiring data samples. The algorithm alternates between two main steps: (1) sampling regions of high model density from a diffusion-based sampler and (2) using these samples to improve the sampler through a stochastic matching objective. iDEM is scalable to high dimensions because the inner matching objective is simulation-free and does not require MCMC samples. By exploiting the fast mode mixing behavior of diffusion, iDEM smooths the energy landscape, enabling efficient exploration and learning of an amortized sampler. The paper evaluates iDEM on various tasks, including standard synthetic energy functions and invariant $n$-body particle systems, demonstrating state-of-the-art performance and significantly faster training compared to previous methods. Notably, iDEM is the first method to successfully scale to the challenging 55-particle Lennard-Jones system using energy-based training.This paper introduces Iterated Denoising Energy Matching (iDEM), an iterative algorithm designed to efficiently generate statistically independent samples from unnormalized probability distributions, particularly equilibrium samples of many-body systems. iDEM leverages a novel stochastic score matching objective that uses only the energy function and its gradient, without requiring data samples. The algorithm alternates between two main steps: (1) sampling regions of high model density from a diffusion-based sampler and (2) using these samples to improve the sampler through a stochastic matching objective. iDEM is scalable to high dimensions because the inner matching objective is simulation-free and does not require MCMC samples. By exploiting the fast mode mixing behavior of diffusion, iDEM smooths the energy landscape, enabling efficient exploration and learning of an amortized sampler. The paper evaluates iDEM on various tasks, including standard synthetic energy functions and invariant $n$-body particle systems, demonstrating state-of-the-art performance and significantly faster training compared to previous methods. Notably, iDEM is the first method to successfully scale to the challenging 55-particle Lennard-Jones system using energy-based training.
Reach us at info@study.space