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

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

cs.LGFeb 28, 2026

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

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

Organizations: Karlsruhe Institute of Technology · Zuse Institute Berlin · dida · NVIDIA

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.

Figures & tables

Panel of Figure 1 — Figure 1: Bridge Matching Sampler (BMS). Illustration of one outer iteration of our fixed-point algorithm. Starting from the current drift $u_{i}$ , we (1) simulate the controlled forward SDE, $\mathrm{d}X_{t}=\sigma(t)\,u_{i}(X_{t},t)\,\mathrm{d}t+\sigma(t)\,\mathrm{d}B_{t}$ , starting from $p_{\mathrm{prior}}$ , to obtain $K$ trajectories with endpoints $(X_{0},X_{T})$ (grey curves); (2) resample the prior endpoints, replacing $X_{0}$ by an independent sample $X_{0}^{\prime}\sim p_{\mathrm{prior}}$ to form the coupling $\Pi^{i}_{0,T}=p_{\mathrm{prior}}\otimes\mathbbm{P}^{u_{i}}_{T}$ ; (3) sample intermediate states $X_{t}\sim\mathbbm{P}_{t\mid 0,T}(\cdot\mid X_{0}^{\prime},X_{T})$ from the associated Brownian bridge (dashed purple); and (4) update the drift to $u_{i+1}$ by regressing it on the path-dependent target $\xi(X,t)$ via the matching objective described in ( 15 ). Iterating steps 1-4 drives the time marginals of the learned process towards $p_{\mathrm{target}}$ while maintaining a single, scalable, and stable training objective.

Coupling		$\xi$	Markovian	Algorithm
General	$\Pi^{*}_{0,T}$	Prop. 2.7	✗
Schrödinger (half)	$\Pi^{}_{0,T}=\mathbbm{P}_{0}\otimes\Pi^{}_{T}$	Prop. 2.8	✓	AS ( Havens et al., 2025 )
Schrödinger	$\Pi^{*}_{0,T}=\Pi^{\mathrm{SB}}_{0,T}$	Prop. 2.10	✓	ASBS ( Liu et al., 2025 )
Independent	$\Pi^{}_{0,T}=\Pi^{}_{0}\otimes\Pi^{*}_{T}$	Prop. 2.11	✗	BMS (ours)

Table 1: Overview: Our framework considers target path measures of the form

\Pi^{*}=\Pi^{*}_{0,T}\mathbbm{P}_{|0,T}

with different couplings

\Pi^{*}_{0,T}

. We provide expressions for

\xi

that satisfy

u^{*}(x,t)=\mathbb{E}_{\Pi^{*}}\left[\xi(X,t)|X_{t}=x\right]

and are used for the proposed fixed-point iteration scheme in Algorithm 1 . For an overview of the different expressions for

\xi

see Table 5 in Appendix E .

u_{i+1}=\argmin_{u\in\mathcal{U}}\mathbb{E}_{\Pi^{i}}\left[\int_{0}^{T}\tfrac{1}{2}\|\xi(X,t)-u(X_{t},t)\|^{2}\mathrm{d}t\right]

Algorithm 1 Generalized Fixed-Point Diffusion Matching

Panel of Figure 2 — Figure 2: Comparison of mode TVD and Wasserstein-2 ( $W_{2}$ ) values on Gaussian mixture models across varying dimensions $d$ and damping values $\eta$ . Figure 2 shows that ASBS exhibits instability at low damping values (N/A) and performance degradation as dimensionality increases. Figure 2 demonstrates that our method maintains stable training and consistently avoids mode collapse up to $d=2500$ , significantly scaling beyond the $d=50$ range common in the recent literature. All values are averaged over three random seeds.

Figure 4: Ramachandran plots with $10^{6}$ samples for Alanine dipeptide for Molecular Dynamics (MD) data, ASBS ( Liu et al., 2025 ) , and our method, BMS, for damping values $\eta\in\{0,10\}$ . ASBS and BMS are trained on all-atom coordinates solely using energy evaluations.

Panel of Table 3 — Table 3: Comparison of performance metrics on Alanine Dipeptide (ALA2) and Tetrapeptide (ALA4). We report the Jensen-Shannon divergence of the joint backbone dihedral angle distribution $(\phi,\psi)$ ( $D_{\mathrm{JS}}$ ), the root-mean-squared error between the corresponding free-energy surfaces ( $\mathrm{RMSE}_{\Delta F}$ ), and the Wasserstein-2 distance between the first two TIC components (TICA- $W_{2}$ ). Runs marked with $\dagger$ diverged during training; for these, results from the final checkpoint are shown. All metrics are averaged over three random seeds. The best-performing method is highlighted in bold.

Figure 5: 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 ( Liu et al., 2025 ) and our proposed method, BMS, for damping values

\eta\in\{0,10\}

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Flexible Prior	No Importance Weights	No Multi-stage Optimization
iEFM ( Woo and Ahn, 2024 )	✓	✗	✓
BNEM ( OuYang et al., 2024 )	✗	✗	✓
DiKL ( He et al., 2024 )	✗	✓	✗
iDEM ( Akhound-Sadegh et al., 2024 )	✗	✗	✓
AS ( Havens et al., 2025 )	✗	✓	✓
IWSM ( Wang et al., 2025 )	✗	✗	✓

