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

TL;DR

The paper introduces the Bridge Matching Sampler, a scalable and stable diffusion-based sampling method that effectively transports arbitrary priors to complex targets, achieving state-of-the-art results in high-dimensional settings.

Contribution

It generalizes fixed-point diffusion matching methods into a scalable, stable framework with a regularized damped iteration for improved sampling performance.

Findings

Enables sampling at unprecedented scales

Preserves mode diversity effectively

Achieves state-of-the-art results on complex benchmarks

Abstract

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.