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

The Bridge Matching Sampler (BMS) develops a new, scalable, and stable method for sampling from unnormalized densities using diffusion models by generalizing fixed-point iterations rooted in Nelson’s relation. This approach enables learning a stochastic transport map between arbitrary prior and target distributions with a single, scalable, and stable objective.

How it works

The core idea of the Bridge Matching Sampler (BMS) is to learn a stochastic transport map between an arbitrary prior distribution and a target distribution by alternating between the construction of valid bridge processes and their Markovianization to update the control fields The key idea pursued here is to Markovianize (see Definition 2.2) the non-Markovian dynamics in (4) by minimizing suitable loss functionals, yielding a Markovian control u with associated path measure Pu as in (1). This method does not require samples from the target distribution, but only access to its unnormalized density In this sense, our method can be viewed as a generalization of recently proposed sampling algorithms based on leastsquares fixed-point iterations (often referred to as matching algorithms) arising from stochastic optimal control theory.

Generalized Fixed-Point Diffusion Matching

The framework proposes a fixed-point iteration that alternates between the construction of valid bridge processes and their Markovianization to update the control fields Conceptually, each iteration i consists of three steps (see also Algorithm 1 for a summary): 1. Simulation and coupling: We sample X0 ∼ pprior from the prior and simulate the SDE in (1) using the current control ui to obtain XT. 2. Reciprocal projection and target drift: Conditioned on the coupling, we use a reference process P0,T to define a bridge process Πi. Our generalized target score identity in Proposition 2.5 can now be used to derive the drift ξ(X, t) for the process in (4) that corresponds to the target measure Π∗; see Table 1. 3. Markovianization: Finally, we obtain ui+1 by minimizing a least-squares regression objective that fits the Markovian control ui to the target drift ξ evaluated at samples drawn from the bridge process.

Damped Fixed-Point Diffusion Matching

To enhance stability and robustness, a damped variant of the fixed-point iteration is introduced, defined by ui+1 = α Φ(ui) + (1 − α)ui This variational perspective shows that the damped iteration regularizes the current drift to remain close to the previous iterate ui in the L2 sense. As shown in our experiments (Section 4), damping substantially improves performance of both baseline methods and our proposed approach, particularly in high-dimensional settings.

Performance on Complex Benchmarks

Empirically, BMS enables sampling at unprecedented scales while preserving mode diversity, achieving state-of-the-art results on complex synthetic densities and high-dimensional molecular benchmarks In Gaussian mixture models, 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. For n-body identical particle systems, our approach consistently outperforms established baselines across both geometric Wasserstein-2 (W2) and energy value distributions (E(·) W2). In molecular benchmarks, we achieve state-of-the-art performance on Alanine dipeptide and tetrapeptide systems.

Theoretical Foundations

The method is grounded in Nelson’s relation and a generalized target score identity (see Proposition 2.5) The path-dependent drift ξ corresponding to the non-Markovian SDE (10) that induces the target measure Π∗ is given by. We introduce a damped fixed-point iteration ui+1 = α Φ(ui) + (1 − α)ui, where Φ(ui):= EΠi [ξ(X, t)Xt]. The damping parameter η = 1−α/α regularizes the update steps.

Variance Reduction and Control Variates

The matching objective decomposes into two distinct components via Lemma A.5 To minimize its contribution, we can optimize c to minimize the variance contribution in (160) by solving c∗ = arg min c EΠ∗ "Z T 0 1/2 u(Xt, t) − ξc(X, t)2dt/h. This leads to the optimal control variate c∗(t) that minimizes the conditional variance of ξc. We propose parameterizing the control variate as a learnable function cϕ(t) = γ(t) + γ(t)(1 − γ(t))NNϕ(t), where NNϕ is a neural network with learnable parameters ϕ.

Conclusion

