Flow Sampling: Amortized Inference for Unnormalized Densities Through Denoising Objectives

Introduction

Sampling from unnormalized probability densities represents a foundational challenge across computational disciplines, from statistical physics to molecular dynamics and Bayesian inference. Unlike the standard generative modeling paradigm where one learns from finite datasets, this setting provides access to an energy function or unnormalized log density r(x), often accompanied by its gradient, but demands efficient sampling from the normalized distribution q(x) = exp(r(x))/Z. Traditional Markov chain Monte Carlo methods, while asymptotically exact, suffer from slow mixing and inherent sequentiality, rendering them impractical for applications requiring rapid batched inference such as materials design or high throughput molecular screening.

In their recent work, Flow Sampling: Learning to Sample from Unnormalized Densities via Denoising Conditional Processes, Havens, Karrer, and Shaul introduce a compelling alternative that reframes amortized sampling through the lens of flow matching and diffusion models. Rather than simulating long Markov chains for each new target distribution, Flow Sampling learns a parametric transport map that can generate independent samples in parallel, achieving state of the art results on molecular conformer generation while reducing computational overhead by factors of four to eight compared to prior diffusion based approaches.

The Mechanism: Denoising Drifts and Energy Gradients

Standard flow matching and diffusion models operate by conditioning on a data sample x1 drawn from the target distribution, then learning to reverse a noising process that corrupts this data toward a simple source distribution p0. The training objective regresses a velocity field onto conditional velocities that depend on the known data point, as shown in equation (5) of the paper.

Flow Sampling inverts this logic in a subtle but crucial manner. Since samples from the target q are unavailable by definition, the method conditions instead on a noise sample x0 from the source distribution p0. The training objective now regresses onto a denoising diffusion drift constructed directly from the energy function r(x). Specifically, the interpolant process Xt = σ_t X0 + α_t X1 is maintained, but rather than observing X1 from a dataset, the method utilizes the energy gradient ∇r evaluated at the interpolant endpoints to construct the regression target.

This construction relies on the insight that for the marginal distribution of the flow to match the target density q, the conditional drift must incorporate information about how the energy landscape shapes the probability mass. The authors derive a closed form expression for this drift that reuses gradients of the negative energy along the interpolant path. By maintaining a detached state of the model during training, they create a replay buffer of energy evaluations that significantly reduces the number of expensive forward passes through the energy function. This stands in sharp contrast to existing diffusion samplers based on stochastic optimal control or Schrödinger bridges, which require complex auxiliary networks or repeated Monte Carlo corrections involving fresh energy evaluations at each training step.

Geometric Extension to Riemannian Manifolds

Perhaps the most technically sophisticated contribution lies in the extension of Flow Sampling to Riemannian manifolds. Many physical systems naturally inhabit non Euclidean geometries; molecular conformers live on tori of torsional angles, while directional data resides on hyperspheres. Prior diffusion samplers were largely confined to flat spaces, limiting their applicability to structured domains.

The authors demonstrate that by replacing the affine interpolant with a geodesic interpolant on the manifold, the framework extends naturally to constant curvature spaces including hyperspheres and hyperbolic manifolds. They derive a closed form formula for the conditional denoising drift on these geometries, enabling the first diffusion based sampler capable of handling unnormalized densities on curved spaces. This is not merely a technical exercise; the ability to sample from mixtures of von Mises Fisher distributions on the sphere, as illustrated in Figure 1 of the paper, opens pathways for applications in protein orientation analysis, geological directional statistics, and hyperbolic representation learning where the underlying geometry constrains the valid probability mass.

Your Take

The efficiency gains reported by Havens et al. suggest a fundamental reassessment of when amortized methods outperform traditional MCMC. The key insight appears to be that Flow Sampling replaces the stochastic exploration of MCMC with a deterministic transport map learned through supervised regression onto energy informed targets. This creates an implicit importance sampling scheme without requiring explicit importance weights or resampling steps that plague prior diffusion samplers.

However, this approach carries distinct limitations that merit honest consideration. The method requires access to gradients of the energy function ∇r(x), which may be unavailable or computationally prohibitive in certain black box simulation environments. Furthermore, the current formulation assumes sufficient smoothness in the energy landscape to permit stable gradient flows; rugged landscapes with sharp metastable states may challenge the interpolant based training dynamics. The Riemannian extension, while elegant, currently applies only to constant curvature manifolds, leaving general Riemannian metrics or discrete spaces as open frontiers.

From a broader perspective, Flow Sampling represents a convergence of two previously distinct threads in machine learning: score based generative models and neural network based variational inference. By treating the energy function as the fundamental object rather than finite