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

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.

• The paper introduces the Bridge Matching Sampler (BMS), a method that unifies diffusion-based samplers using a fixed-point iteration framework based on Nelson’s relation.
• It leverages damping and independent coupling to ensure stable training and superior mode coverage on high-dimensional, multimodal distributions.
• Empirical results demonstrate that BMS outperforms prior methods in achieving accurate density fitting and efficient convergence in benchmarks like Gaussian mixtures and Alanine dipeptide.

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

Problem Formulation and Theoretical Framework

The paper "Bridge Matching Sampler: Scalable Sampling via Generalized Fixed-Point Diffusion Matching" (2603.00530) addresses scalable sampling from unnormalized densities via diffusion models. The central goal is to learn a stochastic transport map between arbitrary prior and target distributions, $p_{\text{prior}}$ and $p_{\text{target}}$, by constructing a Markovian diffusion process whose terminal marginal coincides with the target distribution, using the SDE

$$dX_t = \sigma(t) u(X_t,t) dt + \sigma(t) dB_t, \quad X_0 \sim p_{\text{prior}}$$

where $\sigma(t)$ is the noise schedule and $u$ is the learnable control (drift).

Previous approaches, particularly stochastic optimal control (SOC) schemes and diffusion-based samplers (e.g., Schrödinger Bridge, Adjoint Sampling), rely on matching the path distributions of the forward and backward processes, often involving unstable alternating optimization or restrictive choices for prior/reference distributions. The core innovation here is the unification of these approaches into a generalized fixed-point iteration rooted in Nelson's relation, accompanied by the introduction of the Bridge Matching Sampler (BMS): a method supporting arbitrary priors and reference measures, using a single, computationally efficient objective.

Generalized Fixed-Point Diffusion Matching and Markovianization

The methodology generalizes matching-based samplers by constructing non-Markovian bridge processes through reciprocal projection, followed by Markovianization via minimizing expected squared deviation between the bridge drift and a Markovian function. The fixed-point iteration alternates between:

• Simulation under the current drift,
• Coupling via bridge construction,
• Least-squares regression to fit the Markovian control to the bridge drift.

The drift to match, $\xi$, depends on the endpoints $(X_0, X_T)$ and the transition dynamics.

A principal theoretical result is the generalized target score identity (TSI), which allows for unbiased estimation of $\nabla \log \Pi^*_t$ (the marginal score) using the structure of the bridge coupling. This provides tractable expressions for $\xi$ across a variety of coupling regimes including Schrödinger bridges, half-bridges, and the independent coupling used by BMS.

Tractable Algorithms: AS, ASBS, and BMS

The framework subsumes previous algorithms:

• Adjoint Sampler (AS): Restricts to Dirac priors and memoryless reference process, resulting in instability for broader priors or high diffusion coefficients.
• Adjoint Schrödinger Bridge Sampler (ASBS): Permits arbitrary priors but requires alternating between drift and terminal potential correction, empirically unstable due to interdependence.
• Bridge Matching Sampler (BMS): Uses independent coupling, enabling arbitrary priors and references with tractable drift expressions and eliminating alternating optimization.

BMS leverages damping in its fixed-point iteration to regularize updates and improve stability, especially in high dimensions.

Mode Coverage and Damping: Empirical Results

The empirical evaluation demonstrates that BMS achieves unmatched scalability and mode coverage on synthetic Gaussian mixture targets, $n$-body molecular systems, and complex peptide benchmarks.

For Gaussian mixtures up to $d=2500$, BMS maintains stable training and avoids mode collapse, outperforming ASBS, which degrades or diverges at higher dimensions and lower damping. BMS's efficiency is validated via both mode TVD (empirical mixture weights vs target) and Wasserstein distances.

Figure 1: Comparison of mode TVD and Wasserstein-2 ($W_2$) values on Gaussian mixture models across varying dimensions $d$ and damping $\eta$, showing superior stability and mode coverage for BMS.

In molecular systems such as Alanine dipeptide, BMS-trained samples directly on Cartesian coordinates (not internal or collective variables), achieving superior alignment to joint backbone dihedral distributions (Ramachandran plots) and free-energy surfaces measured by Jensen-Shannon divergence, RMSE, and Wasserstein metrics.

Figure 2: Ramachandran plots with $10^6$ samples for Alanine dipeptide, visualizing mode covering and structural accuracy for BMS and ASBS.

Figure 3: Jensen-Shannon divergence for backbone dihedral angles $(\phi, \psi)$-$D_{\mathrm{JS}}$ for Alanine dipeptide, tracking convergence and stability through training under different damping values.

The strong numerical evidence underscores the criticality of damping: omitting it leads to instability or mode collapse in both ASBS and BMS, while inclusion enables both algorithms to simultaneously achieve mass-covering and accurate density fitting, even at challenging scales.

Practical and Theoretical Implications

Practically, BMS offers a paradigm shift for generative sampling from complex, high-dimensional, and multimodal distributions, as required in statistical physics, molecular modeling, and Bayesian inference. It enables scalable, stable diffusion-based sampling beyond domains previously tractable via matching-based samplers.

Theoretically, the generalized framework unifies Markovian and reciprocal projections, subsuming SOC-based samplers and bridge-based approaches. The fixed-point formulation is inherently linked to forward KL minimization, yielding mass-covering distributions with superior mode diversity—contradicting the mode-seeking bias in conventional reverse KL approaches.

Additional innovations such as control variate tuning for variance reduction and adaptive damping schedules are proposed as future research directions. Convergence guarantees and extension to discrete spaces or normalizing constant estimation remain active open problems.

Speculation on Future AI Developments

The algorithmic stability and scalability of BMS suggest its suitability for deploying diffusion-based samplers in high-dimensional applications such as generative modeling for scientific problems, probabilistic programming, and robust inference. Adaptive damping, importance-weighting, and further theoretical expansion may yield even broader applicability and improved efficiency.

The fixed-point bridge perspective may inform novel approaches to learning in path space, stochastic optimal control, and frontier generative modeling modalities. The mass-covering property, along with computational tractability and stability, will likely be central in next-generation AI systems requiring principled sampling from intractable or highly multimodal targets.

Conclusion

The Bridge Matching Sampler generalizes and subsumes existing matching-based diffusion samplers, enabling scalable, stable sampling from unnormalized densities with arbitrary priors and references. Its fixed-point framework, regularized by damping, achieves state-of-the-art results in mode coverage and density fitting across synthetic and molecular benchmarks. The methodology offers broad practical utility and establishes new theoretical connections between stochastic control, diffusion bridges, and optimal transport, paving the way for future advances in AI sampling and inference.