In this work, we formulate diffusion-based sampling from unnormalized target densities as a generalized fixed-point iteration grounded in Nelson’s relation and a generalized target score identity (see Proposition 2.5) Building on this foundation, we introduce a new scalable and stable algorithm, termed the Bridge Matching Sampler (BMS), which addresses key limitations of existing approaches – most notably alternating objectives and restrictive choices of reference processes or priors. To further improve robustness and avoid mode collapse in high dimensions, we propose a damped fixed-point iteration. Empirically, BMS achieves state-of-the-art performance on challenging synthetic targets and high-dimensional molecular benchmarks, and scales successfully to highly multimodal distributions in dimensions up to d = 2500. The results reported are without using damping, i.e., η is set to zero.

Future Work

Several promising directions remain for future work. First, the performance of our sampler could be further improved by incorporating an adaptive schedule for the damping parameter η, in the spirit of trust-region methods (Blessing et al., 2025a), or by integrating importance weighting strategies. From a theoretical perspective, establishing convergence guarantees for the proposed fixed-point iteration, as well as conditions ensuring existence and uniqueness of its solution, remains an open problem. Finally, while our focus has been on the principled development of scalable and stable diffusion-based samplers, it would be natural to extend this framework to related settings, such as the estimation of normalizing constants (Guo et al., 2025b) and sampling over discrete state spaces.

Impact Statement

This paper presents work whose goal is to advance the field of Machine Learning There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

Acknowledgements

The authors thank Jaemoo Choi, Christopher von Klitzing, and Henrik Schopmans for the many useful discussions. The authors acknowledge support by funding from a Google PhD fellowship in Machine Learning and ML Foundations and other grants.

References

The references listed are as follows: Akhound-Sadegh et al., 2024, Albergo et al., 2025, Berner et al., 2024, Blessing et al., 2025a, Blessing et al. and Neumann, G. Underdamped diffusion bridges with applications to sampling, Brunick & Shreve, 2013, Chen et al., 2025, Conn et al., 2000, De Bortoli et al., 2024 and Guo et al., 2025a, Guo et al. and Chen, Y. Path Integral Sampler: a stochastic control approach for sampling, Havens et al., 2025, He et al., 2024, Henin et al., 2024 and Holderrieth, P. Leaps: A discrete neural sampler via locally equivariant networks Karras et al., 2025a, Kullback, 1968, Leonard, C. A survey of the schr” odinger problem and some of its connections with optimal transport, Liu et al., 2023, Liu et al. and Chen, Y. Discrete adjoint matching, Kingma, D. P. Adam: A method for stochastic optimization Ko & Geffner, 2025, Kohler & Geffner, 2025, Kullback, S. Probability densities with given marginals, Leonard, C. A survey of the schr” odinger problem and some of its connections with optimal transport, Liao et al., 2023, Liu et al., 2023 and Chen, R. T. Generalized schr” odinger bridge matching, Liu et al., and Chen, Y. Adjoint schrodinger bridge sampler, Liu et al.

1. When sampling from unnormalized densities, Bridge Matching Sampler (BMS) enables learning a stochastic transport map between arbitrary prior and target distributions with a single, scalable, and stable objective. This allows AI systems to sample from complex targets like high-dimensional Gaussian mixture models (up to d=2500) and molecular benchmarks with state-of-the-art results on complex synthetic densities and high-dimensional molecular benchmarks.

2. The damped fixed-point iteration in BMS, defined by ui+1 = α Φ(ui) + (1 − α)ui, provides a mechanism to mitigate mode collapse, as demonstrated by the result that our method maintains stable training and consistently avoids mode collapse up to d = 2500.

3. The use of the learned control variate, parameterized as cϕ(t) = γ(t) + γ(t)(1 − γ(t))NNϕ(t), allows the system to dynamically discover the optimal interpolation schedule during optimization, leading to a more robust and accurate sampling process.

4. The formulation of importance sampling via the Probability Flow ODE, where log P u T (XT) = log ppprior(X0) − Z T 0 ∇ · f(Xt, t) dt, allows for the estimation of expectations under the target distribution without needing to know the intractable normalizing constant Z.

5. The joint learning of the drift network u and control variate c via c∗ = arg min u,c EΠ∗ 'Z T 0 1/2 u(Xt, t) − ξ c (X, t)Xt = x enables the model to optimize both the sampling trajectory and its variance reduction schedule simultaneously.

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.