Iterated Energy-based Flow Matching for Sampling from Boltzmann Densities

Dongyeop Woo,Sungsoo Ahn
2024-08-29
Abstract:In this work, we consider the problem of training a generator from evaluations of energy functions or unnormalized densities. This is a fundamental problem in probabilistic inference, which is crucial for scientific applications such as learning the 3D coordinate distribution of a molecule. To solve this problem, we propose iterated energy-based flow matching (iEFM), the first off-policy approach to train continuous normalizing flow (CNF) models from unnormalized densities. We introduce the simulation-free energy-based flow matching objective, which trains the model to predict the Monte Carlo estimation of the marginal vector field constructed from known energy functions. Our framework is general and can be extended to variance-exploding (VE) and optimal transport (OT) conditional probability paths. We evaluate iEFM on a two-dimensional Gaussian mixture model (GMM) and an eight-dimensional four-particle double-well potential (DW-4) energy function. Our results demonstrate that iEFM outperforms existing methods, showcasing its potential for efficient and scalable probabilistic modeling in complex high-dimensional systems.
Machine Learning
What problem does this paper attempt to address?
The problem that this paper attempts to solve is sampling from energy functions or unnormalized densities. Specifically, it is to efficiently generate samples from the Boltzmann distribution. In many scientific applications, such as learning the three - dimensional coordinate distribution of molecules, this is a fundamental task. However, the existing deep generative models are not suitable for such problems because the cost of creating data sets is high, and traditional simulation methods (such as molecular dynamics simulation or Monte Carlo techniques) are computationally expensive and not easily scalable to high - dimensional systems. To this end, the paper proposes iterated energy - based flow matching (iEFM), which is an off - policy method for training continuous normalizing flow (CNF) models from unnormalized densities for the first time. iEFM trains the model by estimating the Monte Carlo estimate of the data - generating vector field, enabling it to predict the marginal vector field constructed from a known energy function. This framework is general and scalable and is applicable to variance - exploding (VE) and optimal transport (OT) conditional probability paths. The paper evaluates the performance of iEFM on a two - dimensional Gaussian mixture model (GMM) and an eight - dimensional four - particle double - well potential (DW - 4) energy function. The results show that iEFM outperforms existing methods, demonstrating its potential for efficient and scalable probability modeling in complex high - dimensional systems.