Simulation-free Unbalanced Dynamic Optimal Transport with General Growth Penalty

Junda Ying; Yuxuan Wang; Bowen Yang; Peijie Zhou; Lei Zhang

cs.LGSep 4, 2026

Simulation-free Unbalanced Dynamic Optimal Transport with General Growth Penalty

Authors: Junda Ying, Yuxuan Wang, Bowen Yang, Peijie Zhou, Lei Zhang

Organizations: Beijing International Center for Mathematical Research, Peking University. · Center for Quantitative Biology, Peking University. · School of Mathematical Sciences, Peking University. · Center for Machine Learning Research, Peking University. ·

Abstract

Inferring cellular dynamics from unpaired single-cell snapshots requires modeling both state transitions and population growth or death. Unbalanced dynamic optimal transport (UDOT) addresses this by penalizing growth along transport paths, making the choice of growth penalty a key way to encode biological priors on proliferation and apoptosis. However, existing UDOT solvers either rely on computationally expensive NeuralODE simulations or depend on analytical solutions of conditional paths, restricting their efficiency solely to quadratic penalties, i.e. Wasserstein-Fisher-Rao (WFR) geodesics. To enable an efficient UDOT solver for general growth penalties, we first show that concave growth penalties lead to degenerate solutions where growth and transport are separated. We then introduce Simulation-free Unbalanced Dynamic Optimal transport (SUDO), a simulation-free framework for UDOT with general non-quadratic convex growth penalties. SUDO learns the conditional paths and transport costs, solves the induced semi-coupling problem, and subsequently leverages unbalanced flow matching to achieve a simulation-free solution. On WFR benchmarks, SUDO matches the accuracy of efficient, analytical solution-driven algorithms while outperforming simulation-based methods in computational speed. Beyond WFR, SUDO supports asymmetric penalties that encode proliferation-dominant priors and produce more plausible trajectories and growth estimates on synthetic and single-cell datasets.

Figures & tables

Method	Simulation (2D)		Dyngen (5D)		Gaussian (1000D)
	$\mathcal{W}_{1}$ ( $\downarrow$ )	RME ( $\downarrow$ )	$\mathcal{W}_{1}$ ( $\downarrow$ )	RME ( $\downarrow$ )	$\mathcal{W}_{1}$ ( $\downarrow$ )	RME ( $\downarrow$ )
TIGON	0.045	0.014	0.512	0.047	2.263	0.127
DeepRUOT	0.043 $\pm$ 0.002	0.017 $\pm$ 0.001	0.623 $\pm$ 0.032	0.065 $\pm$ 0.011	3.785 $\pm$ 0.009	0.303 $\pm$ 0.070
Var-RUOT	0.079 $\pm$ 0.003	0.008 $\pm$ 0.002	0.522 $\pm$ 0.008	0.177 $\pm$ 0.007	2.813 $\pm$ 0.004	0.041 $\pm$ 0.006
WFR-FM	0.019	0.001	0.135	0.005	2.233	0.044
SUDO	0.022	0.002	0.190	0.005	2.315	0.061

Table 1: Mean

\mathcal{W}_{1}

and RME on synthetic datasets. For the methods that exhibit randomness in inference, we report the mean value and standard deviation over 5 runs.

Algorithm	Simulation (2D)	Dyngen (5D)	EMT (10D)	Gaussian (1000D)
WFR-FM	1.0810	7.4620	0.7726	7.6843
SUDO	1.0703	7.7064	0.7741	7.1576
Static Reference	1.0935	7.8403	0.9545	9.9153

Table 2: The UDOT cost of SUDO and WFR-FM. Both algorithms are deterministic.

Algorithm	Metric	Cell number
Algorithm	Metric	10,000	20,000	30,000	40,000	49,302
DeepRUOT	Time (s)	$493.14\pm 58.24$	$818.73\pm 11.89$	$1290.99\pm 63.32$	$1850.74\pm 196.53$	$2204.64\pm 171.37$
DeepRUOT	Mem (GB)	$2.76\pm 0.01$	$6.16\pm 0.00$	$10.22\pm 0.00$	$14.98\pm 0.00$	$19.73\pm 0.00$
VarRUOT	Time (s)	$2707.19\pm 319.11$	$2689.76\pm 277.53$	$2701.46\pm 284.03$	$2754.84\pm 251.65$	$2756.25\pm 281.77$
VarRUOT	Mem (GB)	$2.57\pm 0.00$	$3.55\pm 0.00$	$4.58\pm 0.01$	$5.64\pm 0.01$	$6.63\pm 0.00$
SUDO	Time (s)	$221.42\pm 2.83$	$347.80\pm 5.55$	$489.48\pm 10.65$	$622.74\pm 20.02$	$744.72\pm 12.36$

