Multi-Marginal Schrödinger Bridge Matching

Byoungwoo Park; Juho Lee

stat.MLOct 18, 2025

Multi-Marginal Schrödinger Bridge Matching

Authors: Byoungwoo Park, Juho Lee

Organizations: KAIST

Abstract

Understanding the continuous evolution of populations from discrete temporal snapshots is a critical research challenge, particularly in fields like developmental biology and systems medicine where longitudinal tracking of individual entities is often impossible. Such trajectory inference is vital for unraveling the mechanisms of dynamic processes. While Schrödinger Bridge (SB) offer a potent framework, their traditional application to pairwise time points can be insufficient for systems defined by multiple intermediate snapshots. This paper introduces Multi-Marginal Schrödinger Bridge Matching (MSBM), a novel algorithm specifically designed for the multi-marginal SB problem. MSBM extends iterative Markovian fitting (IMF) to effectively handle multiple marginal constraints. This technique ensures robust enforcement of all intermediate marginals while preserving the continuity of the learned global dynamics across the entire trajectory. Empirical validations on synthetic data and real-world single-cell RNA sequencing datasets demonstrate the competitive or superior performance of MSBM in capturing complex trajectories and respecting intermediate distributions, all with notable computational efficiency.

Figures & tables

Figure 2: Comparison of generated population dynamics using MIOFlow, DMSB and MSBM on a 2-dim petal dataset. All trajectories are generated by simulating the dynamics from

\rho_{t_{0}}

Figure 3: Evaluation results of

{\mathcal{W}}_{2}

and MMD.

	${\mathcal{W}}_{2}$ ↓		Runtime
Methods	$t_{1}$	$t_{3}$	hours
DMSB †	1.10	1.51	15.54
SBIRR †	1.08	$1.33$	0.36 (0.38) ∗
MSBM (Ours)	$1.09$	1.30	0.09

Table 1: Performance on the 5-dim PCA of hESC dataset.

{\mathcal{W}}_{2}

is computed between test

\rho_{t_{i}}

and generated

\hat{\rho}_{t_{i}}

by simulating the dynamics from test

\rho_{t_{0}}

. Best results are highlighted .

	${\mathcal{W}}_{1}\downarrow$
Methods	$t_{1}$	$t_{2}$	$t_{3}$	$t_{4}$	Mean
NLSB † ( Koshizuka and Sato, 2022 )	$0.68$	$0.84$	$0.81$	$0.79$	$0.78$
IPF (GP) † ( Vargas et al., 2021 )	$0.70$	$1.04$	$0.94$	0.98	$0.92$
IPF (NN) † ( Bortoli et al., 2021 )	0.74	0.89	0.84	0.83	0.82
SB-FBSDE † ( Chen et al., 2022a )	0.56	0.80	1.00	1.00	0.84
MSBM (Ours)	0.64	$\mathbf{0.73}$	$\mathbf{0.72}$	$\mathbf{0.73}$	$\mathbf{0.71}$

Table 2: Performance on the 5-dim PCA of EB dataset.

{\mathcal{W}}_{1}

is computed between test

\rho^{\texttt{te}}_{t_{i}}

and generated

\hat{\rho}_{t_{i}}

by simulating the dynamics from previous test

\rho^{\texttt{te}}_{t_{i-1}}

. Best results are highlighted .

Model	CITE		MULTI		Runtime
Model	5dim	100dim	5dim	100dim	hours
OT-CFM †	0.88 $\pm$ 0.06	45.39 $\pm$ 0.42	0.94 $\pm$ 0.05	54.81 $\pm$ 5.86	-
WLF-SB †	0.80 $\pm$ 0.02	46.13 $\pm$ 0.08	0.95 $\pm$ 0.21	55.07 $\pm$ 5.50	0.51
DMSB	0.82 $\pm$ 0.04	50.24 $\pm$ 0.99	1.03 $\pm$ 0.12	58.51 $\pm$ 5.11	2.16
MSBM	0.87 $\pm$ 0.05	45.93 $\pm$ 0.37	0.91 $\pm$ 0.22	54.75 $\pm$ 6.21	0.08

