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.
Figures & tables
Figure 1 : Generation trajectory of score matching and ASBM. Top : Backward drift accumulated over time in pixel level. Bottom : Denoising path in image level. ASBM shows significantly smaller transport cost with straighter path, leading to efficient generation.
Figure 2 : Generated samples from ASBM on pixel space (CIFAR-10) and on latent space (FFHQ).
Method
OC
No PT
FB
FID ↓
Refinement method
DOT ( Tanaka, 2019 )
✓
✗
–
15.78
DGFLOW ( Ansari et al., 2021 )
✓
✗
–
9.63
Flow-based method
FM ( Lipman et al., 2023 )
✗
✓
✓
6.35
Rectified Flow ( Liu et al., 2023b )
✓
✓
✗
6.01
Table 1 : FID evaluation on CIFAR-10. ASBM shows superior generative performance by achieving all three key advantages: optimal coupling (OC), no reliance on pre-training (No PT), and consistent forward–backward dynamics (FB).
Figure 3 : FID comparison along the NFE. We use M and NM to denote the memoryless and non-memoryless condition, respectively. BM denotes empirical bridge-matching pretraining.
NFE
25
50
100
250
500
1000
Score SDE
52.08
19.02
9.84
7.79
6.88
6.63
ASBM (Ours)
8.85
7.64
6.85
6.47
6.38
6.27
Table 2 : FID evaluation on latent space of FFHQ.
Figure 4 : Trajectory efficiency. Left : Ours shows significantly straighter trajectory, leading to low NFE at generation. Right : Ours has lower trajectory variance, implying its better organized path.
Method
Score SDE
SB-FBSDE
DSBM
ASBM
FID
6.72
285.77
39.84
3.74
Table 3 : FID at 25 steps with Heun solver on CIFAR-10.
Figure 5 : Localized prior-data coupling. ASBM trajectories preserve information: reversing from a noised image produces samples similar to the original. In contrast, memoryless dynamics yield completely random samples due to highly noisy trajectories.
Figure 6 : Uncurated one-step generation from distillation on CIFAR-10 . Red boxes highlight repeated patterns indicating mode collapse. DMD still suffers from mode collapse despite the costly regression loss, whereas ASBM achieves diverse generation due to its organized, localized coupling.
Method
FID ↓
Recall ↑
Precision ↑
SDS ( Poole et al., 2023 )
9.36
0.504
0.706
DMD ( Yin et al., 2024b )
8.25
0.513
0.715
Ours
6.68
0.542
0.702
Table 4 : Result of distillation to one-step generator on CIFAR-10.
Figure 7 : Initialization of one-step generator on CIFAR-10. Diffusion-based initialization (left) produces noisy images even with timestep shifting, while ASBM (right) yields clear initial estimates due to its straighter and more organized trajectory
Figure 8 : Ablation on different degree of memorylessness.
Forward
Backward NFE
NFE
25
50
100
250
500
1000
10
23.62
8.03
3.58
3.40
3.28
3.05
20
20.83
5.39
3.05
2.91
2.87
2.74
50
20.73
5.87
3.16
3.01
2.94
2.77
Table 5 : Ablation on different forward NFE.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure I : Initialization of one-step generator (Uncurated generation). As discussed in Sec. 3.3 , ASBM shows clear initialization for one-step generator, indicating its straighter and more organized trajectory. On the other hand, score-based initialization gives much more noisy initialization even with timestep shifting ( Yin et al., 2024b ) .
Figure II : Mode collapse in distillation (Uncurated generation). ASBM shows strong mode coverage in distillation task while the score distillation models still suffer from mode collapse even with costly regression loss.
Scalar (mean, var)
∥X1∥22 (mean, var)
Eigenvalues (min, max)
Gaussian reference
(0,1)
(3072,6144)
(0.566,1.557)
ASBM (learned forward)
(−0.001,1.012)
(3090,6322)
(0.572,1.594)
Appendix
Table I : Prior alignment diagnostics on CIFAR-10. Each metric compares N=50K forward-simulated samples against true Gaussian samples of the same size ( D=3072 ).
Recent advances in generative modeling have enabled the efficient computation of Schrödinger bridges (SB) in high-dimensional settings by leveraging partially simulation-free training methods inspired by flow matching. However, these have not covered SBs with reflecting dynamics, a useful model choice with built-in guarantees that generated samples stay in the data domain. Existing alternatives for reflected SBs instead rely on more complex training based on forward--backward SDE theory, requiring expensive higher-order derivatives and sampling entire paths during training. In this article, we introduce a partially simulation-free framework that allows reflected SBs to be trained similarly to flow matching, using a new sampling method and regression target. We demonstrate our results by coupling pairs of well-known high-dimensional image datasets. Using reflected dynamics incurs negligible additional wall-clock time during both training and inference while maintaining or slightly improving generative performance.
Marcus Häggbom, Viktor Nilsson, Pierre Nyquist +1
SEB Group and Department of Mathematics KTH Royal Institute of Technology Stockholm, Sweden · Department of Mathematical Sciences Chalmers University of Technology and University of Gothenburg and Department of Mathematics KTH Royal Institute of Technology
These notes recapitulate the high level mathematical principles behind different techniques for generative modeling. I show the connections between optimal transport and standard techniques such as Schr{ö}dinger bridge and flow matching.
Over the past few years, diffusion-based Schrödinger bridge models have been proposed to approximate optimal transport dynamics between two prescribed boundary distributions, with successful applications to generative modeling. More precisely, these methods aim to estimate a path measure whose initial and terminal marginals match the two boundary distributions, while minimizing the Kullback-Leibler divergence with respect to a reference Markov process. In this work, we consider the generalized Schrödinger bridge problem, in which the reference process is a twisted Brownian motion, that is, a Feynman-Kac transform of a Brownian motion induced by a time-dependent differentiable potential. Building on the Iterative Markovian Fitting (IMF) paradigm, and in particular on its special case Diffusion Schrödinger Bridge Matching (DSBM), which corresponds to the zero potential case, we introduce Twisted Schrödinger Bridge Matching (TSBM), a diffusion-based method designed to handle both continuous- and discrete-time potentials. Unlike previous approaches, TSBM provides a rigorous extension of the IMF scheme to the generalized Schrödinger bridge problem. This derivation leads to a new bridge-matching loss that depends explicitly on the gradient of the potential and recovers the DSBM objective when the potential vanishes, yielding improved performance. We further introduce trajectory-based variance-reduction techniques that substantially stabilize optimization and may be useful beyond the present setting. Finally, we empirically demonstrate the benefits of TSBM for trajectory inference across increasingly high-dimensional settings, including crowd navigation and single-cell data. Code available at https://github.com/maxencenoble/twisted-sb-matching.
Maxence Noble, Marie Scheid, Yazid Janati +2
CMAP, Ecole polytechnique · Institute of Foundation Models · MBZUAI +1