Table 3: Scalability evaluation on the Mouse dataset (Training time and peak memory consumption)

Panel of Figure 1 — Figure 1: Left: predicted trajectories by WFR-FM and SUDO; Right: predicted growth rate by WFR-FM and SUDO. We also plot the particles and their growth rate at t=0.5.

Panel of Figure 2 — Figure 2: Upper: predicted growth rate by WFR penalty; Lower: predicted growth rate by Only-growth penalty; From left to right: Time point 0, 1, 2. The colorbar is clipped to show the sign of the growth rate more clearly.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Red curve: mass variation learned by path model; Blue curve: True mass variation. Columns: $d\in\{0.01,0.1,1,2,6,12,15\}$ ; Rows: $r\in\{0.001,0.01,0.1,0.5,1,2,5,10,20\}$ .

Figure 4: Red curve: position variation learned by path model; Blue curve: True mass variation. Columns: $d\in\{0.01,0.1,1,2,6,12\}$ ; Rows: $r\in\{0.001,0.01,0.1,0.5,1,2,5,10,20\}$ .

$\delta$	$E_{2}$ (%) $\downarrow$
$\delta$	$N=64$	$N=32$	$N=16$
0.5	8.36	8.29	8.64
1	7.18	7.16	7.25
2	7.63	7.57	7.62
3	6.71	6.65	6.69
5	4.20	3.91	4.01

Appendix

Table 4: Relative

L^{2}

norm error (

E_{2}

) between the

C_{d}(1,r)

and

\mathcal{E}_{\xi}(d,r)

under different training settings.

Parameter	t=1		t=2		t=3		t=4
	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME
$\delta=0.5$	0.036	0.001	0.037	0.010	0.044	0.005	0.045	0.015
$\delta=1.0$	0.023	0.002	0.027	0.006	0.022	0.007	0.022	0.011
$\delta=1.2$	0.024	0.008	0.025	0.001	0.020	0.002	0.019	0.005
$\delta=1.5$	0.024	0.002	0.031	0.009	0.031	0.011	0.029	0.002
$\delta=2.0$	0.032	0.004	0.040	0.008	0.031	0.014	0.044	0.024

Appendix

Table 5: Sensitivity analysis for parameter

\delta

on simulation gene dataset.

Parameter	t=1		t=2		t=3		t=4
	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME
WFR	0.024	0.008	0.025	0.001	0.020	0.002	0.019	0.005
Only-growth	0.027	0.004	0.031	0.009	0.022	0.011	0.019	0.014
Only-death	0.046	0.001	0.057	0.001	0.052	0.006	0.029	0.009
No preference	0.028	0.009	0.032	0.020	0.027	0.021	0.027	0.029

Appendix

Table 6: Sensitivity analysis for penalty

\Psi

on simulation gene dataset.

Method	t=1		t=2		t=3		t=4
	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME
TIGON	0.446	0.033	0.584	0.060	0.415	0.023	0.603	0.071
DeepRUOT	0.454	0.011	0.481	0.070	0.870	0.104	0.688	0.074
Var-RUOT	0.315	0.128	0.548	0.336	0.630	0.222	0.593	0.023
WFR-FM	0.110	0.003	0.098	0.007	0.211	0.008	0.121	0.002
SUDO	0.140	0.003	0.168	0.007	0.280	0.005	0.170	0.006

Appendix

Table 7: Comparison of method performance over time on the Dyngen dataset. Best results are in bold, and the second best are underlined.

WFR		Only-growth
$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME
2.315	0.061	2.333	0.016

Appendix

Table 8: The measure matching results on gaussian 1000d data with different growth penalty

Panel of Figure 5 — Figure 5: Left: predicted trajectories and growth rate by WFR-FM; Right: predicted trajectories and growth rate by SUDO.

