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

2026-05-05Machine Learning

Machine LearningArtificial Intelligence
AI summary

The authors present Flow Sampling, a new method to generate samples from complex probability distributions defined by energy functions, without needing actual data samples. Their approach builds on diffusion models but flips the usual setup by focusing on denoising guided by the energy function, making the sampling more efficient. They also extend their method to work on curved spaces like spheres and hyperbolic geometry, which are important in some scientific problems. Tests show that this method can effectively sample from tough energy distributions and has practical applications, for example, in modeling molecules.

unnormalized densitiesenergy functionsamplingdiffusion modelsflow matchingdenoising diffusionRiemannian manifoldshypersphereshyperbolic spacesmolecular conformer generation
Authors
Aaron Havens, Brian Karrer, Neta Shaul
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.