Organizations: Center for Machine Learning Research, Peking University · School of Mathematical Sciences, Peking University · Center for Quantitative Biology, Peking University · National Engineering Laboratory for Big Data Analysis and Applications, Beijing · AI for Science Institute, Beijing
Inferring continuous system evolution from sparse temporal snapshots is a key challenge in generative modeling and single-cell omics. While Optimal Transport (OT) is popular, existing frameworks are largely restricted to first-order dynamics, assuming memoryless velocity fields. This limits expressiveness, as first-order systems fail to account for regulatory momentum and time-delayed responses inherent in processes like cell differentiation. Here, we introduce TracingFlow, a simulation-free Flow Matching framework generalizing to second-order dynamics. By using neural networks to regress the acceleration field, TracingFlow provides an exact, efficient solution to the Dynamical Optimal Acceleration Transport (DOAT) problem. Unlike first-order methods yielding over-smoothed trajectories, our second-order formulation captures high-curvature transitions and nonlinear evolutions by learning the underlying force fields. Evaluated on complex synthetic and large-scale scRNA-seq datasets, TracingFlow achieves superior accuracy in distributional reconstruction and trajectory faithfulness. Moreover, by integrating lineage tracing priors, it recovers dynamical structures that are both mathematically optimal and biologically plausible.
Figures & tables
Figure 1: An illustration figure for TracingFlow
Method
Dynamics
2D Simulation
Cite 5D
Cite 100D
W1
W2
W1
W2
W1
W2
OT-CFM
1st Order
0.5270 ± 0.0612
0.5922 ± 0.0813
0.8298 ± 0.0162
0.9155 ± 0.0160
10.5498 ± 0.0291
10.6220 ± 0.0310
SF 2 M
0.9926 ± 0.1857
1.1318 ± 0.1734
0.7650 ± 0.0375
0.8956 ± 0.0408
12.3426 ± 0.0735
12.4757 ± 0.0783
3MSBM
2nd Order
1.1855 ± 0.2636
1.2767 ± 0.2215
3.6369 ± 0.3508
3.7703 ± 0.3178
15.7708 ± 0.1473
15.8833 ± 0.1832
MMFM
2.3689 ± 0.1830
2.6466 ± 0.1769
2.5346 ± 0.1520
2.7576 ± 0.1877
12.5353 ± 0.6268
12.7868 ± 0.7061
HRF
1.5390 ± 0.1130
1.6631 ± 0.1214
1.4016 ± 0.0437
1.5313 ± 0.0430
10.1532 ± 0.0643
10.2336 ± 0.0707
Table 1: Average W1 and W2 distances between the generated and ground truth distributions at various time points on the 2D Simulation , Cite 5D, and Cite 100D datasets for different algorithms.
Figure 2: On the Simulation 2D dataset: a) non-crossing paths learned by OT-CFM, and b) crossing paths learned by TracingFlow.
Table 4Figure 5
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Dynamics of simulation datasets. (a)2D simulation data, (b)3D simulation data and (c)3D simulation data projected to XY plane.
Method
Dynamics
t=1
t=2
t=3
t=4
W1
W2
W1
W2
W1
W2
W1
W2
OT-CFM
1st Order
0.1215 ± 0.0164
0.1441 ± 0.0164
0.3644 ± 0.1177
0.4143 ± 0.1168
0.5863 ± 0.2615
0.6238 ± 0.2657
1.0356 ± 0.2325
1.1867 ± 0.2890
SF 2 M
0.2389 ± 0.1217
0.2668 ± 0.1305
0.5700 ± 0.2574
0.6108 ± 0.2654
1.0630 ± 0.2769
1.3557 ± 0.2770
2.0986 ± 1.3198
2.2938 ± 1.3193
3MSBM
2nd Order
1.4029 ± 0.6892
1.4459 ± 0.6596
0.9318 ± 0.2553
1.0181 ± 0.2050
1.8079 ± 0.7761
1.9162 ± 0.7529
0.5993 ± 0.6176
0.7265 ± 0.6935
MMFM
1.3226 ± 0.1774
1.9287 ± 0.3002
1.8387 ± 0.5624
1.1779 ± 0.5114
1.8158 ± 0.1794
1.8740 ± 0.1701
4.4984 ± 0.2937
4.6059 ± 0.2977
HRF
1.2036 ± 0.0425
1.2096 ± 0.0244
0.9955 ± 0.0921
1.0009 ± 0.0873
1.9829 ± 0.1311
2.0040 ± 0.1582
1.9740 ± 0.1873
2.4379 ± 0.2157
Appendix
Table 4: W1 and W2 distances between the generated and ground-truth distributions at each time point on the 2D Simulation dataset.
Method
Dynamics
t=1
t=2
t=3
W1
W2
W1
W2
W1
W2
OT-CFM
1st Order
0.5615 ± 0.0156
0.6484 ± 0.0268
0.8472 ± 0.0463
0.9416 ± 0.0443
1.0536 ± 0.0534
1.1566 ± 0.0573
SF 2 M
0.5932 ± 0.0249
0.6797 ± 0.0424
0.7768 ± 0.0451
0.8670 ± 0.0480
0.9251 ± 0.0279
1.0221 ± 0.0352
3MSBM
2nd Order
5.3735 ± 0.3961
5.4244 ± 0.3582
4.3118 ± 0.4589
4.5458 ± 0.4141
1.2254 ± 0.2998
1.3427 ± 0.3307
MMFM
0.5335 ± 0.0134
0.5911 ± 0.0057
2.2290 ± 0.2151
2.4557 ± 0.1613
4.8413 ± 0.2538
5.2261 ± 0.4061
HRF
0.4225 ± 0.0152
0.4622 ± 0.0185
0.6227 ± 0.0421
0.6814 ± 0.0392
0.7638 ± 0.0738
0.8394 ± 0.0713
Appendix
Table 5: W1 and W2 distances between the generated and ground-truth distributions at each time point on the Cite 5D dataset
Method
Dynamics
t=1
t=2
t=3
W1
W2
W1
W2
W1
W2
OT-CFM
1st Order
10.1136 ± 0.0104
10.1873 ± 0.0135
10.2504 ± 0.0209
10.3145 ± 0.0201
11.2682 ± 0.0938
11.3642 ± 0.0997
SF 2 M
10.3598 ± 0.0051
10.4448 ± 0.0053
11.4647 ± 0.0384
11.5456 ± 0.0402
15.2032 ± 0.1813
15.4367 ± 0.1949
3MSBM
2nd Order
19.9677 ± 0.1373
20.1015 ± 0.1105
15.3218 ± 0.1232
15.4089 ± 0.1343
12.0229 ± 0.2895
12.1395 ± 0.3621
MMFM
9.8726 ± 0.0383
9.9305 ± 0.0361
12.1374 ± 0.2422
12.3076 ± 0.2495
15.5960 ± 1.4972
16.1222 ± 2.0195
HRF
9.9660 ± 0.0215
10.0580 ± 0.0341
9.9451 ± 0.0632
9.9986 ± 0.1526
10.5485 ± 0.1082
10.6442 ± 0.1254
Appendix
Table 6: W1 and W2 distances between the generated and ground-truth distributions at each time point on the Cite 100D dataset
Figure 6: Trajectories learned by TracingFlow on the (a) Cite 5D and (b) Cite 100D datasets, projected to 2D using PCA.
Method
Dynamics
t=1
t=2
t=3
t=4
W1
W2
W1
W2
W1
W2
W1
W2
OT-CFM
1st Order
3.4335 ± 0.1465
3.8462 ± 0.1449
4.8404 ± 0.1496
5.5792 ± 0.1358
5.0713 ± 0.2839
5.6448 ± 0.2861
7.1368 ± 0.3799
7.6938 ± 0.4574
SF 2 M
3.4204 ± 0.0637
3.6713 ± 0.0589
3.5511 ± 0.0394
3.9441 ± 0.0418
4.3175 ± 0.2517
4.8738 ± 0.3247
6.9516 ± 0.0616
7.4727 ± 0.0795
3MSBM
2nd Order
6.1086 ± 1.1215
6.5434 ± 1.2227
7.3626 ± 0.8043
8.0596 ± 1.0252
7.5862 ± 1.1170
8.1734 ± 0.7900
7.7029 ± 0.4241
7.1022 ± 0.3971
MMFM
10.7803 ± 0.6718
11.8376 ± 0.7823
5.7932 ± 0.2162
6.2242 ± 0.1828
9.3326 ± 0.5617
9.9837 ± 0.4994
13.2046 ± 2.1528
14.3464 ± 2.2982
TF (Ours)
3.5611 ± 0.0362
3.8930 ± 0.0243
3.5002 ± 0.0983
3.8035 ± 0.1086
3.5643 ± 0.1140
3.7947 ± 0.4123
7.3639 ± 0.1316
7.6523 ± 0.1976
Appendix
Table 7: W1 and W2 distances between predicted and ground-truth distributions at held-out time points on the EB 5D dataset.
Figure 7: Predicted distributions by TracingFlow at four held-out time points, projected to 2D using PCA.
Method
Dynamics
t=1
t=2
W1
W2
W1
W2
OT-CFM
1st Order
2.6080 ± 0.0029
2.6869 ± 0.0030
1.8686 ± 0.0382
1.9456 ± 0.0403
SF 2 M
2.8262 ± 0.0409
2.8373 ± 0.0462
0.2247 ± 0.0617
0.2533 ± 0.0568
3MSBM
2nd Order
4.0141 ± 0.8788
4.2262 ± 0.8596
0.7548 ± 0.0726
0.6751 ± 0.0573
MMFM
1.8585 ± 0.0350
2.1618 ± 0.0496
3.5695 ± 0.5692
4.1028 ± 0.6037
HRF
2.9150 ± 0.2595
2.9265 ± 0.3901
0.6717 ± 0.5491
0.6905 ± 0.4143
Appendix
Table 8: lineage-weighted W1 and W2 distances between the generated and ground-truth distributions at each time point on the 3D Simulation Lineage dataset
Method
Dynamics
t=1
t=2
W1
W2
W1
W2
OT-CFM
1st Order
13.2972 ± 0.0804
13.7134 ± 0.0831
16.6918 ± 0.1973
17.0807 ± 0.2025
SF 2 M
14.1427 ± 0.0341
14.4199 ± 0.0303
23.2913 ± 0.2590
23.5314 ± 0.2640
3MSBM
2nd Order
17.6487 ± 0.5250
17.8710 ± 0.5204
19.0139 ± 1.4418
19.4144 ± 1.3622
MMFM
18.4040 ± 2.8129
19.8538 ± 2.8906
36.0412 ± 3.0429
40.2282 ± 3.3934
HRF
14.1007 ± 0.0804
14.3122 ± 0.0831
16.8135 ± 0.1973
17.1154 ± 0.2025
Appendix
Table 9: lineage-weighted W1 and W2 distances between the generated and ground-truth distributions at each time point on the Hematopoiesis dataset
Method
Dynamics
t=1
t=2
Average
W1
W2
W1
W2
W1
W2
OT-CFM
First-order
9.8780
10.6132
10.7621
11.4466
10.3201
11.0299
SF 2 M
10.2352
10.8017
13.2507
13.6879
11.7430
12.2428
3MSBM
Second-order
12.8214
13.2501
20.0262
20.6374
16.4238
16.9438
MMFM
10.4305
10.8695
22.3269
23.5322
16.3787
17.2008
HRF
10.6233
11.4510
10.9570
11.6201
10.7902
11.5356
Appendix
Table 10: W1 and W2 distances between the generated and ground-truth distributions at each time point on the Hematopoiesis dataset
Figure 8: Original data (a), and data points generated by TracingFlow on the Hematopoiesis dataset without biological priors (b) and with biological priors (c), projected to 2D using UMAP.
Method
Dynamics
t=1
t=2
t=3
t=4
W1
W2
W1
W2
W1
W2
W1
W2
OT-CFM
First-order
0.0336
0.0375
0.0791
0.0887
0.1367
0.1415
0.1409
0.1474
SF 2 M
0.0768
0.0917
0.0802
0.0964
0.0819
0.0968
0.0809
0.0987
3MSBM
Second-order
0.4407
0.6305
0.5565
0.7617
0.5887
0.9113
0.5231
1.0714
MMFM
0.0346
0.0364
0.2961
0.3008
1.0538
1.0571
1.8755
1.8766
HRF
0.6205
0.6217
1.2004
1.2011
1.5334
1.5337
1.4986
1.5013
Appendix
Table 11: W1 and W2 distances between the generated and ground-truth distributions at each time point on the Mexico Gulf dataset
Method
Simulation
Cite 5D
Cite 100D
SimLineage-3D
RealLineage
OT-CFM
12 s
175 s
218 s
13 s
196 s
SF 2 M
26 s
79 s
84 s
34 s
61 s
MMFM
42 s
96 s
179 s
68 s
201 s
3 MSBM
220 s
495 s
910 s
244 s
855 s
CAF
94 s
274 s
382 s
256 s
264 s
HRF
247 s
271 s
359 s
289 s
347 s
Appendix
Table 12: Training time comparison (in seconds) of different methods across various datasets.
Batch Size
t=1
t=2
t=3
Average
W1
W2
W1
W2
W1
W2
W1
W2
1000
0.4344
0.4798
0.6522
0.7183
0.7658
0.8479
0.6175
0.6820
2000
0.4196
0.4624
0.5640
0.6251
0.7972
0.8820
0.5936
0.6565
Full Dataset
0.2868
0.3185
0.5487
0.5997
0.7689
0.8499
0.5348
0.5894
Appendix
Table 13: Sensitivity analysis of batch size for Minibatch-OT on the Cite-5D dataset.
Understanding cellular trajectories via time-resolved single-cell transcriptomics is vital for studying development, regeneration, and disease. A key challenge is inferring continuous trajectories from discrete snapshots. Biological complexity stems from stochastic cell fate decisions, temporal proliferation changes, and spatial environmental influences. Current methods often use deterministic interpolations treating cells in isolation, failing to capture the probabilistic branching, population shifts, and niche-dependent signaling driving real biological processes. We introduce Manifold Interpolating Optimal-Transport Flow (MIOFlow) 2.0. This framework learns biologically informed cellular trajectories by integrating manifold learning, optimal transport, and neural differential equations. It models three core processes: (1) stochasticity and branching via Neural Stochastic Differential Equations; (2) non-conservative population changes using a learned growth-rate model initialized with unbalanced optimal transport; and (3) environmental influence through a joint latent space unifying gene expression with spatial features like local cell type composition and signaling. By operating in a PHATE-distance matching autoencoder latent space, MIOFlow 2.0 ensures trajectories respect the data's intrinsic geometry. Empirical comparisons show expressive trajectory learning via neural differential equations outperforms existing generative models, including simulation-free flow matching. Validated on synthetic datasets, embryoid body differentiation, and spatially resolved axolotl brain regeneration, MIOFlow 2.0 improves trajectory accuracy and reveals hidden drivers of cellular transitions, like specific signaling niches. MIOFlow 2.0 thus bridges single-cell and spatial transcriptomics to uncover tissue-scale trajectories.
Xingzhi Sun, João Felipe Rocha, Brett Phelan +11
Yale University, USA · University of Wisconsin-Madison, USA · Université de Montréal; Mila - Quebec AI Institute, Canada +2
Single-cell RNA sequencing (scRNA-seq) provides high-dimensional profiles of cellular states, enabling data-driven modeling of cellular dynamics over time. In practice, time-resolved scRNA-seq is collected at only a few discrete time points as unpaired snapshot populations, leaving substantial temporal gaps. This motivates trajectory inference at unmeasured time points. Existing methods mainly follow two directions, optimal-transport (OT) alignment provides distribution-level matching between observed snapshots, while continuous-time generative models support forecasting via learned dynamics. However, two challenges remain: (i) unpaired snapshots render local transitions between adjacent time points ambiguous, leading to unstable supervision; and (ii) long-horizon prediction relies on repeated integration, where small modeling errors compound and cause distribution drift. To address these challenges, we propose single-cell Flow Matching (scFM), a latent generative framework based on coupling-conditioned flow matching. First, we compute entropically regularized OT couplings between adjacent snapshots and use them to construct soft, weighted flow-matching targets for learning time-dependent velocity fields. Second, we learn bidirectional velocity fields and leverage their consistency to refine couplings and improve temporal coherence under sparse supervision. Third, we introduce distribution-level alignment and latent dynamic regularization to anchor long rollouts and mitigate drift. Experiments on real-world time-series scRNA-seq datasets show that scFM consistently improves distributional prediction performance for both temporal interpolation and extrapolation. Moreover, scFM yields more accurate trajectory reconstruction and temporally coherent visualizations where intermediate time points are absent, indicating a more faithful recovery of underlying temporal gene expression dynamics.
Siyu Pu, Qingqing Long, Xiaohan Huang +7
Computer Network Information Center, Chinese Academy of Sciences · University of California, Los Angeles
Single-cell trajectory inference from destructive time-course snapshots is fundamentally ill-posed: neither cross-time cell correspondences nor continuous trajectories are observed, so the snapshot distributions alone do not uniquely determine the underlying dynamics. Existing optimal transport and flow-based methods typically couple cells by Euclidean proximity at observed clock times, which can misalign trajectories when development is asynchronous and cells sampled at the same experimental time occupy different latent pseudotime stages. We propose PACE, a trajectory inference framework that recovers geometry-consistent continuous transport dynamics from destructive time-course snapshots through three coupled components. First, PACE constructs a state- and time-dependent anisotropic Riemannian metric that assigns low transport cost along locally supported tangent directions while penalizing normal velocity components. Second, it alternates between refining cross-time couplings under the induced path-action cost and fitting endpoint-preserving neural bridges between adjacent snapshots. Third, it distills the learned bridge dynamics into a global continuous-time velocity field over cellular states. Across seven controlled and biological datasets covering nine held-out reconstruction experiments, PACE achieves the strongest overall reconstruction performance, reducing MMD, Wasserstein-1 distance, and Wasserstein-2 distance by 23.7% on average relative to the strongest competing baseline. PACE also improves RNA-velocity alignment by 15.4% on an embryoid body differentiation benchmark, without requiring explicit cell pairing, lineage tracing, or RNA-velocity supervision during training. Code is available at https://github.com/AI4Science-WestlakeU/PACE.
Chenglei Yu, Chuanrui Wang, Bangyan Liao +1
Zhejiang University · Department of Artificial Intelligence, School of Engineering, Westlake University