WASABI: Wasserstein Sampling via Bridge Interpolants

Sebastian Salazar; Dilum Fernando; Michal Kucer; Emily Casleton; Luhuan Wu; David M. Blei

stat.MLOct 30, 2026

WASABI: Wasserstein Sampling via Bridge Interpolants

Authors: Sebastian Salazar, Dilum Fernando, Michal Kucer, Emily Casleton, Luhuan Wu, David M. Blei

Organizations: Columbia University · Los Alamos National Laboratory · Johns Hopkins University

Abstract

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

Panel of Figure 1 — Figure 1: Overview of WASABI . The endpoint distributions evolve from an initialization $q_{0}$ , distinct from the shared diffusion prior $p_{0}$ . Top: we optimize a sequence of sampling distributions $q_{n}$ by evolve them toward the target $\pi$ through Wasserstein gradient descent of the KL divergence. Bottom: the sampling distribution $q_{n}$ is parameterized by a diffusion model with a prior $p_{0}$ for all $n$ ; WASABI translates Wasserstein descent into updates of the diffusion model by updating its denoising representation $\bar{m}_{t}^{(n)}$ , gradually moving it toward the target across iterations $n$ . The final panel shows the ideal diffusion sampler, with denoising representation $\bar{m}_{t}^{\star}$ associated with $\pi$ .

Figure 2: DW-4. (Left) Pooled pairwise distances

\{r_{ij}\}_{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.

W_{2}

compares particle configurations (Config.) and unfiltered potential energies (Energy).

Figure 3: LJ-13. (Left) Energy and pooled pair-distance

\{r_{ij}\}_{i<j}

densities: WASABI (green, three-seed mean) and reference (dashed). (Right) Configuration

W_{2}

compares particle arrangements after matching and alignment; energy

W_{2}

compares energy distributions. Results are mean

\pm

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 $\pm$ 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 $D_{\rm JS}$ in nats; — denotes unreported values.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Term	Energy expression	Physics
$U_{\mathrm{bond}}$	$\displaystyle\frac{1}{2}\sum_{e\in E}k_{e}^{\mathrm{b}}(\ell_{e}-\ell_{e}^{0})^{2}$	Covalent bond stretching.
$U_{\mathrm{angle}}$	$\displaystyle\frac{1}{2}\sum_{a\in A}k_{a}^{\mathrm{a}}(\alpha_{a}-\alpha_{a}^{0})^{2}$	Valence-angle bending.
$U_{\mathrm{torsion}}$	$\displaystyle\sum_{\tau\in T}\sum_{n}V_{\tau n}[1+\cos(n\phi_{\tau}-\gamma_{\tau n})]$	Rotation barriers about bonds.
$U_{\mathrm{LJ}}$	$\displaystyle 4\sum_{i<j}^{\star}\epsilon_{ij}(\rho_{ij}^{12}-\rho_{ij}^{6})$	Steric repulsion and dispersion.
$U_{\mathrm{Coulomb}}$	$\displaystyle k_{e}\sum_{i<j}^{\star}\frac{q_{i}q_{j}}{r_{ij}}$	Direct charge interactions.
$U_{\mathrm{GB}}$	$\displaystyle-\frac{k_{e}\Delta_{\epsilon}}{2}\sum_{i,j}\frac{q_{i}q_{j}}{f_{ij}^{\mathrm{GB}}}$	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

\pm

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_{\eta}=1.00$ fixed		$\lambda_{\rm cov}=0.25$ fixed
$\lambda_{\rm cov}$	$D_{\rm JS}\downarrow$	$s_{\eta}$	$D_{\rm JS}\downarrow$
$0.0000$	$0.0595\pm 0.0212$	$0.10$	0.0312 $\pm$ 0.0023
$0.0667$	$0.0520\pm 0.0125$	$0.16$	$0.0387\pm 0.0218$
$0.1333$	$0.0479\pm 0.0137$	$0.22$	$0.0408\pm 0.0173$
$0.2000$	$0.0377\pm 0.0061$	$0.28$	$0.0512\pm 0.0167$

Appendix

Table 5: Covariance-weight and step-scale sweeps. Values are mean

\pm

SD over three seeds; bold marks the lowest mean in each sweep.

Boltzmann Sampling

See all 18 papers

One-Step Generative Modeling via Wasserstein Gradient Flows
May 12, 2026Jiaqi Han, Puheng Li, Qiushan Guo +3Wasserstein Gradient Flows Image Generation
Stanford University · ByteDance
Bridge Matching Sampler: Scalable Sampling via Generalized Fixed-Point Diffusion Matching
Feb 28, 2026Denis Blessing, Lorenz Richter, Julius Berner +2Diffusion Sampling
Karlsruhe Institute of Technology · Zuse Institute Berlin · dida +1
Flow Sampling: Learning to Sample from Unnormalized Densities via Denoising Conditional Processes
May 5, 2026Aaron Havens, Brian Karrer, Neta ShaulSampling from Unnormalized Distributions Diffusion Sampling
FAIR at Meta · Weizmann Institute of Science

May 12, 2026J/K move · Enter open · S save

One-Step Generative Modeling via Wasserstein Gradient Flows

Jiaqi Han, Puheng Li, Qiushan Guo, Renyuan Xu+2

Stanford University · ByteDance

Wasserstein Gradient Flows Image Generation One-Step Generative Modeling Generative Modeling Flow-Based Generative Modeling

Open paper AskPDF arXiv

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

\times

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

\times

faster sampling. These results show that Wasserstein gradient flows provide a principled and effective foundation for fast and high-fidelity generative modeling.

Abstract

Boltzmann Sampling

Explore similar work

One-Step Generative Modeling via Wasserstein Gradient Flows

Bridge Matching Sampler: Scalable Sampling via Generalized Fixed-Point Diffusion Matching

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