Method	t=1		t=2		t=3
	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME
TIGON	0.2433	0.002	0.2661	0.003	0.2847	0.001
DeepRUOT	0.2902	0.001	0.3193	0.011	0.3291	0.002
Var-RUOT	0.2540	0.075	0.2670	0.014	0.2683	0.041
WFR-FM	0.2099	0.001	0.2272	0.002	0.2346	0.001
SUDO	0.2305	0.002	0.2485	0.003	0.2556	0.002

Appendix

Table 9: Comparison of method performance over time on the 10D EMT dataset.

Method	t=1		t=2		t=3		t=4
	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME
TIGON	10.547	0.014	12.926	0.052	13.897	0.107	14.945	0.096
DeepRUOT	10.256	0.002	11.103	0.074	11.529	0.136	12.406	0.047
Var-RUOT	11.746	0.091	12.237	0.024	12.957	0.150	13.335	0.074
WFR-FM	9.941	0.009	11.040	0.006	11.516	0.008	12.664	0.005
SUDO	10.231	0.015	11.389	0.009	11.914	0.006	13.118	0.002

Appendix

Table 10: Comparison of method performance over time on the 100D EB dataset. Best results are in bold, and the second best are underlined.

Batch Size	t=1		t=2		t=3		t=4
	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME
500	10.258	0.010	11.408	0.016	11.993	0.002	13.360	0.008
1000	10.250	0.013	11.364	0.014	11.902	0.005	13.151	0.014
2000	10.231	0.015	11.389	0.009	11.914	0.006	13.118	0.002
3000	10.213	0.017	11.436	0.014	12.010	0.004	13.348	0.010
4000	10.174	0.015	11.379	0.010	11.898	0.001	13.177	0.004

Appendix

Table 11: Sensitivity analysis for batch size of mini-batch OT chunk size on the 100D EB dataset.

Method	t=1		t=2
	$\mathcal{W}_{1}$	RME	$\mathcal{W}_{1}$	RME
TIGON	6.140	0.382	6.973	0.326
DeepRUOT	6.052	0.062	6.757	0.041
Var-RUOT	7.951	0.131	10.862	0.154
WFR-FM	5.486	0.012	6.211	0.011
SUDO	5.894	0.027	6.824	0.019

Appendix

Table 12: Comparison of method performance over time on the 50D Mouse dataset. Best results are in bold, and the second best are underlined.

Figure 6: Red curve: mass variation learned by only-growth penalty; Blue curve: mass variation induced by WFR penalty. Columns: $d\in\{0.1,0.5,1,2,3,5\}$ ; Rows: $r\in\{1,1.2,1.5,2,5\}$ .

Figure 7: Red curve: mass variation learned by only-death penalty; Blue curve: mass variation induced by WFR penalty. Columns: $d\in\{0.1,0.5,1,2,3,5\}$ ; Rows: $r\in\{0.2,0.5,0.8,1\}$ .

Optimal Transport

Optimal transport compares and transforms probability distributions by finding low-cost ways to move mass between them.

See all 296 papers

Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots
May 1, 2026Junda Ying, Yuxuan Wang, Bowen Yang +2Schrödinger Bridges Single-Cell Transcriptomics
Beijing International Center for Mathematical Research, Peking University · Center for Data Science, Peking University · School of Mathematical Sciences, Peking University +4
Multiscale Supervised Unbalanced Optimal Transport Flow Matching
May 15, 2026Qiangwei Peng, Lezhi Chen, Peijie ZhouFlow Matching Single-Cell Transcriptomics
College of Electrical Engineering Sichuan University, Sichuan, China · AI for Science Institute, Beijing, China
Learning Population Dynamics with Variational Doob-Wasserstein Flow
Oct 24, 2026Weizhu QianOptimal Transport Dynamical Systems
School of Computer Science and Technology Soochow University, China

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

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

Junda Ying, Yuxuan Wang, Bowen Yang, Peijie Zhou+1

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

Schrödinger Bridges Single-Cell Transcriptomics

Open paper AskPDF arXiv

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.

Abstract

Optimal Transport

Explore similar work

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

Multiscale Supervised Unbalanced Optimal Transport Flow Matching

Learning Population Dynamics with Variational Doob-Wasserstein Flow