Table 4: Performance on the 5-dim and 100-dim PCA of CITE/MULTI dataset.

{\mathcal{W}}_{1}

is computed between test

\rho^{\texttt{te}}_{t_{i}}

and generated

\hat{\rho}_{t_{i}}

by simulating the dynamics from test

\rho^{\texttt{te}}_{t_{i-1}}

. Best results are highlighted .

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Learning Rate	Iteration (N)	Training step (S)	# of Discretization	Batch Size	$T$
Petal	$1\times 10^{-3}$	20	1000	30	256	4
hESC	$1\times 10^{-3}$	100	1000	30	256	5
EB(100-dim)	$2\times 10^{-4}$	10	1000	100	256	4
EB(5-dim)	$2\times 10^{-4}$	3	50000	100	256	4
CITE(5-dim)	$1\times 10^{-4}$	50	2000	100	256	3
CITE(100-dim)	$1\times 10^{-4}$	50	2000	100	256	3

Appendix

Table 5: Training Hyper-parameters

	MMD ↓				SWD ↓
Methods	Full	${t_{1}}$	${t_{2}}$	${t_{3}}$	Full	${t_{1}}$	${t_{2}}$	${t_{3}}$
NLSB † ( Koshizuka and Sato, 2022 )	$0.66$	$0.38$	$0.37$	$0.37$	$0.54$	$0.55$	$0.54$	$0.55$
MIOFlow † ( Huguet et al., 2022 )	$0.23$	$0.23$	$0.90$	$0.23$	$0.35$	$0.49$	$0.72$	$0.50$
DMSB † ( Chen et al., 2023 )	${\color[rgb]{0,0,1}{0.03}}$	$\mathbf{0.04}$	$\mathbf{0.04}$	$\mathbf{0.04}$	${\color[rgb]{0,0,1}{0.16}}$	${\color[rgb]{0,0,1}{0.20}}$	${\color[rgb]{0,0,1}{0.19}}$	$\mathbf{0.18}$
MSBM	$\mathbf{0.018}\pm$ 4e-4	$\mathbf{0.049}\pm$ 3e-2	$\mathbf{0.038}\pm$ 5e-4	${\color[rgb]{0,0,1}{0.05}}\pm$ 9e-4	$\mathbf{0.129}\pm$ 3e-3	$\mathbf{0.1895}\pm$ 1e-2	$\mathbf{0.1772}\pm$ 2e-3	${\color[rgb]{0,0,1}{0.1997}}\pm$ 4e-3

Appendix

Table 6: Full results over

3

different seeds. Performance on the 100-dim PCA of EB dataset. MMD and SWD are computed between test

\rho^{\texttt{te}}_{t_{i}}

and generated

\hat{\rho}_{t_{i}}

by simulating the dynamics from test

\rho^{\texttt{te}}_{t_{0}}

	${\mathcal{W}}_{2}$ ↓		Runtime
Methods	$t_{1}$	$t_{3}$	hours
TrajectoryNet †	1.30	1.93	10.19
DMSB †	1.10	1.51	15.54
SBIRR †	1.08	1.33	0.36 (0.38) ∗
MSBM (Ours)	1.083 $\pm$ 7e-3	1.304 $\pm$ 3e-2	0.09 $\pm$ 1e-2

Appendix

Table 7: Full results over

3

different seeds. Performance on the 5-dim PCA of hESC dataset.

{\mathcal{W}}_{2}

is computed between test

\rho_{t_{i}}

and generated

\hat{\rho}_{t_{i}}

by simulating the dynamics from test

\rho_{t_{0}}

	${\mathcal{W}}_{1}\downarrow$
