Learning Explicit Single-Cell Dynamics Using ODE Representations

TL;DR

Cell-MNN introduces a locally linear latent ODE within an end-to-end encoder–decoder framework to model single-cell differentiation from snapshot data. By learning a state/time-conditioned linear operator in PCA-embedded latent space, it achieves competitive interpolation performance without OT preprocessing and enables direct extraction of gene interactions interpretable in the original gene space. The approach supports scalable amortized training across multiple datasets and demonstrates robustness to noise, outperforming baselines on large inflations and offering biologically plausible interactions validated by TRRUST. This combination of predictive accuracy and mechanistic interpretability advances trajectory reconstruction and GRN discovery in single-cell genomics, with potential implications for hypothesis generation and perturbation design.

Abstract

Modeling the dynamics of cellular differentiation is fundamental to advancing the understanding and treatment of diseases associated with this process, such as cancer. With the rapid growth of single-cell datasets, this has also become a particularly promising and active domain for machine learning. Current state-of-the-art models, however, rely on computationally expensive optimal transport preprocessing and multi-stage training, while also not discovering explicit gene interactions. To address these challenges we propose Cell-Mechanistic Neural Networks (Cell-MNN), an encoder-decoder architecture whose latent representation is a locally linearized ODE governing the dynamics of cellular evolution from stem to tissue cells. Cell-MNN is fully end-to-end (besides a standard PCA pre-processing) and its ODE representation explicitly learns biologically consistent and interpretable gene interactions. Empirically, we show that Cell-MNN achieves competitive performance on single-cell benchmarks, surpasses state-of-the-art baselines in scaling to larger datasets and joint training across multiple datasets, while also learning interpretable gene interactions that we validate against the TRRUST database of gene interactions.

Figures (10)

• Figure 1: (a) Single-cell interpolation: trajectories are evaluated by the earth mover’s distance (EMD) between predictions and the marginal distribution at a held-out time $t_{\text{val}}$. (b) Like a hypernetwork, Cell-MNN predicts a linear operator ${\bm{A}}_\theta({\bm{z}}, t)$ that approximates the local dynamics explicitly, whereas Neural ODEs (NODE) and Flow Matching (FM) models only output a velocity.
• Figure 2: Visualization of the meta-learning task of Cell-MNN's encoder: Rather than directly predicting the velocity at a given operating point, as in the Neural ODE framework, the MLP of Cell-MNN maps to the space of linear operators. Conditioned on the current system state, it predicts local linear approximations to the global dynamics.
• Figure 3: (a) Model comparison across the synthetically inflated datasets. We report mean $\pm$ standard deviation of the MMD metric, along with the average across datasets. Lower values indicate better performance. Standard deviation is computed over left-out time points.(b) Comparison of models jointly trained on Cite and Multi datasets to test potential for amortization. We report mean $\pm$ standard deviation of the EMD metric, along with the average across datasets.
• Figure 4: (a) Strongest predicted gene interactions by Cell-MNN for days 12–17 of the EB dataset, normalized to the range $[-1, 1]$. (b) UMAP projection of operators predicted by Cell-MNN on the EB dataset, showing that the model learns distinct dynamics at different time points. (c) Validation of predicted gene interactions by two Cell-MNN versions: For each source gene $j$, we classify each TRRUST edge $j\!\rightarrow\! i$ as activating or repressing using the sign of Cell-MNN's learned weight $w_{j\rightarrow i}$.
• Figure 5: UMAPs of the predicted operators by Cell-MNN across the five time ranges of EB. Points are colored by whether joint expression of the EN-1 marker genes FOXA2 and SOX17 is above the 95th percentile. Clustering indicates that Cell-MNN learns distinct dynamics for the EN-1 cell type.

Theorems & Definitions (2)

• Proposition 1: Extension of Proposition 1 of cimen2010sdre
• proof