MIOFlow 2.0: A unified framework for inferring cellular stochastic dynamics from single cell and spatial transcriptomics data

Xingzhi Sun; João Felipe Rocha; Brett Phelan; Dhananjay Bhaskar; Guillaume Huguet; Yanlei Zhang; Alexander Tong; Ke Xu

cs.LGMar 23, 2026

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,

Organizations: Yale University, USA · University of Wisconsin-Madison, USA · Université de Montréal; Mila - Quebec AI Institute, Canada · University of Georgia, USA ·

Abstract

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\pm 0.070$	$0.243\pm 0.202$
SF2M (ODE)	$0.134\pm 0.122$	$0.278\pm 0.240$
SF2M (SDE)	$0.115\pm 0.091$	$0.320\pm 0.285$
OT-CFM	$0.118\pm 0.095$	$0.298\pm 0.250$

Table 1. Trajectory reconstruction error (mean

\pm

std) from held-out test samples to the nearest point on predicted trajectories.

Method	Held-out $t$	W1 $\downarrow$	MMD-G $\downarrow$	MMD-M $\downarrow$
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.

Trajectory Inference

Trajectory inference reconstructs continuous paths and underlying dynamics from sparse, noisy, or incomplete observations.

See all 37 papers

TracingFlow: A Simulation-Free Trajectory Inference Framework Based on Second-Order Dynamics
Aug 21, 2026Yuhao Sun, Zekun Wu, Zixun Huang +1Trajectory Inference Flow-Based Generative Modeling
Center for Machine Learning Research, Peking University · School of Mathematical Sciences, Peking University · Center for Quantitative Biology, Peking University +2
From Snapshots to Trajectories: Learning Single-Cell Gene Expression Dynamics via Conditional Flow Matching
May 21, 2026Siyu Pu, Qingqing Long, Xiaohan Huang +7Latent Dynamics Modeling Conditional Flow Matching
Computer Network Information Center, Chinese Academy of Sciences · University of California, Los Angeles
Modeling Temporal scRNA-seq Data with Latent Gaussian Process and Optimal Transport
May 20, 2026Mehmet Yigit Balik, Harri LähdesmäkiLatent Dynamics Modeling Conditional Generative Modeling
Department of Computer Science, Aalto University, Espoo, Finland.

Aug 21, 2026J/K move · Enter open · S save

TracingFlow: A Simulation-Free Trajectory Inference Framework Based on Second-Order Dynamics

Yuhao Sun, Zekun Wu, Zixun Huang, Peijie Zhou

Center for Machine Learning Research, Peking University · School of Mathematical Sciences, Peking University · Center for Quantitative Biology, Peking University+2

Trajectory Inference Flow-Based Generative Modeling Dynamical Systems Flow Matching Optimal Transport

Open paper AskPDF arXiv

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.

Abstract

Trajectory Inference

Explore similar work

TracingFlow: A Simulation-Free Trajectory Inference Framework Based on Second-Order Dynamics

From Snapshots to Trajectories: Learning Single-Cell Gene Expression Dynamics via Conditional Flow Matching

Modeling Temporal scRNA-seq Data with Latent Gaussian Process and Optimal Transport