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

Flow Sampling: Learning to Sample from Unnormalized Densities via Denoising Conditional Processes
May 5, 2026Aaron Havens, Brian Karrer, Neta ShaulSampling from Unnormalized Distributions Diffusion Sampling
FAIR at Meta · Weizmann Institute of Science
Efficient Generative Modeling beyond Memoryless Diffusion via Adjoint Schrödinger Bridge Matching
Feb 17, 2026Jeongwoo Shin, Jinhwan Sul, Joonseok Lee +2Diffusion Model Sampling Schrödinger Bridges
Seoul National University, Seoul, Korea · Georgia Institute of Technology, Atlanta, GA, USA · Sungkyunkwan University, Seoul, Korea
Conditional Diffusion Sampling
May 5, 2026Francisco M. Castro-Macías, Pablo Morales-Álvarez, Saifuddin Syed +3Markov Chain Monte Carlo Sampling from Unnormalized Distributions
University of Granada · Work done while visiting the University of Cambridge. · University of British Columbia +2

May 5, 2026J/K move · Enter open · S save

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

Aaron Havens, Brian Karrer, Neta Shaul

FAIR at Meta · Weizmann Institute of Science

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

Open paper AskPDF arXiv

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.

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