Methods	$t_{1}$	$t_{2}$	$t_{3}$	$t_{4}$	Mean
Neural SDE † ( Li et al., 2020 )	$0.69$	$0.91$	$0.85$	$0.81$	$0.82$
TrajectoryNet † ( Tong et al., 2020 )	$0.73$	$1.06$	$0.90$	$1.01$	$0.93$
IPF (GP) † ( Vargas et al., 2021 )	$0.70$	$1.04$	$0.94$	$0.98$	$0.92$
IPF (NN) † ( Bortoli et al., 2021 )	$0.73$	$0.89$	$0.84$	$0.83$	$0.82$
SB-FBSDE † ( Chen et al., 2022a )	$\mathbf{0.56}$	$0.80$	$1.00$	$1.00$	$0.84$

Appendix

Table 8: Full results over

3

different seeds. Performance on the 5-dim PCA of EB dataset.

{\mathcal{W}}_{1}

is computed between test

\rho^{\texttt{te}}_{t_{i}}

and generated

\hat{\rho}_{t_{i}}

by simulating the dynamics from previous test

\rho^{\texttt{te}}_{t_{i-1}}

Explore similar work

Jul 18, 2026stat.ML

Twisted Schrödinger Bridge 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

Schrödinger Bridges Differentiable Optimal Transport Stochastic Differential Equations

Ask PDF

Oct 1, 2025cs.LG

Multi-Marginal Flow Matching with Adversarially Learnt Interpolants

Learning the dynamics of a process given sampled observations at several time points is an important but difficult task in many scientific applications. When no ground-truth trajectories are available, but one has only snapshots of data taken at discrete time steps, the problem of modelling the dynamics, and thus inferring the underlying trajectories, can be solved by multi-marginal generalisations of flow matching algorithms. This paper proposes a novel flow matching method that overcomes the limitations of existing multi-marginal trajectory inference algorithms. Our proposed method, ALI-CFM, uses a GAN-inspired adversarial loss to fit neurally parametrised interpolant curves between source and target points such that the marginal distributions at intermediate time points are close to the observed distributions. The resulting interpolants are smooth trajectories that, as we show, are unique under mild assumptions. These interpolants are subsequently marginalised by a flow matching algorithm, yielding a trained vector field for the underlying dynamics. We showcase the versatility and scalability of our method by outperforming the existing baselines on spatial transcriptomics and cell tracking datasets, while performing on par with them on single-cell trajectory prediction. Code: https://github.com/mmacosha/adversarially-learned-interpolants.

Oskar Kviman, Kirill Tamogashev, Nicola Branchini +3

KTH · University of Edinburgh

Conditional Flow Matching Interpolation Stochastic Interpolants

Ask PDF

May 1, 2026cs.LG

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots

Inferring cellular trajectories from destructive snapshots is complicated by the challenges of stochasticity and non-conservative mass dynamics such as cell proliferation and apoptosis. Existing unbalanced Optimal Transport (OT) methods treat mass as a continuous fluid, performing inference at the population level. However, this macroscopic view often fails to capture the discrete, jump-like nature of birth-death events at single-cell resolution, which is essential for understanding lineage branching and fate decisions. We present Unbalanced Schrödinger Bridge (USB), a simulation-free framework for learning underlying dynamics that effectively integrates both stochastic and unbalanced effects which also models the discrete, jump-like birth-death dynamics at single-cell resolution. Theoretically, USB provides a tractable solution to the Branching Schrödinger Bridge (BSB) problem, offering a rigorous microscopic interpretation where individual cells undergo both Brownian motion and discrete birth-death jumps. Technically, the method implements an efficient solver by introducing a simulation-free training objective that effectively scales to high-dimensional omics data. Empirically, we demonstrate on both simulated and real-world datasets that USB not only achieves trajectory reconstruction performance better than or comparable to deterministic baselines but also uniquely enables realistic discrete simulation of birth-death dynamics at single-cell resolution.

Junda Ying, Yuxuan Wang, Bowen Yang +2

Beijing International Center for Mathematical Research, Peking University · Center for Data Science, Peking University · School of Mathematical Sciences, Peking University +4

Schrödinger Bridges Population Dynamics Discrete

Ask PDF