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. · National Engineering Laboratory for Big Data Analysis and Applications, Beijing. · AI for Science Institute, Beijing. · Institute for Artificial Intelligence, Peking University.
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)
W1 ( ↓ )
RME ( ↓ )
W1 ( ↓ )
RME ( ↓ )
W1 ( ↓ )
RME ( ↓ )
TIGON
0.045
0.014
0.512
0.047
2.263
0.127
DeepRUOT
0.043 ± 0.002
0.017 ± 0.001
0.623 ± 0.032
0.065 ± 0.011
3.785 ± 0.009
0.303 ± 0.070
Var-RUOT
0.079 ± 0.003
0.008 ± 0.002
0.522 ± 0.008
0.177 ± 0.007
2.813 ± 0.004
0.041 ± 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 W1 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
10,000
20,000
30,000
40,000
49,302
DeepRUOT
Time (s)
493.14±58.24
818.73±11.89
1290.99±63.32
1850.74±196.53
2204.64±171.37
Mem (GB)
2.76±0.01
6.16±0.00
10.22±0.00
14.98±0.00
19.73±0.00
VarRUOT
Time (s)
2707.19±319.11
2689.76±277.53
2701.46±284.03
2754.84±251.65
2756.25±281.77
Mem (GB)
2.57±0.00
3.55±0.00
4.58±0.01
5.64±0.01
6.63±0.00
SUDO
Time (s)
221.42±2.83
347.80±5.55
489.48±10.65
622.74±20.02
744.72±12.36
Table 3: Scalability evaluation on the Mouse dataset (Training time and peak memory consumption)
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.
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∈{0.01,0.1,1,2,6,12,15} ; Rows: r∈{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∈{0.01,0.1,1,2,6,12} ; Rows: r∈{0.001,0.01,0.1,0.5,1,2,5,10,20} .
δ
E2 (%) ↓
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 L2 norm error ( E2 ) between the Cd(1,r) and Eξ(d,r) under different training settings.
Parameter
t=1
t=2
t=3
t=4
W1
RME
W1
RME
W1
RME
W1
RME
δ=0.5
0.036
0.001
0.037
0.010
0.044
0.005
0.045
0.015
δ=1.0
0.023
0.002
0.027
0.006
0.022
0.007
0.022
0.011
δ=1.2
0.024
0.008
0.025
0.001
0.020
0.002
0.019
0.005
δ=1.5
0.024
0.002
0.031
0.009
0.031
0.011
0.029
0.002
δ=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 δ on simulation gene dataset.
Parameter
t=1
t=2
t=3
t=4
W1
RME
W1
RME
W1
RME
W1
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 Ψ on simulation gene dataset.
Method
t=1
t=2
t=3
t=4
W1
RME
W1
RME
W1
RME
W1
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
W1
RME
W1
RME
2.315
0.061
2.333
0.016
Appendix
Table 8: The measure matching results on gaussian 1000d data with different growth penalty
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
W1
RME
W1
RME
W1
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
W1
RME
W1
RME
W1
RME
W1
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
W1
RME
W1
RME
W1
RME
W1
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
W1
RME
W1
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∈{0.1,0.5,1,2,3,5} ; Rows: r∈{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∈{0.1,0.5,1,2,3,5} ; Rows: r∈{0.2,0.5,0.8,1} .
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
Unbalanced optimal transport (UOT) provides a principled framework for modeling single-cell transitions and birth-death dynamics, but its high computational cost limits scalability to large-scale datasets. Although single-cell data often contain hierarchical annotations and known transition priors, existing UOT approximations rarely exploit this multiscale structure or prior knowledge. We introduce Multiscale Supervised Unbalanced Optimal Transport Flow Matching (MUST-FM), a simulation-free framework that scales UOT by leveraging hierarchical data structure. MUST-FM further supports an optional supervised formulation that incorporates transition priors, such as cell lineages, to guide the learning of displacement fields and mass variations. Experiments show that MUST-FM reduces computational overhead while achieving robust and biologically meaningful trajectory inference, enabling dynamic modeling of atlas-scale single-cell datasets.
Qiangwei Peng, Lezhi Chen, Peijie Zhou
College of Electrical Engineering Sichuan University, Sichuan, China · AI for Science Institute, Beijing, China
This paper introduces a novel simulation-free framework for learning unbalanced population dynamics on the Riemannian Wasserstein-Fisher-Rao (WFR) manifold, termed Variational Doob-Wasserstein Flow (VDWF), built upon the theory of the unbalanced mean-field Schrödinger bridge (MFSB). First, by independently learning a geometry model equipped with a data-dependent metric, our model captures the intrinsic geometric structure of population dynamics. Furthermore, leveraging Doob's h-transform, the proposed method incorporates boundary information to improve flow-based trajectory inference. For unbalanced transport, a growth model is employed to learn the dynamics while constructing intermediate mass samples in the WFR space without simulation. Additionally, we analyze the convergence rate of the proposed framework using a uniform Log-Sobolev inequality (LSI). Finally, we validate the proposed approach on multiple real-world datasets, demonstrating competitive performance in learning both balanced and unbalanced population dynamics compared with state-of-the-art methods.
Weizhu Qian
School of Computer Science and Technology Soochow University, China