Iterated Denoising Energy Matching for Sampling from Boltzmann Densities

TL;DR

The paper introduces iDEM, an iterative, diffusion-based algorithm that efficiently generates independent samples from unnormalized distributions without data or MCMC, improving training speed and scalability for complex systems.

Contribution

iDEM is a novel, simulation-free method that leverages energy functions and diffusion to train scalable samplers faster than existing approaches.

Findings

Achieves state-of-the-art performance on various energy functions.

Trains 2-5 times faster than previous methods.

First to train on the 55-particle Lennard-Jones system.

Abstract

Efficiently generating statistically independent samples from an unnormalized probability distribution, such as equilibrium samples of many-body systems, is a foundational problem in science. In this paper, we propose Iterated Denoising Energy Matching (iDEM), an iterative algorithm that uses a novel stochastic score matching objective leveraging solely the energy function and its gradient -- and no data samples -- to train a diffusion-based sampler. Specifically, iDEM alternates between (I) sampling regions of high model density from a diffusion-based sampler and (II) using these samples in our stochastic matching objective to further improve the sampler. iDEM is scalable to high dimensions as the inner matching objective, is simulation-free, and requires no MCMC samples. Moreover, by leveraging the fast mode mixing behavior of diffusion, iDEM smooths out the energy landscape enabling efficient exploration and learning of an amortized sampler. We evaluate iDEM on a suite of tasks ranging from standard synthetic energy functions to invariant $n$-body particle systems. We show that the proposed approach achieves state-of-the-art performance on all metrics and trains $2-5\times$ faster, which allows it to be the first method to train using energy on the challenging $55$-particle Lennard-Jones system.