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, Avisek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera, Siamak Ravanbakhsh, Youshua Bengio, Nikolay Malkin, Alexander Tong
This paper introduces ITERATED DENOISING ENERGY MATCHING (iDEM), a novel algorithm for sampling from Boltzmann distributions using only the energy function and its gradient. iDEM is an iterative algorithm that alternates between sampling from a diffusion-based model and using these samples to improve the model through a stochastic matching objective. The algorithm is scalable, simulation-free, and does not require MCMC samples. By leveraging the fast mixing behavior of diffusion, iDEM smooths the energy landscape, enabling efficient exploration and learning of an amortized sampler. iDEM is evaluated on various tasks, including synthetic energy functions and invariant n-body particle systems. It achieves state-of-the-art performance and trains 2-5 times faster than previous methods, making it the first to successfully scale to the challenging 55-particle Lennard-Jones system using energy-based training. The algorithm is also designed to incorporate symmetries inherent in physical systems, making it well-suited for scientific applications. iDEM uses a bi-level iterative scheme, where the inner loop trains a diffusion sampler using a stochastic regression objective on the energy function, and the outer loop uses reverse SDE simulations to amortize sampling and improve model performance. The algorithm is shown to be effective in accurately modeling complex energy landscapes and outperforms other methods in terms of sample quality and computational efficiency.This paper introduces ITERATED DENOISING ENERGY MATCHING (iDEM), a novel algorithm for sampling from Boltzmann distributions using only the energy function and its gradient. iDEM is an iterative algorithm that alternates between sampling from a diffusion-based model and using these samples to improve the model through a stochastic matching objective. The algorithm is scalable, simulation-free, and does not require MCMC samples. By leveraging the fast mixing behavior of diffusion, iDEM smooths the energy landscape, enabling efficient exploration and learning of an amortized sampler. iDEM is evaluated on various tasks, including synthetic energy functions and invariant n-body particle systems. It achieves state-of-the-art performance and trains 2-5 times faster than previous methods, making it the first to successfully scale to the challenging 55-particle Lennard-Jones system using energy-based training. The algorithm is also designed to incorporate symmetries inherent in physical systems, making it well-suited for scientific applications. iDEM uses a bi-level iterative scheme, where the inner loop trains a diffusion sampler using a stochastic regression objective on the energy function, and the outer loop uses reverse SDE simulations to amortize sampling and improve model performance. The algorithm is shown to be effective in accurately modeling complex energy landscapes and outperforms other methods in terms of sample quality and computational efficiency.
Reach us at info@study.space