🤖 machine learning

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

This paper introduces Flow Sampling, a data-free framework leveraging diffusion models and flow matching to efficiently learn samplers for unnormalized densities by regressing onto denoising drifts derived from energy functions, while also extending naturally to Riemannian manifolds with closed-form solutions for constant curvature spaces.

Original authors: Aaron Havens, Brian Karrer, Neta Shaul

Published 2026-05-06

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Aaron Havens, Brian Karrer, Neta Shaul

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find the best seats in a massive, dark theater. You know exactly where the "good" seats are because you have a map that tells you the "energy" or "comfort level" of every single seat. However, there's a catch: you don't know how many good seats there are in total (the normalization constant), and you can't just walk through the theater to see them all because the map is too expensive to read for every single seat.

This is the problem scientists face when trying to sample from unnormalized densities. They have a rule (an energy function) that says what a "good" sample looks like, but they can't easily generate those samples.

Enter Flow Sampling, a new method introduced by Aaron Havens, Brian Karrer, and Neta Shaul. Here is how it works, explained through simple analogies.

The Old Way: The Slow Hiker

Previously, scientists used methods like MCMC (Markov Chain Monte Carlo). Imagine a hiker trying to find the best seats. The hiker takes one step, checks the map, takes another step, checks the map again, and so on. They eventually find the good seats, but they do it one by one, very slowly. If you need a million samples, the hiker takes a million steps. This is too slow for big problems like designing new materials or drugs.

The New Way: The "Denoising" Sculptor

Flow Sampling is like hiring a sculptor who learns to carve the perfect statue (the target distribution) without needing to see the original statue first. Instead, they learn by watching how a block of clay turns into a statue if you slowly remove the noise.

Here is the step-by-step magic:

1. The Reverse Movie (Denoising)

Imagine you have a clear, high-definition photo of a beautiful landscape (the target data). If you add static noise to it, it becomes a blurry mess.

Standard AI usually learns by starting with the blurry mess and trying to guess the original photo.
Flow Sampling flips the script. It starts with a random noise sample (a blank canvas) and learns how to "clean" it up to match the landscape, guided by the energy map.

2. The "Detached" Student (Saving Money)

The biggest cost in this process is reading the energy map. Reading it is like paying a toll. If you have to read the map millions of times to train the AI, it becomes too expensive.

Flow Sampling introduces a clever trick called the "Detached State" or a Replay Buffer:

Imagine a student (the AI model) trying to learn. Instead of asking the teacher (the energy map) for help every single time the student makes a mistake, the student takes a "snapshot" of their current work, asks the teacher for feedback once on that snapshot, and writes the answer down in a notebook (the Replay Buffer).
The student then practices using that notebook over and over again without bothering the teacher.
This means the expensive "energy map" is only consulted a few times, while the student learns from the notes. This cuts the cost of training by 4 to 8 times compared to previous methods.

3. The "Flow" (Smooth Movement)

Instead of the hiker taking tiny, shaky steps, Flow Sampling uses a "flow." Imagine a river. The water (the data) flows smoothly from a source (random noise) to a destination (the target distribution). The AI learns the direction of the river current so it can guide the water exactly where it needs to go, efficiently and quickly.

Going Beyond Flat Ground (Curved Spaces)

Most AI assumes the world is flat (like a sheet of paper). But sometimes, data lives on curved surfaces, like the surface of a sphere (think of the Earth) or a saddle shape.

Flow Sampling is special because it naturally understands these curves.

The Analogy: If you are walking on a flat floor, you walk in a straight line. If you are walking on a sphere (like the Earth), the "straightest" path is a curve (a great circle).
The paper provides a mathematical formula that tells the AI exactly how to walk these curved paths (geodesics) without getting lost. This allows the method to work on spheres and hyperbolic spaces, which are crucial for things like robotics and 3D molecular structures.

What Did They Prove?

The authors tested this method on several challenges:

Synthetic Puzzles: They created fake energy landscapes and showed Flow Sampling found the best spots faster and more accurately than other top methods.
Protein Folding: They used it to find the shapes of small proteins (peptides). It successfully predicted the 3D shapes of these molecules, matching the results of much slower, traditional simulations.
Molecule Generation: They used it to generate thousands of different shapes for large molecules. It was much faster than the competition while still finding high-quality shapes.
Sphere Sampling: They demonstrated it could sample data distributed on a sphere, proving it works on curved geometries.

The Bottom Line

Flow Sampling is a new, efficient way to teach computers to generate complex data when you only have a rule for what "good" looks like, but not a list of examples. By using a "notebook" to remember expensive calculations and by understanding curved spaces, it allows scientists to generate high-quality samples (like new drug molecules) much faster and cheaper than before.

Technical Summary: Flow Sampling

Problem Statement

The paper addresses the challenge of sampling from high-dimensional probability distributions defined only up to a normalization constant (unnormalized densities). Such distributions are ubiquitous in computational sciences, defined by an energy function $r(x)$ where the target density is $q(x) \propto \exp(r(x))$ . While evaluating $r(x)$ and its gradient $\nabla r(x)$ is possible, it is often computationally expensive, and direct access to samples from $q(x)$ is unavailable.

Traditional Markov Chain Monte Carlo (MCMC) methods, such as Langevin dynamics, converge asymptotically but suffer from slow mixing and sequential generation, making them inefficient for applications requiring fast, reusable samplers (e.g., materials design). Existing learning-based approaches, including diffusion models and flow matching, typically require data samples from the target distribution or rely on stochastic corrections (e.g., importance weighting, resampling, or auxiliary MCMC) that incur significant computational costs due to repeated energy function evaluations.

