Sampling from Boltzmann distributions is a fundamental problem in physics, chemistry, and biology. We introduce WAsserstein SAmpling via Bridge Interpolants (WASABI), a method that trains a diffusion model to generate samples from a target Boltzmann distribution given access only to the energy function. WASABI minimizes the KL divergence between the distribution generated by its diffusion model and the Boltzmann target using Wasserstein gradient flows. This approach builds on a principle underlying classical sampling methods such as Langevin dynamics. However, applying Wasserstein gradient flows to diffusion models is challenging because the updates require the gradient of the model's own log density, which is generally intractable. We derive a practical training algorithm that avoids computing this gradient. Once trained, WASABI can generate new samples without further evaluations of the energy or its gradient. On alanine dipeptide, WASABI outperforms the best reported baselines using an estimated 3.3 times fewer energy evaluations. It also achieves competitive performance on interacting particle system benchmarks.
Figures & tables
Figure 1: Overview of WASABI . The endpoint distributions evolve from an initialization q0 , distinct from the shared diffusion prior p0 . Top: we optimize a sequence of sampling distributions qn by evolve them toward the target π through Wasserstein gradient descent of the KL divergence. Bottom: the sampling distribution qn is parameterized by a diffusion model with a prior p0 for all n ; WASABI translates Wasserstein descent into updates of the diffusion model by updating its denoising representation mˉt(n) , gradually moving it toward the target across iterations n . The final panel shows the ideal diffusion sampler, with denoising representation mˉt⋆ associated with π .
Figure 2: DW-4. (Left) Pooled pairwise distances {rij}i<j from 10,000 reference and 10,000 WASABI samples; dashed lines mark the pair-potential minima. (Right) Matched-training comparison using samples pooled over rounds 625–1,000. W2 compares particle configurations (Config.) and unfiltered potential energies (Energy).
Figure 3: LJ-13. (Left) Energy and pooled pair-distance {rij}i<j densities: WASABI (green, three-seed mean) and reference (dashed). (Right) Configuration W2 compares particle arrangements after matching and alignment; energy W2 compares energy distributions. Results are mean ± SD across three seeds; baselines are quoted from Blessing et al. (2026, Table 2) .
Figure 4: ALA2 at 300 K. (Left) Ramachandran plots comparing 20k WASABI and 20k REMD samples; gray bins are unobserved. (Right) Jensen–Shannon divergence and estimated computational cost. WASABI results are mean ± SD across three seeds; baselines are quoted from Blessing et al. (2026) . Training cost is in TFLOPs per outer iteration, energy/force calls in millions, and DJS in nats; — denotes unreported values.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Term
Energy expression
Physics
Ubond
21e∈E∑keb(ℓe−ℓe0)2
Covalent bond stretching.
Uangle
21a∈A∑kaa(αa−αa0)2
Valence-angle bending.
Utorsion
τ∈T∑n∑Vτn[1+cos(nϕτ−γτn)]
Rotation barriers about bonds.
ULJ
4i<j∑⋆ϵij(ρij12−ρij6)
Steric repulsion and dispersion.
UCoulomb
kei<j∑⋆rijqiqj
Direct charge interactions.
UGB
−2keΔϵi,j∑fijGBqiqj
Implicit-solvent electrostatics.
Appendix
Table 4: Energy terms in the alanine-dipeptide Hamiltonian. The formulas display functional forms; ff96 exclusions and special 1–4 parameters are included in the starred nonbonded sums. Coefficients include any thermal rescaling used to define U .
Figure 5: ALA2 ablations and replay comparison. (a,c) JS mean ± SD over three seeds; gray dots show individual scores and dashed lines published BMS/ASBS means ( Blessing et al., 2026 ) . (b,d) Matched seeds with smaller and larger replay settings; dashed curves show REMD. Arrows highlight mitigation of seed 2’s central-peak underrepresentation. Both capacity and refresh increase. The sweeps and replay comparison use different checkpoints and sample counts; evaluation protocols are described below.
Covariance sweep
Step-scale sweep
sη=1.00 fixed
λcov=0.25 fixed
λcov
DJS↓
sη
DJS↓
0.0000
0.0595±0.0212
0.10
0.0312 ± 0.0023
0.0667
0.0520±0.0125
0.16
0.0387±0.0218
0.1333
0.0479±0.0137
0.22
0.0408±0.0173
0.2000
0.0377±0.0061
0.28
0.0512±0.0167
Appendix
Table 5: Covariance-weight and step-scale sweeps. Values are mean ± SD over three seeds; bold marks the lowest mean in each sweep.
Diffusion models and flow-based methods have shown impressive generative capability, especially for images, but their sampling is expensive because it requires many iterative updates. We introduce W-Flow, a framework for training a generator that transforms samples from a simple reference distribution into samples from a target data distribution in a single step. This is achieved in two steps: we first define an evolution from the reference distribution to the target distribution through a Wasserstein gradient flow that minimizes an energy functional; second, we train a static neural generator to compress this evolution into one-step generation. We instantiate the energy functional with the Sinkhorn divergence, which yields an efficient optimal-transport-based update rule that captures global distributional discrepancy and improves coverage of the target distribution. We further prove that the finite-sample training dynamics converge to the continuous-time distributional dynamics under suitable assumptions. Empirically, W-Flow sets a new state of the art for one-step ImageNet 256×256 generation, achieving 1.29 FID, with improved mode coverage and domain transfer. Compared to multi-step diffusion models with similar FID scores, our method yields approximately 100× faster sampling. These results show that Wasserstein gradient flows provide a principled and effective foundation for fast and high-fidelity generative modeling.
Sampling from unnormalized densities using diffusion models has emerged as a powerful paradigm. However, while recent approaches that use least-squares `matching' objectives have improved scalability, they often necessitate significant trade-offs, such as restricting prior distributions or relying on unstable optimization schemes. By generalizing these methods as special forms of fixed-point iterations rooted in Nelson's relation, we develop a new method that addresses these limitations, called Bridge Matching Sampler (BMS). Our approach enables learning a stochastic transport map between arbitrary prior and target distributions with a single, scalable, and stable objective. Furthermore, we introduce a damped variant of this iteration that incorporates a regularization term to mitigate mode collapse and further stabilize training. Empirically, we demonstrate that our method enables sampling at unprecedented scales while preserving mode diversity, achieving state-of-the-art results on complex synthetic densities and high-dimensional molecular benchmarks.
Denis Blessing, Lorenz Richter, Julius Berner +2
Karlsruhe Institute of Technology · Zuse Institute Berlin · dida +1
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.