Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Denis Blessing, Lorenz Richter, Julius Berner, Egor Malitskiy, Gerhard Neumann

TL;DR

This work develops a new method that enables sampling at unprecedented scales while preserving mode diversity, achieving state-of-the-art results on complex synthetic densities and high-dimensional molecular benchmarks, called Bridge Matching Sampler (BMS).

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.

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

TL;DR

This work develops a new method that enables sampling at unprecedented scales while preserving mode diversity, achieving state-of-the-art results on complex synthetic densities and high-dimensional molecular benchmarks, called Bridge Matching Sampler (BMS).

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.
Paper Structure (60 sections, 26 theorems, 197 equations, 15 figures, 5 tables, 2 algorithms)

This paper contains 60 sections, 26 theorems, 197 equations, 15 figures, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Generalized Fixed-Point Diffusion Matching
  3. Notation.
  4. Overview.
  5. General couplings
  6. Schrödinger half bridges
  7. Schrödinger bridges
  8. Independent couplings
  9. Damped fixed-point diffusion matching
  10. Related work
  11. Diffusion-based sampling.
  12. Diffusion bridges.
  13. Target score estimation.
  14. Numerical Experiments
  15. Gaussian mixture models
  16. ...and 45 more sections

Key Result

Lemma 1

Suppose $X$ solves the path-dependent SDE eq: optimal non-Markov SDE where $\xi(X, t)$ is a functional of the path $X$ and $\Pi^*$ is the associated non-Markovian path measure. The drift $u^*$ of the corresponding Markovian process is given by the conditional expectation Furthermore, it can be shown that

Figures (15)

  • Figure 1: ASBS
  • Figure 2: BMS (ours)
  • Figure 4: Samples obtained with BMS, projected on the first two dimensions of a GMM with $d=100$ and 100 modes.
  • Figure 5: Ramachandran plots with $10^6$ samples for Alanine dipeptide for Molecular Dynamics (MD) data, ASBS liu2025adjoint, and our method, BMS, for damping values $\eta\in \{0,10\}$. ASBS and BMS are trained on all-atom coordinates solely using energy evaluations.
  • Figure 6: Jensen-Shannon divergence for the joint distribution of backbone dihedral angles $(\phi,\psi)$-$D_{\mathrm{JS}}$ for Alanine Dipeptide over the course of training for ASBS liu2025adjoint and our proposed method, BMS, for damping values $\eta\in \{0,10\}$.
  • ...and 10 more figures

Theorems & Definitions (51)

  • Definition 1: Reciprocal class and reciprocal projection
  • Definition 2: Markovian projection
  • Lemma 1: Formula for Markov control
  • Proposition 1: Markovian projections
  • Proposition 2: Generalized target score identity
  • Proposition 3: Path-dependent drift of general target measure
  • Proposition 4: Path-dependent drift for Schrödinger half bridges
  • Proposition 5: Schrödinger system
  • Proposition 6: Path-dependent drift for Schrödinger bridges
  • Proposition 7: Target score and drift for independent couplings
  • ...and 41 more