Wasserstein Sampling via Bridge Interpolants

Langevin dynamics provides a classical route to sampling from unnormalized densities by evolving a distribution along the Wasserstein gradient flow of its Kullback–Leibler divergence to the target. We introduce Wasserstein Sampling via Bridge Interpolants (WASABI), a method for amortizing this sampling process using diffusion models. WASABI translates Wasserstein descent steps in distribution space into updates of the diffusion model's drift, learned through regression on samples generated by the model itself. We derive unbiased estimators that exhibit low variance empirically, making these updates trainable without evaluating the current model's density or score. The resulting fixed-point algorithm alternates between regression updates using a frozen teacher and sample generation from the updated model, with a replay buffer reusing generated samples and their expensive energy-gradient evaluations. One alanine dipeptide, WASABI outperforms reported baselines using an estimated $3.3$ times fewer energy evaluations, while remaining competitive on interacting particle systems benchmarks.