Methodology: Flow Sampling

The authors introduce Flow Sampling, a framework for learning amortized samplers directly from unnormalized density functions without requiring target data samples. The method is built upon diffusion models and flow matching but adapts the training objective to the data-free setting.

Core Concept

Unlike standard flow matching, which conditions on a data sample $X_1 \sim q$ to learn a noising drift, Flow Sampling conditions on a noise sample $X_0 \sim p_0$ (from an easy-to-sample source distribution) to learn a denoising drift. The marginal distribution of this learned process is designed to match the target $q$ by construction.

Training Objective

The framework utilizes a conditional probability path defined as a push-forward of the target distribution conditioned on the source:
$p_{t|0}(x|x_0) = \frac{1}{\alpha_t^d} q\left(\frac{x - \sigma_t x_0}{\alpha_t}\right)$
The training objective involves regressing a learned drift $u_\theta$ onto a conditional drift $u_{t|0}$ constructed from the energy function. A key theoretical insight (Proposition 3.2) provides a closed-form expression for this conditional drift along an interpolant process $X_t = \sigma_t X_0 + \alpha_t X_1$ :
$u_{t|0}(X_t|x_0) = \dot{\alpha}_t X_1 + \dot{\sigma}_t x_0 + \frac{g_t^2}{2\alpha_t} \nabla r(X_1)$
This formulation allows the model to reuse the energy gradient $\nabla r(X_1)$ evaluated at a single point $X_1$ for all time steps $t \in [0, 1]$ , significantly reducing the number of energy evaluations required during training.

Algorithm: Fixed-Point Iteration

Since target samples $X_1 \sim q$ are unavailable, the authors propose a fixed-point iteration training algorithm (Algorithm 1) consisting of two phases:

Exploration Phase: Using the current detached model state $\bar{\theta}$ , the algorithm generates samples $X_1^{\bar{\theta}}$ via simulation (Euler–Maruyama), evaluates the energy gradient $\nabla r(X_1^{\bar{\theta}})$ , and stores the pairs $(X_1^{\bar{\theta}}, \nabla r(X_1^{\bar{\theta}}))$ in a replay buffer.
Optimization Phase: The model $u_\theta$ is trained by minimizing the mean squared error between its predicted drift and the closed-form target derived from the samples in the replay buffer. This process repeats until convergence.

Extension to Riemannian Manifolds

The framework naturally extends to non-Euclidean geometries. By replacing the affine interpolant with a geodesic interpolant and deriving a closed-form formula for the conditional drift on constant curvature manifolds (including hyperspheres and hyperbolic spaces), the method enables diffusion-based sampling on Riemannian manifolds without reparameterization.

Key Contributions

Flow Sampling Framework: A principled method for learning diffusion samplers from unnormalized densities by matching conditional denoising diffusion drifts, eliminating the need for target data samples.
Efficient Training Objective: Derivation of a simple closed-form regression target that reuses endpoint energy gradients along the interpolant process, enabling an efficient replay-buffer training algorithm that drastically reduces energy function evaluations.
Riemannian Extension: Extension of diffusion-based samplers to geometries beyond Euclidean space, including a closed-form formula for the conditional drift on constant curvature manifolds.
Empirical Performance: Demonstration of state-of-the-art performance on synthetic energy benchmarks and molecular conformer generation, with a reported reduction in simulation training costs by 4–8 times compared to prior methods.

Experimental Results

The authors evaluate Flow Sampling on four distinct benchmarks:

Synthetic Energy Benchmarks: On Double-Well (DW-4) and Lennard-Jones (LJ-13, LJ-55) potentials, Flow Sampling achieves lower geometric and energy Wasserstein-2 ( $W_2$ ) distances compared to state-of-the-art baselines like iDEM, PIS, DDS, and Adjoint Sampling (AS).
Peptide Sampling: On alanine dipeptide (Ala2) and tetrapeptide (Ala4), the method accurately recovers ground truth densities (measured by KL divergence and Jensen-Shannon divergence) and demonstrates robustness when reducing the number of function evaluations (NFE) during training, outperforming ASBS which degrades significantly at lower NFE.
Amortized Conformer Generation: On large-scale molecular datasets (SPICE and GEOM-DRUGS), Flow Sampling achieves high coverage and precision in generating molecular conformers, outperforming RDKit and ASBS baselines while requiring fewer training steps.
Riemannian Sampling: On the 2D unit sphere ( $S^2$ ) with a mixture of von Mises–Fisher distributions, the method successfully learns to sample from the multimodal target, validating the theoretical extension to curved geometries.

Significance and Claims

The paper claims that Flow Sampling offers a flexible and easy-to-implement alternative to complex stochastic optimal control or Schrödinger bridge formulations. By avoiding optimal couplings and auxiliary corrector networks, it simplifies the training scheme while maintaining high efficiency.

The authors highlight that the method's ability to reduce the number of energy function evaluations (via the replay buffer and closed-form drift) makes it particularly scalable for large models and datasets where energy evaluation is the primary bottleneck. Furthermore, the extension to Riemannian manifolds provides a unified approach for sampling in diverse geometric settings, such as robotics and molecular dynamics, without the need for specialized reparameterizations. The work positions Flow Sampling as a significant step toward efficient, amortized sampling for unnormalized densities in both Euclidean and non-Euclidean spaces.