Learning Explicit Single-Cell Dynamics Using ODE Representations

The paper proposes Cell-Mechanistic Neural Networks (Cell-MNN), an end-to-end encoder-decoder architecture that utilizes locally linearized ODEs to efficiently model single-cell differentiation dynamics and explicitly learn interpretable, biologically consistent gene interactions, outperforming current state-of-the-art methods in scalability and interpretability.

The Gist

Imagine you are trying to understand how a lump of clay (a stem cell) transforms into a specific shape, like a bird or a car (a specialized tissue cell). In biology, this process is called differentiation.

The problem is that scientists can't watch this happen in real-time. To see what's inside the clay, they have to smash the cell open and take a snapshot. By the time they take the next picture, the cell is dead. So, they end up with a pile of photos: some of "baby" cells, some of "teenage" cells, and some of "adult" cells, but no video showing the transformation.

Reconstructing the "movie" from these scattered photos is a huge challenge. Current methods are like trying to guess the movie by calculating the distance between every single frame of every other movie in existence. It's slow, expensive, and doesn't tell you why the clay changed shape.

Enter Cell-MNN, a new AI tool proposed by researchers at ISTA and the Chan Zuckerberg Initiative. Here is how it works, explained simply:

1. The "Local Map" Analogy

Imagine you are driving a car through a massive, foggy city. You don't know the whole map of the city, and you can't see the destination.

• Old AI (Neural ODEs): Tries to guess the entire route from start to finish in one giant, complex calculation. It's like trying to memorize the whole city's traffic patterns at once.
• Cell-MNN: Takes a different approach. It says, "I don't need to know the whole city. I just need to know which way to turn right now based on where I am."

Cell-MNN looks at a cell's current state and asks: "If I take a tiny step forward in time, what happens?" It creates a local linear map (a simple, straight-line rule) for that specific moment. It's like a GPS that only gives you the next turn, but it does it so accurately that if you keep asking for the next turn, you can reconstruct the entire journey.

2. The "Gene Orchestra" Metaphor

Inside a cell, thousands of genes are talking to each other. Some say "Turn on the muscle gene," while others say "Shut down the brain gene." This is a chaotic orchestra.

• The Problem: We don't know who is conducting the orchestra or which instruments are playing together.
• The Cell-MNN Solution: Because Cell-MNN uses simple math rules (linear equations) to predict the next step, it can actually write down the sheet music.
  • It doesn't just predict the future; it tells you which gene is pulling the strings.
  • If Gene A makes Gene B turn on, Cell-MNN highlights that connection. It's like the AI not only predicting the song but also telling you, "The violinist (Gene A) is making the drummer (Gene B) hit harder."

3. Why It's a Game Changer

The paper highlights three main superpowers of Cell-MNN:

• Speed & Scale (The "No Traffic Jam" Factor):
  Old methods tried to connect every single cell to every other cell to figure out the path. If you have a million cells, that's a trillion connections to calculate. It's like trying to schedule a meeting between every person on Earth. Cell-MNN skips this traffic jam. It learns the rules of the road directly, so it can handle massive datasets without crashing the computer.

• One-Size-Fits-All Training (The "Universal Translator"):
  Usually, if you want to study two different types of cells (like liver cells and heart cells), you have to train two separate AI models. Cell-MNN is like a universal translator. You can feed it data from many different experiments at once, and it learns a single, robust set of rules that works for all of them. This is a step toward a "foundation model" for biology.

• Transparency (The "Glass Box"):
  Most AI models are "black boxes"—they give an answer, but you don't know how they got there. Cell-MNN is a "glass box." Because it uses simple math rules, scientists can look inside and say, "Ah, the AI thinks Gene X controls Gene Y." They can then check this against existing biology books (like the TRRUST database) to see if the AI is right. And guess what? It turns out the AI is often right!

The Bottom Line

Cell-MNN is a new way to watch the movie of life by looking at still photos. Instead of guessing the whole plot at once, it figures out the tiny rules that govern each moment.

It's faster, it handles huge amounts of data, and most importantly, it doesn't just predict the future; it explains why it's happening by revealing the hidden conversations between genes. This could help scientists design better drugs, understand cancer, and figure out how to heal wounds by knowing exactly which "switches" to flip in the cell's control panel.

Technical Summary

Here is a detailed technical summary of the paper "Learning Explicit Single-Cell Dynamics Using ODE Representations" (Cell-MNN), published as a conference paper at ICLR 2026.

1. Problem Statement

The paper addresses two fundamental challenges in modeling single-cell differentiation dynamics:

1. Data Scarcity and Snapshot Nature: Single-cell sequencing destroys the cell during measurement, resulting in "snapshot" data (a single point in time per cell) rather than continuous trajectories. Reconstructing continuous dynamics from these snapshots is an ill-posed inverse problem.
2. Limitations of State-of-the-Art (SOTA) Methods:
   • Computational Bottleneck: Current SOTA methods (e.g., OT-MFM, DeepRUOT) rely on Optimal Transport (OT) preprocessing to infer trajectories. The Sinkhorn algorithm used for OT scales quadratically ($O(n^2)$) with the number of samples, making it intractable for large datasets.
   • Multi-Stage Training: These methods often require complex, multi-stage training pipelines, hindering amortized training (training a single model across multiple datasets).
   • Lack of Interpretability: Existing models typically learn black-box velocity fields (e.g., Neural ODEs) that predict cell movement but do not explicitly reveal the underlying gene regulatory interactions (GRNs) driving the dynamics.

