Flow Sampling: Learning to Sample from Unnormalized Densities via Denoising Conditional Processes
Aaron Havens, Brian Karrer, Neta Shaul
TLDR
Flow Sampling uses diffusion models and flow matching to efficiently sample from unnormalized densities, even on Riemannian manifolds.
Key contributions
- Introduces Flow Sampling, a diffusion-based framework for efficient sampling from unnormalized densities.
- Utilizes an interpolant process to minimize costly energy function evaluations during training.
- Extends sampling to Riemannian manifolds, deriving conditional drifts for constant curvature spaces.
Why it matters
This paper introduces an efficient and scalable method for sampling from unnormalized densities, a core problem in generative modeling and scientific computing. By extending diffusion-based sampling to Riemannian manifolds, it significantly broadens the applicability of these powerful models. This could advance fields like molecular conformer generation and statistical physics.
Original 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.
📬 Weekly AI Paper Digest
Get the top 10 AI/ML arXiv papers from the week — summarized, scored, and delivered to your inbox every Monday.