MIOFlow 2.0: A unified framework for inferring cellular stochastic dynamics from single cell and spatial transcriptomics data
Authors: Xingzhi Sun, João Felipe Rocha, Brett Phelan, Dhananjay Bhaskar, Guillaume Huguet, Yanlei Zhang, Alexander Tong, Ke Xu, Oluwadamilola Fasina, Mark Gerstein, Natalia Ivanova, Christine L. Chaffer, Guy Wolf, Smita Krishnaswamy
Organizations: Yale University, USA · University of Wisconsin-Madison, USA · Université de Montréal; Mila - Quebec AI Institute, Canada · University of Georgia, USA · Garvan Institute of Medical Research, Australia
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.
Figures & tables
Figure 1. Overview of the MIOFlow model. A. We initialize with scRNA-seq data and spatial transcriptomics, then concatenate both feature sets into a jointly embedded latent space. Each data point in this latent space represents a cell embedding informed by its neighbors. B. The embedding serves as input to three networks: a proliferation network predicting the proliferation rate, and drift and diffusion networks comprising the SDE/ODE model. C. The resulting trajectories can be visualized over the latent space.
Figure 2. features Extraction for MIOFlow 2.0. A. Build the neighborhood graph using knn graph or Voronoi polygons B. Compute local cell type frequency from neighborhood. Colors indicate cell types. C. Compute ligand-receptors signalling strength from neighbors to target cell. D. Local Expression Niche: The mean PCA embedding vector of neighboring cells E. Concatenate cell features with their spatial neighbors information
Figure 3. Comparison of trajectory inference methods on synthetic SERGIO datasets. (Top) Trifurcation dataset with three terminal fates. (Bottom) S-shaped dataset with cyclic progression and bifurcation. Cells are colored by ground truth timepoint. Predicted mean branch trajectories for each method are shown as colored curves. MIOFlow 2.0 maintains close adherence to the data manifold, while baseline methods deviate into unpopulated regions or fail to capture fine-scale curvature.
Dataset
Method
S-shaped
Trifurcation
MIOFlow
0.082±0.070
0.243±0.202
SF2M (ODE)
0.134±0.122
0.278±0.240
SF2M (SDE)
0.115±0.091
0.320±0.285
OT-CFM
0.118±0.095
0.298±0.250
Table 1. Trajectory reconstruction error (mean ± std) from held-out test samples to the nearest point on predicted trajectories.
Method
Held-out t
W1 ↓
MMD-G ↓
MMD-M ↓
MIOFlow
1
0.8072
0.1396
0.7700
MIOFlow
2
0.4689
0.0255
0.1973
MIOFlow
3
0.4101
0.0209
0.3646
OT-CFM
1
0.7924
0.1442
0.7709
OT-CFM
2
0.4741
0.0296
0.2034
OT-CFM
3
0.4612
0.0299
0.4343
Table 2. Leave-one-out trajectory inference accuracy on the EB dataset. Metrics are computed in raw PCA space (50D). W1: Wasserstein-1; MMD-G: Gaussian MMD; MMD-M: mean discrepancy. Best values are in bold.
Figure 4. We evaluate MIOFlow 2.0 across three simulated datasets designed to mimic key biological characteristics: branching, population decline (dying), and proliferation (growing). By comparing the base MIOFlow 2.0 framework with variants incorporating the growth-rate model and Neural SDEs, we visualize the resulting trajectories. The results demonstrate that these integrated biological priors more accurately capture the geometry of the data manifold and the underlying dynamics compared to the base interpolation.
Figure 5. Gene and spatial feature embeddings colored by (A) cell type annotation, (B) pseudotime point. (C) NCAN:SDC3 ligand-receptor signalling spatial feature. (D) Trajectories traced on top of gene embedding (left) and joint projection of gene and spatial embeddings (right) colored by decoded NCAN:SDC3 signalling feature. (E) Spatial organization of cells in Stereo-seq dataset colored by NCAN:SDC3 signalling feature. Note that not all cells shown are included in the trajectories, but all are used for computing spatial features. (F) Decoded NCAN:SDC3 signalling feature over time, averaged across trajectories with +/- one standard deviation. (G) Decoded gene trajectories for a selection of highly variable genes.
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.
Yuhao Sun, Zekun Wu, Zixun Huang +1
Center for Machine Learning Research, Peking University · School of Mathematical Sciences, Peking University · Center for Quantitative Biology, Peking University +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 RNA sequencing provides insights into gene expression at single-cell resolution, yet inferring temporal processes from these static snapshot measurements remains a fundamental challenge. Current approaches utilizing neural differential equations and flows are sensitive to overfitting and lack careful considerations of biological variability. In this work, we propose a generative framework that models population trends using a latent heteroscedastic Gaussian process (GP) approximated by Hilbert space methods. To address the absence of genuine cell trajectories, we leverage an optimal transport (OT) objective that aligns generated and observed population distributions. Our method explicitly captures biological heterogeneity by incorporating cell-specific latent time and cell type conditioning to disentangle temporal asynchrony and trajectories to different cell types. We demonstrate state-of-the-art performance on complex interpolation and extrapolation benchmarks and introduce a novel gradient-based strategy for inferring perturbation trajectories.
Mehmet Yigit Balik, Harri Lähdesmäki
Department of Computer Science, Aalto University, Espoo, Finland.