Bridge Matching Sampler: Scalable Sampling via Generalized Fixed-Point Diffusion Matching
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| n-body particle system sampling | DW-4 d = 8 | W2 Distance0.38 | 20 | |
| n-body particle system sampling | LJ-13 (d = 39) | W2 Distance1.54 | 13 | |
| n-body particle system sampling | LJ-55 d = 165 | W23.8 | 10 | |
| Molecular conformation sampling | Alanine Dipeptide (ALA2) d = 66 | DJS (phi, psi)0.068 | 4 | |
| Molecular conformation sampling | Tetrapeptide ALA4 (d = 126) | (phi, psi)-DJS0.228 | 4 |