Appendix

Table 4: Characterization of diffusion-based sampling methods, based on their ability to use arbitrary prior distributions

p_{\mathrm{prior}}

, whether they require importance weights or multiple stages during optimization.

Coupling		Formula for $\sigma(t)^{-1}\xi(X)$
Bridge dep.	$\Pi^{*}_{0,T}$	$\frac{1-c(t)}{1-\gamma(t)}\nabla_{X_{0}}\log\Pi^{\ast}_{0,T}(X_{0},X_{T})+\frac{c(t)}{\gamma(t)}\,\nabla_{X_{T}}\log\Pi^{\ast}_{0,T}(X_{0},X_{T})-\nabla_{X_{t}}\log\mathbbm{P}_{t\|0}(X_{t}\|X_{0})$
SHB	$\Pi^{}_{0,T}=\mathbbm{P}_{0}\otimes\Pi^{}_{T}$	$\frac{\gamma(t)-c(t)}{1-\gamma(t)}\nabla_{X_{0}}\log\mathbbm{P}_{0}(X_{0})+\frac{c(t)}{\gamma(t)}\nabla_{X_{T}}\log\frac{\Pi_{T}^{*}(X_{T})}{\mathbbm{P}_{T}(X_{T})}+\frac{\gamma(t)-c(t)}{\gamma(t)(1-\gamma(t))}\nabla_{X_{0}}\log\mathbbm{P}_{0\|T}(X_{0}\|X_{T})$
SB	$\Pi^{*}_{0,T}=\Pi^{\mathrm{SB}}_{0,T}$	$\frac{\gamma(t)-c(t)}{1-\gamma(t)}\nabla_{X_{0}}\log\frac{\Pi_{0}^{}(X_{0})}{{\varphi}_{0}(X_{0})}+\frac{c(t)}{\gamma(t)}\nabla_{X_{T}}\log\frac{\Pi_{T}^{}(X_{T})}{{\widehat{\varphi}}_{T}(X_{T})}+\frac{\gamma(t)-c(t)}{\gamma(t)(1-\gamma(t))}\nabla_{X_{0}}\log\mathbbm{P}_{T\|0}(X_{T}\|X_{0})$
Bridge indep.	$\Pi^{}_{0,T}=\Pi^{}_{0}\otimes\Pi^{*}_{T}$	$\frac{1-c(t)}{1-\gamma(t)}\nabla_{X_{0}}\log\Pi^{\ast}_{0}(X_{0})+\frac{c(t)}{\gamma(t)}\,\nabla_{X_{T}}\log\Pi^{\ast}_{T}(X_{T})-\nabla_{X_{t}}\log\mathbbm{P}_{t\|0}(X_{t}\|X_{0})$

Appendix

Table 5: Overview: our framework considers target path measures of the form

\Pi^{*}=\Pi^{*}_{0,T}\mathbbm{P}_{|0,T}

with different couplings

\Pi^{*}_{0,T}

. We provide expressions for

\xi

that satisfy

u^{*}(x,t)=\sigma(t)\mathbb{E}_{\Pi^{*}}\left[\xi(X,t)|X_{t}=x\right]

and are used for the proposed fixed-point iteration scheme in Algorithm 1 .

Panel of Figure 6 — Figure 6: Comparison of sliced TVD values on Gaussian mixture models across varying dimensions $d$ and damping values $\eta$ . Figure 6(a) shows that ASBS exhibits instability at low damping values (N/A) and performance degradation as dimensionality increases. Figure 6(b) demonstrates that our method maintains stable training and consistently avoids mode collapse up to $d=2500$ , significantly scaling beyond the $d=50$ range common in recent literature. All values are averaged over three random seeds.

Panel of Figure 7 — Figure 7: Comparison of mode TVD , sliced TVD and Wasserstein-2 $W_{2}$ on a Gaussian mixture model target in $d=100$ with different number of modes. All values are averaged over three random seeds.

Panel of Figure 8 — Figure 8: Visualization of interatomic distances and energy histograms of different $n$ -body identical particle systems. The blue and red curves are generated from our BMS method, and compared to the orange and green curves obtained from a long run MCMC.

Panel of Figure 9 — Figure 9: Comparison of mode TVD and sliced TVD on Gaussian mixture models across varying damping values $\eta$ for $d=250$ and $d=2500$ visualized over training iterations. All values are averaged over three random seeds.

Figure 10: TICA plots with $10^{6}$ samples for Alanine Dipeptide for molecular dynamics (MD) data, ASBS ( Liu et al., 2025 ) , and our proposed method, BMS, for damping values $\eta\in\{0,10\}$ .

Panel of Figure 11 — Figure 11: Ramachandran ( Figure 11(a) ) and TICA ( Figure 11(b) ) plots with $10^{6}$ samples for Alanine Tetrapeptide for molecular dynamics (MD) data, ASBS ( Liu et al., 2025 ) , and our proposed method, BMS, for damping values $\eta=50$ .