2. Methodology: Cell-Mechanistic Neural Networks (Cell-MNN)

The authors propose Cell-MNN, an end-to-end encoder-decoder architecture that models cellular evolution using a locally linearized Ordinary Differential Equation (ODE).

Core Architecture

• Encoder (PCA + MLP):

  • Input gene expression vectors $x_t$ are projected into a low-dimensional latent space $z_t$ using Principal Component Analysis (PCA).
  • A Multi-Layer Perceptron (MLP) acts as a hypernetwork. Instead of predicting a velocity vector directly, it predicts a linear operator $A_\theta(z, t)$ conditioned on the current state $z$ and time $t$.
  • The dynamics are modeled as a locally linear ODE: $\dot{z} = A_\theta(z, t) z$.
  • This formulation allows the model to approximate complex global dynamics by stitching together local linear approximations, avoiding the intractable search space of global non-linear ODE discovery.

• Decoder (Analytical Solution):

  • Because the local dynamics are linear, the ODE admits a closed-form analytical solution: $z(t + \Delta t) = \exp(A_\theta \Delta t) z(t)$.
  • This eliminates the need for numerical ODE solvers (which are slow and unstable) and allows for exact, differentiable trajectory prediction.
  • The predicted latent state is projected back to the gene expression space via the PCA inverse: $\hat{x}_{t+\Delta t} = V_{PCA} z_{t+\Delta t}$.

Parameterization and Training

• Operator Parametrization: The linear operator $A_\theta$ is predicted in an eigen-decomposed form ($P \Lambda P^{-1}$) to facilitate matrix exponentiation. A regularizer ensures the invertibility of $P$.
• Loss Function: The model is trained end-to-end using Maximum Mean Discrepancy (MMD) to match the model-induced distribution $q_\theta$ with the empirical data distribution $p_t$ at various time points.
  • $L_{total} = L_{MMD} + \lambda_{kin} L_{kin} + \lambda_{inv} L_{inv}$
  • $L_{kin}$ regularizes kinetic energy to encourage smooth trajectories (optimal transport flows).
  • $L_{inv}$ ensures the stability of the eigen-decomposition.

Gene Interaction Discovery

• The linearity of the model allows the dynamics to be projected back into the original gene space: $\dot{x} = (V_{PCA} A_\theta V_{PCA}^\top) x$.
• The matrix $W = V_{PCA} A_\theta V_{PCA}^\top$ represents the gene regulatory network. The entry $W_{ij}$ quantifies the interaction strength from gene $j$ to gene $i$, providing an explicit, interpretable map of gene regulation.

3. Key Contributions

1. Novel Architecture: Introduction of Cell-MNN, the first end-to-end model for single-cell dynamics that uses a locally linearized ODE representation without OT preprocessing.
2. Scalability: By eliminating OT preprocessing, Cell-MNN avoids the $O(n^2)$ bottleneck, enabling training on datasets with hundreds of thousands of cells where SOTA methods fail (Out-of-Memory errors).
3. Amortized Training: The single-stage, end-to-end design allows the model to be trained jointly across multiple datasets, demonstrating superior transfer learning capabilities compared to multi-stage baselines.
4. Interpretability: The model explicitly learns gene-gene interaction weights. The authors validate these predictions against the TRRUST database (a curated collection of human transcriptional regulatory interactions), showing high accuracy in predicting activation/repression relationships.
5. Performance: Achieves State-of-the-Art (SOTA) performance on single-cell interpolation benchmarks (measured by Earth Mover's Distance) across three real-world datasets (Embryoid Body, CITE-seq, Multiome).

4. Experimental Results

• Interpolation Accuracy: On the EB, Cite, and Multi datasets, Cell-MNN achieved the lowest average EMD scores, outperforming baselines like OT-MFM, DeepRUOT, and OT-CFM. It was the only method to outperform a simple "OT-Interpolate" baseline (which treats OT maps as ground truth).
• Scalability: In experiments with synthetically inflated datasets (250k cells), OT-based baselines (OT-CFM, DeepRUOT) crashed due to memory constraints. Cell-MNN successfully trained and outperformed a mini-batch OT baseline (Batch-OT-CFM).
• Amortization: When jointly trained on Cite and Multi datasets, Cell-MNN outperformed OT-CFM and I-CFM, proving its ability to learn a shared dynamic prior across different biological contexts.
• Gene Interaction Validation:
  • Using an unsupervised classification task against TRRUST, Cell-MNN achieved an average F1 score of 69%, significantly outperforming SCODE (46%) and the Jacobian of OT-CFM (48%).
  • Introducing an inductive bias (fixing one eigenvalue to zero to represent static directions) further improved the F1 score to 69% while maintaining competitive interpolation performance.

5. Significance and Impact

• Bridging Prediction and Interpretation: Cell-MNN uniquely bridges the gap between high-accuracy trajectory prediction and the discovery of biologically meaningful gene regulatory networks. Previous methods had to choose between predictive power (OT-based) or interpretability (static GRN methods).
• Foundation for Control: By modeling dynamics in a locally linear form, the paper opens the door to applying control theory to biology. This could enable the design of specific gene perturbations (e.g., CRISPR edits) to steer cell states toward desired configurations (e.g., regenerating tissue or halting cancer progression).
• Efficiency: The removal of OT preprocessing makes the method accessible for large-scale single-cell atlases, a critical step as datasets continue to grow exponentially.

In summary, Cell-MNN represents a paradigm shift in single-cell dynamics modeling, moving from computationally expensive, black-box trajectory inference to efficient, interpretable, and scalable mechanistic modeling.