Iterated Energy-based Flow Matching for Sampling from Boltzmann Densities

TL;DR

This paper introduces iterated energy-based flow matching (iEFM), a novel off-policy method for training continuous normalizing flows from unnormalized energy functions, improving efficiency in high-dimensional probabilistic modeling.

Contribution

We propose the first off-policy approach, iEFM, for training CNF models from unnormalized densities, extending to VE and OT paths, with demonstrated superior performance.

Findings

iEFM outperforms existing methods on GMM and DW-4 energy functions.

The framework is general and adaptable to various probabilistic paths.

Results show improved scalability and efficiency in high-dimensional systems.

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.