Panel of Figure 12 — Figure 12: Comparison of mode TVD and Wasserstein-2 ( $W_{2}$ ) values on Gaussian mixture models across varying dimensions $d$ and Metropolis adjusted Langevin algorithm (MALA) step sizes using the same number of target evaluations as the neural sampling methods. All values are averaged over three random seeds. SMC suffers from mode collapse due to resampling which is avoided when using AIS. PT-MCMC works best if the MALA step size is tuned properly.

Diffusion Sampling

Diffusion sampling studies how to turn a trained diffusion model into high-quality samples efficiently and reliably.

See all 265 papers

May 5, 2026

Flow Sampling: Learning to Sample from Unnormalized Densities via Denoising Conditional Processes

Sampling from unnormalized densities is analogous to the generative modeling problem, but the target distribution is defined by a known energy function instead of data samples. Because evaluating the energy function is often costly, a primary challenge is to learn an efficient sampler. We introduce Flow Sampling, a framework built on diffusion models and flow matching for the data-free setting. Our training objective is conditioned on a noise sample and regresses onto a denoising diffusion drift constructed from the energy function. In contrast, diffusion models' objective is conditioned on a data sample and regresses onto a noising diffusion drift. We utilize the interpolant process to minimize the number of energy function evaluations during training, resulting in an efficient and scalable method for sampling unnormalized densities. Furthermore, our formulation naturally extends to Riemannian manifolds, enabling diffusion-based sampling in geometries beyond Euclidean space. We derive a closed-form formula for the conditional drift on constant curvature manifolds, including hyperspheres and hyperbolic spaces. We evaluate Flow Sampling on synthetic energy benchmarks, small peptides, large-scale amortized molecular conformer generation, and distributions supported on the sphere, demonstrating strong empirical performance.

Aaron Havens, Brian Karrer, Neta Shaul

FAIR at Meta · Weizmann Institute of Science

Sampling from Unnormalized Distributions Diffusion Sampling Energy-Based Models Diffusion Models

Ask PDF

Feb 17, 2026

Efficient Generative Modeling beyond Memoryless Diffusion via Adjoint Schrödinger Bridge Matching

Diffusion models often yield highly curved trajectories and noisy score targets due to an uninformative, memoryless forward process that induces independent data-noise coupling. We propose Adjoint Schrödinger Bridge Matching (ASBM), a generative modeling framework that recovers optimal trajectories in high dimensions via two stages. First, we view the Schrödinger Bridge (SB) forward dynamic as a coupling construction problem and learn it through a data-to-energy sampling perspective that transports data to an energy-defined prior. Then, we learn the backward generative dynamic with a simple matching loss supervised by the induced optimal coupling. By operating in a non-memoryless regime, ASBM produces significantly straighter and more efficient sampling paths. Compared to prior works, ASBM scales to high-dimensional data with notably improved stability and efficiency. Extensive experiments on image generation show that ASBM improves fidelity with fewer sampling steps. We further showcase the effectiveness of our optimal trajectory via distillation to a one-step generator.

Jeongwoo Shin, Jinhwan Sul, Joonseok Lee +2

Seoul National University, Seoul, Korea · Georgia Institute of Technology, Atlanta, GA, USA · Sungkyunkwan University, Seoul, Korea

Diffusion Model Sampling Schrödinger Bridges Few-Step Diffusion Sampling Diffusion Models

Ask PDF

May 5, 2026

Conditional Diffusion Sampling

Sampling from unnormalized multimodal distributions with limited density evaluations remains a fundamental challenge in machine learning and natural sciences. Successful approaches construct a bridge between a tractable reference and the target distribution. Parallel Tempering (PT) serves as the gold standard, while recent diffusion-based approaches offer a continuous alternative at the cost of neural training. In this work, we introduce Conditional Diffusion Sampling (CDS), a framework that combines these two paradigms. To this end, we derive Conditional Interpolants, a class of stochastic processes whose transport dynamics are governed by an exact, closed-form stochastic differential equation (SDE), requiring no neural approximation. Although these dynamics require sampling from a non-trivial initialization distribution, we show both theoretically and empirically that the cost of this initialization diminishes for sufficiently short diffusion times. CDS leverages this by a two-stage procedure: (1) PT is used to efficiently sample the initial distribution, and then (2) samples are transported via the transport SDE. This combination couples the robust global exploration of PT with efficient local transport. Experiments suggest that CDS has the potential to achieve a superior trade-off between sample quality and density evaluation cost compared to state-of-the-art samplers.

Francisco M. Castro-Macías, Pablo Morales-Álvarez, Saifuddin Syed +3

University of Granada · Work done while visiting the University of Cambridge. · University of British Columbia +2

Markov Chain Monte Carlo Sampling from Unnormalized Distributions Diffusion Sampling

Ask PDF

Abstract

Diffusion Sampling

Explore similar work

Flow Sampling: Learning to Sample from Unnormalized Densities via Denoising Conditional Processes

Efficient Generative Modeling beyond Memoryless Diffusion via Adjoint Schrödinger Bridge Matching

Conditional Diffusion Sampling