WASABI: Wasserstein Sampling via Bridge Interpolants
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.
Related
- DAWIS: Data Assimilation with Windowed Inverse Sampling via Multitask Interpolants
- Neural Sampling with Reweighted Normalizing Flows via the Wasserstein--Fisher--Rao JKO Scheme
- Wasserstein Bounds for Euler and Randomized Midpoint Langevin Sampling with Noisy Gradients
- Transport along Local Annealing Paths: Wasserstein Geometry and Multimodal Sampling
- Wasserstein Gradient Flows and Forward-Only Diffusion Are Not Enough for Multimodal Sampling
BibTeX
@misc{salazar2026wasabi,
title = {{WASABI: Wasserstein Sampling via Bridge Interpolants}},
author = {Sebastian Salazar and Dilum Fernando and Michal Kucer and Emily Casleton and Luhuan Wu and David M. Blei},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.41958},
eprint = {2610.41958},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.41958},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - WASABI: Wasserstein Sampling via Bridge Interpolants
AU - Sebastian Salazar
AU - Dilum Fernando
AU - Michal Kucer
AU - Emily Casleton
AU - Luhuan Wu
AU - David M. Blei
PY - 2026
DA - 2026/10/30
PB - arXiv
DO - 10.48550/arxiv.2610.41958
UR - https://arxiv.org/abs/2610.41958
AB - 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.
ER -