Flow Sampling: Learning to Sample from Unnormalized Densities via Denoising Conditional Processes

TL;DR

Flow Sampling introduces a diffusion-based framework for efficiently sampling from unnormalized densities, extending to Riemannian manifolds and demonstrating strong empirical results across various applications.

Contribution

It presents a novel training objective conditioned on noise for energy-based sampling, reducing energy evaluations and enabling sampling on complex geometries.

Findings

Efficient sampling with fewer energy function evaluations.

Extension of diffusion models to Riemannian manifolds.

Strong empirical performance on diverse benchmarks.

Abstract

Sampling from unnormalized densities is analogous to the generative modeling problem, but the target distribution is defined by a known energy function instead of data samples. Because evaluating the energy function is often costly, a primary challenge is to learn an efficient sampler. We introduce Flow Sampling, a framework built on diffusion models and flow matching for the data-free setting. Our training objective is conditioned on a noise sample and regresses onto a denoising diffusion drift constructed from the energy function. In contrast, diffusion models' objective is conditioned on a data sample and regresses onto a noising diffusion drift. We utilize the interpolant process to minimize the number of energy function evaluations during training, resulting in an efficient and scalable method for sampling unnormalized densities. Furthermore, our formulation naturally extends to Riemannian manifolds, enabling diffusion-based sampling in geometries beyond Euclidean space. We derive a closed-form formula for the conditional drift on constant curvature manifolds, including hyperspheres and hyperbolic spaces. We evaluate Flow Sampling on synthetic energy benchmarks, small peptides, large-scale amortized molecular conformer generation, and distributions supported on the sphere, demonstrating strong empirical performance.