Energy-Weighted Flow Matching: Unlocking Continuous Normalizing Flows for Efficient and Scalable Boltzmann Sampling
This addresses the computational bottleneck for scientists and engineers in fields like molecular dynamics by providing a more efficient method for Boltzmann sampling, though it is incremental as it builds on existing flow matching techniques.
The paper tackled the problem of sampling from Boltzmann distributions, which is computationally challenging due to complex energy landscapes, by introducing Energy-Weighted Flow Matching (EWFM) to enable continuous normalizing flows to model these distributions using only energy evaluations, achieving sample quality competitive with state-of-the-art methods while reducing energy evaluations by up to three orders of magnitude.
Sampling from unnormalized target distributions, e.g. Boltzmann distributions $μ_{\text{target}}(x) \propto \exp(-E(x)/T)$, is fundamental to many scientific applications yet computationally challenging due to complex, high-dimensional energy landscapes. Existing approaches applying modern generative models to Boltzmann distributions either require large datasets of samples drawn from the target distribution or, when using only energy evaluations for training, cannot efficiently leverage the expressivity of advanced architectures like continuous normalizing flows that have shown promise for molecular sampling. To address these shortcomings, we introduce Energy-Weighted Flow Matching (EWFM), a novel training objective enabling continuous normalizing flows to model Boltzmann distributions using only energy function evaluations. Our objective reformulates conditional flow matching via importance sampling, allowing training with samples from arbitrary proposal distributions. Based on this objective, we develop two algorithms: iterative EWFM (iEWFM), which progressively refines proposals through iterative training, and annealed EWFM (aEWFM), which additionally incorporates temperature annealing for challenging energy landscapes. On benchmark systems, including challenging 55-particle Lennard-Jones clusters, our algorithms demonstrate sample quality competitive with state-of-the-art energy-only methods while requiring up to three orders of magnitude fewer energy evaluations.