Learning Explicit Single-Cell Dynamics Using ODE Representations

## Abstract

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.

## 1 Introduction

The process by which stem cells differentiate into specialized tissue cells is poorly understood, and prediction of cellular fate remains an open problem in systems biology.
Deeper understanding of the differentiation dynamics is essential for advancing treatment of diseases such as cancer [9], neurodegenerative diseases [10], and to improving wound healing [44].
While all cells in an organism share the same genome, the level of expression of genes varies over time as differentiation progresses.
During this process, genes activate or repress the expression of other genes through complex regulatory mechanisms, causing the cell to differentiate.

Today, only a small subset out of the large number of possible gene interactions has been thoroughly studied.
This is due to both the vast combinatorial search space, with $\sim 10^{8}$ theoretically possible gene interactions, and the experimental effort required to validate specific mechanisms.
However, recent advances in single-cell sequencing technology [35, 58] have enabled high-throughput measurements that were previously prohibitively expensive, producing datasets that are growing at a pace exceeding Moore’s law [25].
This rapid growth, coupled with the limitations of direct experimental approaches, presents a unique opportunity to apply ML methods to study single-cell dynamics.

In this work we propose Cell-MNN, a method to jointly tackle the challenges of predicting cell fate and discovering gene regulatory interactions.
Cell-MNN is an end-to-end encoder-decoder architecture whose representation is a locally linear ordinary differential equation (ODE) that governs the dynamics of cellular evolution from stem to tissue cells.
The ODE representation of Cell-MNN can learn explicit, biologically consistent, and interpretable gene interactions.

Figure: (a) Single-cell interpolation problem
Refer to caption: x1.png

A key challenge in modeling single-cell dynamics is that cells are destroyed by measurement, resulting in datasets that contain a single point along each cell’s trajectory [50], i.e., a snapshot observation.
This motivated a line of work on reconstructing trajectories from snapshot data:
The best-performing methods in this setting rely on optimal transport (OT) preprocessing to create label trajectories [52, 57, 24, 54], which becomes a computational bottleneck for large datasets due to quadratic scaling of the Sinkhorn algorithm with the number of samples [11].
In contrast, Cell-MNN eliminates OT preprocessing entirely and is designed to be end-to-end.
Another bottleneck of state-of-the-art (SOTA) models such as OT-MFM [24] and DeepRUOT [57] is that they involve multiple training stages and networks, making amortized training across datasets challenging, whereas Cell-MNN is trained in a single stage, enabling straightforward amortized training across multiple datasets.
Furthermore, existing SOTA methods focus primarily on accurate interpolation of empirical distributions and do not learn explicit gene regulatory interactions.
By comparison, Cell-MNN learns biologically interpretable interactions through its ODE representation, which explicitly models the interactions governing the predicted cellular evolution.
While there are dedicated methods for discovering gene regulatory interactions [32], to the best of our knowledge, no such method achieves SOTA predictive performance on single-cell interpolation benchmarks. Cell-MNN addresses both challenges simultaneously, bridging the gap between predictive performance and interpretable gene regulatory modeling.

#### Contributions.

Our main contributions are: (i) we propose Cell-MNN, an architecture that models single-cell dynamics via a locally linearized ODE representation; (ii) we demonstrate SOTA average performance on three benchmark datasets; (iii) we show that eliminating OT preprocessing enables scalability, with Cell-MNN outperforming all baselines on upsampled datasets; (iv) we leverage the end-to-end design for amortized training across datasets, surpassing a strong amortized baseline; and (v) we exploit the explicit ODE representation to extract gene interactions and quantitatively validate them against the TRRUST database [18] of gene interactions.

## 2 Learning the Dynamics of Cells

#### Formalizing the Problem.

We assume a data-generating process consisting of a cell state $\displaystyle{\bm{c}}(t)\in\mathcal{C}$ evolving over time in a high-dimensional state space $\mathcal{C}$ that includes all relevant molecular, physical, and biochemical variables, and an observation function mapping this state to data. The measurement process observes only a subset of the full state mapping it to the gene expression vector of $\displaystyle d_{x}$ genes $\displaystyle{\bm{x}}_{t}\in\mathbb{R}^{d_{x}}$ via an unknown, potentially noisy measurement process $\displaystyle{\bm{m}}:\mathcal{C}\rightarrow\mathbb{R}^{d_{x}}$, so that $\displaystyle{\bm{x}}(t)={\bm{m}}({\bm{c}}(t))$. Measuring the system involves deconstructing the observed cell, which implies that each measurement corresponds to a single point along its trajectory, i.e., a snapshot observation. We assume time $t\in\mathbb{R}$ to be a continuous variable and denote an arbitrary time interval by $\Delta t\in\mathbb{R}$. In practice, the lab schedules a discrete set of experimental time points $\displaystyle\mathcal{T}=\{t_{1},t_{2},\dots,t_{K}\}$ at which cell populations are sampled. We denote by $\displaystyle p_{t}$ the distribution of $\displaystyle{\bm{x}}_{t}$ at time $\displaystyle t$, and by $\displaystyle\mu_{t}$ its empirical estimate from the observations. The dataset of snapshot observations is $\displaystyle\mathcal{D}=\{{\bm{x}}^{(i)},t^{(i)}\}_{i=1}^{N}$ with $t^{(i)}\in\mathcal{T}$, and our goal is to learn a best-fit mechanistic model for the dynamics of the observable $\displaystyle{\bm{x}}_{t}$ that is consistent with the family of marginals $\displaystyle\{p_{t}\}_{t\in\mathcal{T}}$.

### 2.1 Cell-MNN

SOTA models on single-cell interpolation benchmarks face scalability issues from OT preprocessing and do not learn interpretable gene interactions that can be cross-validated against biological evidence.
Our goal is to design a scalable mechanistic model of single-cell dynamics using an ODE representation, enabling accurate forecasting and discovery of interpretable gene interactions.

The Mechanistic Neural Network (MNN) is a recent architecture that Pervez et al. [41] showed to outperform Neural ODEs on tasks such as solar system dynamics and the $n$-body problem, while also being able to learn explicit models of the underlying dynamics. This motivates us to design an MNN-inspired architecture for the single-cell setting. However, this domain presents unique challenges that make the vanilla MNN not directly applicable: for ODE discovery, the MNN has only been applied with full trajectories and not yet in biological contexts. Moreover, when identifying a latent space ODE with the MNN, there is typically no way to interpret that ODE in the input space. In contrast, single-cell dynamics require learning latent space dynamics from snapshot data. To discover gene interactions, the learned ODE must furthermore be interpretable in the input space. We therefore adapt the MNN architecture to this setting and refer to the resulting version as Cell-MNN. Cell-MNN is an encoder–decoder model, learning a mechanistic map

$$ ${\bm{x}}_{t+\Delta t}=\text{Cell-MNN}_{\theta}({\bm{x}}_{t},t,\Delta t),$ $$

which maps a gene expression vector $\displaystyle{\bm{x}}_{t}$ at time $\displaystyle t$ to a predicted state $\displaystyle{\bm{x}}_{t+\Delta t}$ after an *arbitrary* time interval $\Delta t$. We define the model-induced distribution at time $\displaystyle t+\Delta t$ as $\displaystyle q_{t+\Delta t}^{\theta}$, which is the distribution of $\text{Cell-MNN}_{\theta}({\bm{x}}_{t},t,\Delta t)$ when ${\bm{x}}_{t}$ is drawn from $p_{t}$.
As a core part of the architecture, Cell-MNN maps to a compressed representation $\displaystyle{\bm{z}}\in\mathbb{R}^{d_{z}}$, with $d_{z}\ll d_{x}$, of the high-dimensional gene expression vector $\displaystyle{\bm{x}}\in\mathbb{R}^{d_{x}}$, and learns the dynamics in the latent space.
Following prior work [52], we obtain this latent representation by applying principal component analysis (PCA), with projection matrix ${\bm{V}}_{\text{PCA}}\in\mathbb{R}^{d_{x}\times d_{z}}$, so that ${\bm{z}}={\bm{V}}_{\text{PCA}}^{\top}{\bm{x}}$.

#### Locally Linearizing the Latent ODE.

The latent vector $\displaystyle{\bm{z}}\in\mathbb{R}^{d_{z}}$ in the PCA subspace is assumed to follow non-autonomous, non-linear dynamics $\displaystyle\dot{{\bm{z}}}={\bm{f}}({\bm{z}},t)$. In practice, this ODE is often highly complex, and learning an *explicit* form that globally approximates it would be intractable due to the combinatorial search space of basis functions that grows with the latent space dimension $d_{z}$.

To address this, we decompose the intractable global ODE discovery problem into smaller subproblems: at the current state $\displaystyle({\bm{z}}^{(i)},t^{(i)})$, which we also call the operating point, we approximate the dynamics by a linear ODE in a small neighborhood. The learning task is then to predict these local dynamics models from the operating point $\displaystyle({\bm{z}}^{(i)},t^{(i)})$ using an encoder.

Figure: 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.
Refer to caption: x3.png

$$ $\displaystyle\dot{{\bm{z}}}$ $\displaystyle={\bm{f}}({\bm{z}},t)$ $\displaystyle={\bm{A}}({\bm{z}},t)\,{\bm{z}},$ $\displaystyle\text{if }{\bm{f}}(\mathbf{0},t)=\mathbf{0},\;\forall t\in\mathbb{R},$ $\displaystyle\approx{\bm{A}}_{\theta}\big({\bm{z}}^{(i)},t^{(i)}\big)\,{\bm{z}}.$ $$

We predict the linear operator $\displaystyle{\bm{A}}_{\theta}\in\mathcal{A}$
using a multilayer perceptron $\text{MLP}_{\theta}:\ \mathbb{R}^{d_{z}+1}\rightarrow\mathcal{A}$. Here $\mathcal{A}:=\mathcal{L}(\mathbb{R}^{d_{z}},\mathbb{R}^{d_{z}})\cong\mathbb{R}^{d_{z}\times d_{z}}$ represents the space of linear operators acting on $\mathbb{R}^{d_{z}}$.
Note that, while the operator governing the local dynamics is linear, it is a non-linear function of the current latent state $\displaystyle{\bm{z}}^{(i)}$ and time $\displaystyle t^{(i)}$. In Appendix [D](https://arxiv.org/html/2510.02903v1#A4), we show that the reparametrization of the right-hand side ${\bm{f}}({\bm{z}},t)={\bm{A}}({\bm{z}},t)\,{\bm{z}}$ always exists under mild assumptions.

This approach is conceptually orthogonal to Neural ODEs [8], which learn an *unconditional black-box* approximation to $\displaystyle{\bm{f}}$.
In the Cell-MNN setting, the MLP functions more like a hypernetwork [17], outputting a conditional white-box linear function ${\bm{g}}_{\theta}({\bm{z}},t|{\bm{z}}^{(i)},t^{(i)})={\bm{A}}_{\theta}({\bm{z}}^{(i)},t^{(i)})\,{\bm{z}}$ that locally approximates $\displaystyle{\bm{f}}$ at the operating point $({\bm{z}}^{(i)},t^{(i)})$.
Unlike most neural operators [31, 28] that learn a single global operator, Cell-MNN predicts a state-conditioned linear operator for each operating point.
This makes the learned dynamics explicit and enables amortization across arbitrarily many states and datasets within a single network.

#### Decoding by Analytically Solving the ODE.

Decoding the ODE representation involves solving the ODE system.
The locally linearized formulation of the dynamics has the advantage that the latent space ODE admits a local closed-form solution.
For fixed $\displaystyle{\bm{A}}_{\theta}$ at the operating point, the system $\displaystyle\dot{{\bm{z}}}={\bm{A}}_{\theta}\,{\bm{z}}$ is a linear, time-invariant ODE with solution

$$ ${\bm{z}}(t^{(i)}+\Delta t)=\exp\big({\bm{A}}_{\theta}\Delta t\big)\,{\bm{z}}_{t}^{(i)}.$ $$

Predictions in the gene expression space are obtained by projecting back $\displaystyle{\bm{x}}(t+\Delta t)={\bm{V}}_{\text{PCA}}{\bm{z}}(t+\Delta t)$.

#### Parametrization of the Operator.

For more fine-grained control over the parametrization of ${\bm{A}}_{\theta}$, we let the MLP predict the matrix in an eigen-decomposed form ${\bm{A}}_{\theta}={\bm{P}}_{\theta}\,\text{diag}(\bm{\lambda}_{\theta})\,{\bm{P}}_{\theta}^{-1}$, which is also beneficial to compute the matrix exponential. To ensure invertibility of ${\bm{P}}_{\theta}$, we train with the additional regularizer $\mathcal{L}^{\text{inv}}(\theta)=1/(\text{det}({\bm{P}}_{\theta})+\epsilon)$, which is practical if the latent space is small. This also lets us introduce inductive bias by selectively fixing eigenvalues, for example to zero, if needed.

#### Optimization.

We train the MLP parameters $\theta$ by minimizing the Maximum Mean Discrepancy (MMD, Gretton et al. [16])(^1^11A full definition of the MMD is given in Appendix [B](https://arxiv.org/html/2510.02903v1#A2)) between the model-induced marginals $q_{t}^{\theta}$ and the empirical marginals, thereby fitting a mechanistic model whose dynamics align with the target marginals $p_{t}$ under a future discounting factor $\gamma$. All discrepancies are computed in latent space via the pullback kernel

$$ $k_{x}({\bm{x}},{\bm{x}}^{\prime}):=k_{z}\!\big({\bm{V}}_{\text{PCA}}^{\top}{\bm{x}},\;{\bm{V}}_{\text{PCA}}^{\top}{\bm{x}}^{\prime}\big),$ $$

so that $\mathrm{MMD}^{2}(q_{t}^{\theta},p_{t};k_{x})=\mathrm{MMD}^{2}(q_{t}^{\theta,z},p_{t}^{z};k_{z})$. Here, $p_{t}^{z}$ and $q_{t}^{\theta,z}$ denote the distributions of the gene expression marginals in the latent space. The MMD loss is:

$$ $\mathcal{L}^{\text{MMD}^{2}}(\theta)=\mathbb{E}_{t}\!\left[\sum_{t^{\prime}=t}^{t_{K}}\gamma^{t^{\prime}}\,\mathrm{MMD}^{2}\big(q_{t^{\prime}}^{\theta},p_{t^{\prime}};k_{x}\big)\right].$ $$

Following Tong et al. [50], we also regularize the kinetic energy to improve generalization:

$$ $\mathcal{L}^{\text{kin}}(\theta)=\mathbb{E}_{t,\,{\bm{z}}_{t}\sim q_{t}^{\theta}}\!\left[\|\dot{{\bm{z}}}_{t}\|^{2}\right]=\mathbb{E}_{t,\,{\bm{z}}_{t}\sim q_{t}^{\theta}}\!\left[\|{\bm{A}}_{\theta}({\bm{z}}_{t},t)\,{\bm{z}}_{t}\|^{2}\right],$ $$

which serves as a soft constraint encouraging trajectories close to optimal transport flows in the sense of the Benamou and Brenier [1] formulation. Our final loss then becomes:

$$ $\mathcal{L}^{\text{total}}(\theta)=\mathcal{L}^{\text{MMD}^{2}}(\theta)+\lambda_{\text{kin}}\mathcal{L}^{\text{kin}}(\theta)+\lambda_{\text{inv}}\mathcal{L}^{\text{inv}}(\theta).$ (1) $$

#### Computational Complexity.

With ${\bm{A}}_{\theta}$ given in eigendecomposed form at an operating point, evaluating the analytical solution (Eq. [2.1](https://arxiv.org/html/2510.02903v1#S2.Ex3)) at $T$ time points has time complexity $\mathcal{O}(T\,d_{z}^{2})$ and space complexity $\mathcal{O}(d_{z}^{2})$, where $d_{z}$ is the latent space dimensionality. This improves the time and space complexity over the Scalable Mechanistic Neural Network (S-MNN) [7]. Forming the full operator requires computing ${\bm{P}}_{\theta}^{-1}$, incurring a one-time $\mathcal{O}(d_{z}^{3})$ cost per operating point.

#### Limitations.

The cubic time complexity in the latent dimension can become a challenge for high-dimensional latent spaces but could be mitigated by imposing sparsity assumptions on ${\bm{A}}_{\theta}$. In our application to single-cell dynamics, we follow the practice [50, 52] of using a 5-dimensional PCA space, which we find expressive enough to capture meaningful gene interactions in the high-dimensional gene expression space as presented later in the paper. Note that OT preprocessing on two time points, when using the Sinkhorn algorithm, scales as $\mathcal{O}(d_{z}\,n^{2})$ with the number of samples $n$, which becomes a bottleneck for large datasets, as $n$ is usually much larger than $d_{z}$. However, approximate batch approaches are also possible to address this [52].
A separate limitation of predicting the local dynamics laws is that evolving the system too far causes it to leave the regime where the linear ODE is accurate, which would require a new forward pass through the encoder to update the ODE. In our experiments, however, we did not encounter this issue.

#### Uncovering Local Gene Regulatory Interactions.

Combining the linear projection to the PCA subspace $\displaystyle{\bm{z}}={\bm{V}}_{\text{PCA}}^{\top}{\bm{x}}$ with locally linear dynamics around an operating point $\displaystyle\dot{{\bm{z}}}={\bm{A}}_{\theta}\,{\bm{z}}$ enables projecting the predicted local dynamics back into the gene expression space with:

$$ $\frac{d}{dt}\,{\bm{z}}={\bm{A}}_{\theta}{\bm{z}}\;\;\Longleftrightarrow\;\;\frac{d}{dt}\big({\bm{V}}_{\text{PCA}}^{\top}{\bm{x}}\big)={\bm{A}}_{\theta}{\bm{V}}_{\text{PCA}}^{\top}{\bm{x}}\;\;\Longleftrightarrow\;\;\frac{d}{dt}\,{\bm{x}}={\bm{V}}_{\text{PCA}}{\bm{A}}_{\theta}{\bm{V}}_{\text{PCA}}^{\top}{\bm{x}}.$ $$

which gives direct access to an explicit form of the predicted local dynamics in the gene expression space, essentially uncovering the predicted local gene regulatory interactions. We interpret:

$$ $w_{j\rightarrow i}({\bm{x}},t):=\big[{\bm{V}}_{\text{PCA}}{\bm{A}}_{\theta}({\bm{x}},t){\bm{V}}_{\text{PCA}}^{\top}\big]_{i,j}\cdot{\bm{x}}_{j},$ $$

as the interaction weight of gene $j$ to gene $i$. It essentially represents the contribution of gene $j$’s expression to the time derivative of ${\bm{x}}_{i}$. This makes our proposed approach fully interpretable, as we can inspect the learned gene interactions directly.

## 3 Related Works

#### Single-cell Interpolation.

The single-cell trajectory inference problem, as formalized by Lavenant et al. [29], entails reconstructing continuous dynamics from snapshot data. Early work based on recurrent neural networks [19] was followed by NeuralODE-based methods [50, 51, 56, 27, 21], in which a neural network outputs the velocity field governing the dynamics. In contrast, Cell-MNN predicts an explicit local dynamics model, which not only facilitates the learning of gene interactions but also circumvents the need for numerical ODE solvers. A separate line of work avoids simulation by relying on OT preprocessing to approximate cell trajectories [45, 4], which were also used to train flow-matching models such as by [52, 24, 57, 54, 48]. However, solving the OT coupling with the Sinkhorn algorithm scales quadratically in the number of samples, creating a major bottleneck for large datasets, which is why Tong et al. [52] proposed batch-wise approximation. To address this scalability bottleneck, Cell-MNN is designed to eliminate OT preprocessing entirely. Furthermore, SOTA OT-based models such as OT-MFM and DeepRUOT rely on multiple training stages beyond a standard PCA dimensionality reduction, which complicates amortized training across datasets. In contrast, Cell-MNN involves only a single training stage while achieving competitive performance on single-cell benchmarks. Finally, Action Matching [38] also avoids OT preprocessing, but unlike Cell-MNN, it does not learn an explicit form of the underlying dynamics.

#### Gene Regulatory Network Discovery.

A complementary line of work assumes that the interactions governing cell differentiation can be represented as a graph, known as a gene regulatory network (GRN) [13]. Tong et al. [53] demonstrated that such GRNs can to some extent be recovered from flow-matching models in the setting of low-dimensional synthetic data as simulated by Pratapa et al. [42]. In contrast, we show that Cell-MNN learns biologically plausible gene interactions directly from real single-cell data, validating them against the literature-curated TRRUST database. Additional approaches for GRN discovery include tree-based methods [22, 36], information-theoretic approaches [6], regression-based time-series models [34], Gaussian processes [59] and ODE-based models such as PerturbODE [32]. However, unlike Cell-MNN, these methods typically learn a static GRN and, to our knowledge, they are either inapplicable to single-cell interpolation benchmarks or do not deliver competitive performance.

#### ODE Discovery.

The idea to learn an explicit ODE representation of the cell differentiation dynamics as pursued by PerturbODE [32] and Cell-MNN relates directly to the broader problem of ODE discovery. A seminal method in this area is SINDy [3], which infers governing equations from data but requires access to full trajectories, making it unsuitable for the snapshot-based single-cell setting. Similar limitations apply to more recent approaches such as MNN and related methods such as ODEFormer [41, 7, 55, 12], which extend ODE discovery to amortized settings by using neural networks to predict the underlying dynamics from observed trajectories. In contrast, Cell-MNN is qualitatively distinct in learning dynamics from population data. It furthermore learns them in locally linear form, an idea with strong precedents in physics and control theory such as the Apollo navigation filter [46], the control of a 2-link 6-muscle arm model [30, 49] and rocket landing [47]. The locally linear parameterization theoretically imposes learning control-oriented structure and, in principle, supports the design of performant controllers as described by [43], which could enable design of gene perturbations.

## 4 Experiments

In the following, we present four experiments to evaluate Cell-MNN in terms of predictive accuracy, suitability for amortized training, scalability, and assessment of the predicted gene interactions.

#### Datasets.

For our experiments, we use 3 commonly studied real single-cell datasets. Following Tong et al. [50], we include the Embryoid Body (EB) dataset from Moon et al. [37], which after preprocessing contains $\sim$16 K human embryoid cells measured at five time points over 25 days. For EB, we model the time grid with $\mathcal{T}=\{0,1,...4\}$. We also use the CITE-seq (Cite) and Multiome (Multi) datasets from Burkhardt et al. [5], as repurposed by Tong et al. [52]. Both consist of gene expression measurements at four time points of cells developing over seven days, with Cite containing $\sim$31 K cells and Multi $\sim$33 K cells after preprocessing. Here we model the time grid with the days of measurement, namely $\mathcal{T}=\{0,1,2,3,7\}$. We use the datasets as preprocessed by Tong et al. [50, 52], which involves filtering for outliers and normalizing the data.

#### Training.

We use the same hyperparameters for all experiments unless stated otherwise. Following Tong et al. [50], we project gene expression to 5D PCA before training. The MLP used to parameterize ${\bm{A}}_{\theta}$ has depth $4$, width $96$, leaky ReLU activations, and Kaiming normal initialization [20]. For stability, we scale the MLP’s last layer by $0.01$ at initialization so that predictions of ${\bm{A}}_{\theta}$ start near zero.
For the MMD, we use the Laplacian kernel $k(z,z^{\prime})=\text{exp}[-\frac{\text{max}(||z-z^{\prime}||_{1},\epsilon)}{\sigma\cdot d_{z}}]$ with parameters $\sigma=1$ and $\epsilon=10^{-8}$. We optimize the final loss (Eq. [1](https://arxiv.org/html/2510.02903v1#S2.E1)) with a batch size per time point of $200$, future discount factor $\gamma=0.1$, initialization scale $0.01$, and regularization weights $\lambda_{\mathrm{kin}}=0.1$ and $\lambda_{\mathrm{inv}}=1$.
Optimization is performed using AdamW [26, 33] with a learning rate of $2\times 10^{-4}$ and weight decay $1\times 10^{-5}$.
Hyperparameters are selected according to grid search and all experiments are run with three random seeds. We validate every $10$ steps, with a patience of $40$ validation checks and a maximum training time of $200$ minutes. All training runs are performed with one NVIDIA GeForce RTX 2080 Ti per model (11 GB of RAM).

### 4.1 Single Cell Interpolation

**Table 1: Model comparison for single-cell interpolation across the Cite, EB, and Multi datasets, sorted by best average performance. We report the mean $\pm$ standard deviation of the EMD metric, along with the average across datasets. Standard deviation is computed over left-out time points. Lower values indicate better performance. Values marked * are computed by us.**
| Method | Cite | EB | Multi | Average $\downarrow$ |
| --- | --- | --- | --- | --- |
| TrajectoryNet [50] | – | 0.848 | – | – |
| WLF-UOT [39] | – | 0.800  $\pm$ 0.002 | – | – |
| NLSB [27] | – | 0.777  $\pm$ 0.021 | – | – |
| SB-CFM [52] | 1.067  $\pm$ 0.107 | 1.221  $\pm$ 0.380 | 1.129  $\pm$ 0.363 | 1.139  $\pm$ 0.077 |
| $[\text{SF}]^{2}$M-Sink [53] | 1.054  $\pm$ 0.087 | 1.198  $\pm$ 0.342 | 1.098  $\pm$ 0.308 | 1.117  $\pm$ 0.074 |
| $[\text{SF}]^{2}$M-Geo [53] | 1.017  $\pm$ 0.104 | 0.879  $\pm$ 0.148 | 1.255  $\pm$ 0.179 | 1.050  $\pm$ 0.190 |
| I-CFM [52] | 0.965  $\pm$ 0.111 | 0,872  $\pm$ 0.087 | 1.085  $\pm$ 0.099 | 0.974  $\pm$ 0.107 |
| DSB [14] | 0.965  $\pm$ 0.111 | 0.862  $\pm$ 0.023 | 1.079  $\pm$ 0.117 | 0.969  $\pm$ 0.109 |
| I-MFM [24] | 0.916  $\pm$ 0.124 | 0.822  $\pm$ 0.042 | 1.053  $\pm$ 0.095 | 0.930  $\pm$ 0.116 |
| $[\text{SF}]^{2}$M-Exact [53] | 0.920  $\pm$ 0.049 | 0.793  $\pm$ 0.066 | 0.933  $\pm$ 0.054 | 0.882  $\pm$ 0.077 |
| OT-CFM [52] | 0.882  $\pm$ 0.058 | 0.790  $\pm$ 0.068 | 0.937  $\pm$ 0.054 | 0.870  $\pm$ 0.074 |
| DeepRUOT [57]* | 0.845  $\pm$ 0.167 | 0.776  $\pm$ 0.079 | 0.919  $\pm$ 0.090 | 0.846  $\pm$ 0.071 |
| OT-Interpolate* | 0.821  $\pm$ 0.004 | 0.749  $\pm$ 0.019 | 0.830  $\pm$ 0.053 | 0.800  $\pm$ 0.044 |
| OT-MFM [24] | 0.724  $\pm$ 0.070 | 0.713  $\pm$ 0.039 | 0.890  $\pm$ 0.123 | 0.776  $\pm$ 0.099 |
| Cell-MNN (ours)* | 0.791  $\pm$ 0.022 | 0.690  $\pm$ 0.073 | 0.742  $\pm$ 0.100 | 0.741  $\pm$ 0.050 |

Following Tong et al. [50, 52], we evaluate model performance by measuring how closely it reproduces the marginal distribution of a held-out time point. Each intermediate day is left out in cross-validation fashion to obtain one comprehensive score per dataset.

#### Metric.

For easy comparison with SOTA methods, we follow [50] and report results in terms of the 1-Wasserstein distance in the PCA subspace ($W_{1}$ or otherwise EMD). We use the exact linear programming EMD from the POT (Python Optimal Transport) package [15]. The EMD metric measures the minimum cost of transporting probability mass to transform one distribution into another, where a lower score represents a closer match of distributions.

#### Baselines.

We compare with the 3 SOTA methods for this task, namely OT-MFM [24], OT-CFM [52] and DeepRUOT [57]. We found the pre-processing of DeepRUOT to be different from the other approaches, which is why we reran the experiments for DeepRUOT with exactly the same input data as the other methods. We also report the performance of other relevant previous works on this problem as indicated in the results Table [1](https://arxiv.org/html/2510.02903v1#S4.T1).
As an intuitive bar to cross, we additionally compute the performance of solely interpolating the optimal transport map between two consecutive time points and refer to it as OT-Interpolate.

#### Results.

Table [1](https://arxiv.org/html/2510.02903v1#S4.T1) summarizes the results on all three datasets. Cell-MNN achieves the best performance on EB and Multi, and ranks second on Cite, leading to the highest average performance across datasets. Notably, Cell-MNN is the only method that outperforms our proposed OT-Interpolate benchmark on all datasets. We think this is an important additional result, because any method that trains on velocities that are derived from the OT map implicitly treats OT-Interpolate as the ground truth. This also explains the strong performance of OT-Interpolate. Given the above, Cell-MNN delivers highly competitive predictive performance for single-cell interpolation.

### 4.2 Amortized Training

Figure: (a) Scaling experiment

Foundation models have shown strong transfer learning capabilities across datasets in a variety of domains [2, 40, 2]. However, current SOTA methods for single-cell interpolation, such as OT-MFM and DeepRUOT, rely on multi-stage training or dataset-specific regularizers, making them less suitable for building foundation models. In contrast, the end-to-end nature of Cell-MNN enables amortized training across multiple datasets. We design an experiment to assess which models are promising for amortized training in the single-cell interpolation setting by jointly training on datasets with the same time scale, namely Cite and Multi.

#### Training.

Our amortized training setup follows the single-cell interpolation experiment described in Section [4.1](https://arxiv.org/html/2510.02903v1#S4.SS1), with the only differences being that (i) we iteratively sample batches from Cite and Multi, (ii) we use a wider network with width $128$, and (iii) we pass an additional dataset index into the model. We do not sample from the marginals at the left-out time point for either dataset. Since each dataset contains a different set of genes, we use the same PCA embeddings as in the previous experiment (Section [4.1](https://arxiv.org/html/2510.02903v1#S4.SS1)) and merge datasets in the PCA subspace.

#### Baselines.

We use OT-CFM as a baseline as it is the best-performing alternative model on the single-cell interpolation task that involves only a single training stage, making it easy to adapt to the amortized training setting. For each dataset, we compute the OT map on the entire dataset separately to ensure that the derived velocity labels are accurate. We use the hyperparameters specified by Tong et al. [52] and first reproduce the original results for the separate-dataset setting to verify our setup. In amortized training, we find that passing the dataset index as input does not affect OT-CFM’s performance. For additional reference, we also report the performance of I-CFM.

#### Results.

As shown in Figure [3(b)](https://arxiv.org/html/2510.02903v1#S4.F3.sf2), Cell-MNN outperforms both OT-CFM and I-CFM in the amortized setting and achieves performance comparable to training on each dataset separately. Since the gene sets differ between datasets, transfer learning may be difficult in this setup. Nevertheless, these results suggest that for datasets with shared structure, Cell-MNN could enable transfer learning.

### 4.3 Scalability and Robustness to Noise

Beyond leveraging multiple datasets to train a single model, the practical usefulness of a method depends on its ability to handle the increasingly large datasets available in the single-cell dynamics domain. In this context, performing OT preprocessing over all samples from two consecutive days, as required by OT-CFM, DeepRUOT, or OT-Interpolate, can become a significant bottleneck due to the quadratic time and space complexity of the Sinkhorn algorithm [11]. To experimentally compare the scalability of different methods, we conduct the following scaling experiment.

#### Training.

We synthetically inflate the dataset size of EB, Cite, and Multi to $250{,}000$ cells each by resampling from each dataset and adding noise drawn from $\mathcal{N}(0,0.1)$ to the PCA embeddings. Cell-MNN is trained with the same hyperparameters as in our first experiment in Section [4.1](https://arxiv.org/html/2510.02903v1#S4.SS1).

#### Baselines.

We run OT-CFM and DeepRUOT on the inflated datasets and observe that both methods encounter out-of-memory (OOM) errors on our hardware (NVIDIA GeForce RTX 2080 Ti per model, 11 GB RAM) due to the quadratic memory complexity of the Sinkhorn algorithm. To mitigate this, Tong et al. [52] proposed a minibatch variant of optimal transport for OT-CFM, which achieves competitive image generation quality compared to dataset-wide OT preprocessing. We therefore use this mini-batch version as a baseline, denoted Batch-OT-CFM. As I-CFM does not require OT preprocessing, we also train it on the inflated datasets and report its performance.

#### Metric.

For the larger datasets, computing the EMD metric becomes impractical, as it also requires estimating the OT map. We therefore use the MMD metric with a Laplacian kernel to compute the validation score. Since MMD can be computed in a batch-wise fashion, it is more practical for this experiment. We use the same hyperparameters for the MMD metric as for the training loss.

#### Results.

We present the validation scores for models trained on inflated datasets in Table [3(a)](https://arxiv.org/html/2510.02903v1#S4.F3.sf1). Performing dataset-wide OT preprocessing, as required by standard OT-CFM or DeepRUOT, leads to OOM errors due to the quadratic space complexity of the Sinkhorn algorithm. Batch-OT-CFM outperforms I-CFM, demonstrating the gains from minibatch OT. However, Cell-MNN achieves the best performance on all three inflated datasets, highlighting its scalability and robustness to noise.

### 4.4 Discovering Gene Interactions

Figure: (a) Gene interactions
Refer to caption: x4.png

Figure: 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.
Refer to caption: x7.png

Cell-MNN predicts the local dynamics of the cell differentiation process in gene expression space, thereby learning interaction weights $w_{j\rightarrow i}$ from each gene to every other gene as described in Section [2.1](https://arxiv.org/html/2510.02903v1#S2.SS1.SSS0.Px7). This corresponds to unsupervised learning of local gene interactions.
To assess whether these learned interactions are biologically meaningful, we validate them against the literature-curated TRRUST database [18], which contains 8,444 regulatory relationships synthesized from 11,237 PubMed articles. While TRRUST represents only a small subset of all potential relationships, it provides a valuable reference signal for evaluating Cell-MNN’s predictions.

#### Unsupervised Classification.

Using labels from the TRRUST database, we design an unsupervised classification task for a source gene $j$: for every interaction $j\rightarrow i$ listed in TRRUST, we predict whether the relationship is activating or repressing. Since Cell-MNN outputs $w_{j\rightarrow i}$ per cell, we average these values over the dataset to obtain a single prediction for each interaction, classifying it as activating if $\sum_{({\bm{x}},t)\in\mathcal{D}}w_{j\rightarrow i}({\bm{x}},t)>0$ and repressing if smaller than zero. We report results for the genes $j$ that are in the top 10 most active ones (as by summing over all interaction weights per gene) for any time point and that have more than 10 matching TRRUST interactions within the EB gene set.

#### Training and Inductive Bias.

We train an ensemble of Cell-MNN models for single-cell interpolation on EB, each with a different left-out marginal. The setup matches Section [4.1](https://arxiv.org/html/2510.02903v1#S4.SS1) but uses a less preprocessed EB version, retaining gene names for downstream interaction analysis. To highlight the benefit of having access to an explicit dynamics model, we also introduce an additional model version with inductive bias on the parametrization of the linear operator ${\bm{A}}_{\theta}$: in particular, we know that only some genes vary over time, implying that at least one eigenvalue of ${\bm{A}}_{\theta}$ should be zero. During training, we therefore fix one eigenvalue to zero, forcing the model to learn static directions in the gene expression space. While this inductive bias reduces predictive performance in the single-cell interpolation setting by approximately $1\%$, it significantly improves gene interaction discovery performance on TRRUST (see Appendix [7](https://arxiv.org/html/2510.02903v1#A5.T7) for ablation results).

#### Results.

We compute precision, recall, and F1 scores for the unsupervised classification task, with results presented in Figure [4](https://arxiv.org/html/2510.02903v1#S4.F4) with all numerical values in Table [4](https://arxiv.org/html/2510.02903v1#A5.T4) and [5](https://arxiv.org/html/2510.02903v1#A5.T5). For all but two of the source genes, the vanilla Cell-MNN achieves better-than-average performance, indicating that for the tested gene regulatory interactions, the model can meaningfully discover activation or repression in a fully unsupervised manner. Interestingly, the Cell-MNN variant with one eigenvalue set to zero significantly improves classification performance, demonstrating that the inductive bias introduced on the operator effectively constrains the solution space. Note that gene interactions are context-dependent (e.g., varying by cell type or other factors), and therefore the labels in TRRUST may not fully apply to the context of the EB dataset. Nevertheless, we view the agreement for the most dominant source genes as a meaningful signal that the mechanisms learned by Cell-MNN are biologically plausible.

#### Visualizing Operators.

To visualize the learned operators, we plot UMAP projections of ${\bm{A}}_{\theta}$ computed jointly across all time ranges (Figure [4(b)](https://arxiv.org/html/2510.02903v1#S4.F4.sf2)) and separately for each time range (Figure [5](https://arxiv.org/html/2510.02903v1#S4.F5)) in the EB dataset. The joint projection shows that Cell-MNN captures distinct dynamics across time, while the separate projections highlight differences between cell types. Additional UMAP visualizations for the cell types reported in the original EB study [37] are provided in Appendix [G](https://arxiv.org/html/2510.02903v1#A7).

## 5 Conclusion

We introduced Cell-MNN, an encoder-decoder architecture whose representation is a locally linear latent ODE at the operating point of the cell differentiation dynamics. The formulation explicitly captures gene interactions conditioned on the time and gene expression. Empirically, we show that Cell-MNN achieves competitive performance on single-cell benchmarks, as well as in scaling and amortization experiments. Importantly, the gene interactions learned from real single-cell data exhibit consistency with the literature-curated TRRUST database. Thus, Cell-MNN jointly addresses the challenges of trajectory reconstruction from snapshot data and gene interaction discovery.

Having shown that Cell-MNN learns biologically plausible gene interactions, a natural next step is to use it as a hypothesis generation engine for less-studied genes, guiding which interactions to test experimentally. Moreover, since Cell-MNN models dynamics in locally linear form, it may be possible to leverage the rich control theory literature on controller design for such locally linear systems. In principle, this could enable steering gene expression states toward desired configurations via perturbations, which could for example inform CRISPR-based gene edits [23].

## Acknowledgements

This work was supported by the Chan Zuckerberg Initiative
(CZI) through the External Residency Program. We thank
CZI for the opportunity to participate in this program, which
enabled close collaboration and access to critical resources.
We also thank Alexander Tong and Lazar Atanackovic for sharing code and technical clarifications, which were essential for reproducing their results and conducting our subsequent experiments. Similarly, we thank the first authors of DeepRUOT [57] and VGFM [54], Dongyi Wang and Zhenyi Zhang, for providing their code and assistance with its use. Finally, we thank the Causal Learning and Artificial Intelligence group at ISTA for their valuable feedback and discussions throughout the project.

## Ethics Statement

All datasets used in this work (EB, Cite, Multi) are publicly available and were preprocessed by prior works Tong et al. [52, 50].
While the gene regulatory interactions predicted by Cell-MNN are partially validated against the TRRUST database, they should not be interpreted as definitive biological ground truth without further experimental validation. It is important to note that the model provides hypotheses for experimental follow-up, not direct medical recommendations.
Insights from the model could eventually inform gene perturbation studies or therapeutic research. To mitigate potential misuse, we emphasize that the work is intended for advancing computational methodology in machine learning and computational biology, not for direct clinical application.

## References

- Benamou and Brenier [2000]
Jean-David Benamou and Yann Brenier.
A computational fluid mechanics solution to the monge-kantorovich mass transfer problem.
*Numerische Mathematik*, 84(3):375–393, 2000.
doi: 10.1007/s002110050002.
URL [https://doi.org/10.1007/s002110050002](https://doi.org/10.1007/s002110050002).
- Bodnar et al. [2025]
Cristian Bodnar, Wessel P. Bruinsma, Ana Lucic, Megan Stanley, Anna Allen, Johannes Brandstetter, Patrick Garvan, Maik Riechert, Jonathan A. Weyn, Haiyu Dong, Jayesh K. Gupta, Kit Thambiratnam, Alexander T. Archibald, Chun-Chieh Wu, Elizabeth Heider, Max Welling, Richard E. Turner, and Paris Perdikaris.
A foundation model for the earth system.
*Nature*, 641(8065):1180–1187, 2025.
ISSN 1476-4687.
doi: 10.1038/s41586-025-09005-y.
URL [https://doi.org/10.1038/s41586-025-09005-y](https://doi.org/10.1038/s41586-025-09005-y).
- Brunton et al. [2016]
Steven L. Brunton, Joshua L. Proctor, and J. Nathan Kutz.
Discovering governing equations from data by sparse identification of nonlinear dynamical systems.
*Proceedings of the National Academy of Sciences*, 113(15):3932–3937, 2016.
doi: 10.1073/pnas.1517384113.
URL [https://www.pnas.org/doi/abs/10.1073/pnas.1517384113](https://www.pnas.org/doi/abs/10.1073/pnas.1517384113).
- Bunne et al. [2021]
Charlotte Bunne, Laetitia Meng-Papaxanthos, Andreas Krause, and Marco Cuturi.
Jkonet: Proximal optimal transport modeling of population dynamics.
*CoRR*, abs/2106.06345, 2021.
URL [https://arxiv.org/abs/2106.06345](https://arxiv.org/abs/2106.06345).
- Burkhardt et al. [2022]
Daniel Burkhardt, Malte Luecken, Andrew Benz, Peter Holderrieth, Jonathan Bloom, Christopher Lance, Ashley Chow, and Ryan Holbrook.
Open problems - multimodal single-cell integration.
[https://kaggle.com/competitions/open-problems-multimodal](https://kaggle.com/competitions/open-problems-multimodal), 2022.
Kaggle.
- Chan et al. [2017]
Thalia E. Chan, Michael P.H. Stumpf, and Ann C. Babtie.
Gene regulatory network inference from single-cell data using multivariate information measures.
*Cell Systems*, 5(3):251–267.e3, 2017.
ISSN 2405-4712.
doi: https://doi.org/10.1016/j.cels.2017.08.014.
URL [https://www.sciencedirect.com/science/article/pii/S2405471217303861](https://www.sciencedirect.com/science/article/pii/S2405471217303861).
- Chen et al. [2025]
Jiale Chen, Dingling Yao, Adeel Pervez, Dan Alistarh, and Francesco Locatello.
Scalable mechanistic neural networks.
In *The Thirteenth International Conference on Learning Representations*, 2025.
URL [https://openreview.net/forum?id=Oazgf8A24z](https://openreview.net/forum?id=Oazgf8A24z).
- Chen et al. [2018]
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud.
Neural ordinary differential equations.
In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, *Advances in Neural Information Processing Systems*, volume 31. Curran Associates, Inc., 2018.
URL [https://proceedings.neurips.cc/paper_files/paper/2018/file/69386f6bb1dfed68692a24c8686939b9-Paper.pdf](https://proceedings.neurips.cc/paper_files/paper/2018/file/69386f6bb1dfed68692a24c8686939b9-Paper.pdf).
- Chu et al. [2024]
Xianjing Chu, Wentao Tian, Jiaoyang Ning, Gang Xiao, Yunqi Zhou, Ziqi Wang, Zhuofan Zhai, Guilong Tanzhu, Jie Yang, and Rongrong Zhou.
Cancer stem cells: advances in knowledge and implications for cancer therapy.
*Signal Transduction and Targeted Therapy*, 9(1):170, 2024.
ISSN 2059-3635.
doi: 10.1038/s41392-024-01851-y.
URL [https://doi.org/10.1038/s41392-024-01851-y](https://doi.org/10.1038/s41392-024-01851-y).
- Cuomo et al. [2023]
Anna S. E. Cuomo, Aparna Nathan, Soumya Raychaudhuri, Daniel G. MacArthur, and Joseph E. Powell.
Single-cell genomics meets human genetics.
*Nature Reviews Genetics*, 24(8):535–549, August 2023.
ISSN 1471-0064.
doi: 10.1038/s41576-023-00599-5.
URL [https://doi.org/10.1038/s41576-023-00599-5](https://doi.org/10.1038/s41576-023-00599-5).
- Cuturi [2013]
Marco Cuturi.
Sinkhorn distances: Lightspeed computation of optimal transport.
In C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q. Weinberger, editors, *Advances in Neural Information Processing Systems*, volume 26. Curran Associates, Inc., 2013.
URL [https://proceedings.neurips.cc/paper_files/paper/2013/file/af21d0c97db2e27e13572cbf59eb343d-Paper.pdf](https://proceedings.neurips.cc/paper_files/paper/2013/file/af21d0c97db2e27e13572cbf59eb343d-Paper.pdf).
- d’Ascoli et al. [2024]
Stéphane d’Ascoli, Sören Becker, Philippe Schwaller, Alexander Mathis, and Niki Kilbertus.
ODEFormer: Symbolic regression of dynamical systems with transformers.
In *The Twelfth International Conference on Learning Representations*, 2024.
URL [https://openreview.net/forum?id=TzoHLiGVMo](https://openreview.net/forum?id=TzoHLiGVMo).
- Davidson et al. [2002]
Eric H. Davidson, Jonathan P. Rast, Paola Oliveri, Andrew Ransick, Cristina Calestani, Chiou-Hwa Yuh, Takuya Minokawa, Gabriele Amore, Veronica Hinman, Cesar Arenas-Mena, Ochan Otim, C. Titus Brown, Carolina B. Livi, Pei Yun Lee, Roger Revilla, Alistair G. Rust, Zheng Jun Pan, Maria J. Schilstra, Peter J. C. Clarke, Maria I. Arnone, Lee Rowen, R. Andrew Cameron, David R. McClay, Leroy Hood, and Hamid Bolouri.
A genomic regulatory network for development.
*Science*, 295(5560):1669–1678, 2002.
ISSN 0036-8075, 1095-9203.
doi: 10.1126/science.1069883.
- De Bortoli et al. [2021]
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet.
Diffusion schrödinger bridge with applications to score-based generative modeling.
In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, *Advances in Neural Information Processing Systems*, volume 34, pages 17695–17709. Curran Associates, Inc., 2021.
URL [https://proceedings.neurips.cc/paper_files/paper/2021/file/940392f5f32a7ade1cc201767cf83e31-Paper.pdf](https://proceedings.neurips.cc/paper_files/paper/2021/file/940392f5f32a7ade1cc201767cf83e31-Paper.pdf).
- Flamary et al. [2021]
RÃ©mi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, AurÃ©lie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, LÃ©o Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer.
Pot: Python optimal transport.
*Journal of Machine Learning Research*, 22(78):1–8, 2021.
URL [http://jmlr.org/papers/v22/20-451.html](http://jmlr.org/papers/v22/20-451.html).
- Gretton et al. [2012]
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola.
A kernel two-sample test.
*Journal of Machine Learning Research*, 13(25):723–773, 2012.
URL [http://jmlr.org/papers/v13/gretton12a.html](http://jmlr.org/papers/v13/gretton12a.html).
- Ha et al. [2017]
David Ha, Andrew M. Dai, and Quoc V. Le.
Hypernetworks.
In *International Conference on Learning Representations*, 2017.
URL [https://openreview.net/forum?id=rkpACe1lx](https://openreview.net/forum?id=rkpACe1lx).
- Han et al. [2018]
Heonjong Han, Jae-Won Cho, Sangyoung Lee, Ayoung Yun, Hyojin Kim, Dasom Bae, Sunmo Yang, Chan Yeong Kim, Muyoung Lee, Eunbeen Kim, Sungho Lee, Byunghee Kang, Dabin Jeong, Yaeji Kim, Hyeon-Nae Jeon, Haein Jung, Sunhwee Nam, Michael Chung, Jong-Hoon Kim, and Insuk Lee.
TRRUST v2: an expanded reference database of human and mouse transcriptional regulatory interactions.
*Nucleic Acids Research*, 46(D1):D380–D386, January 2018.
ISSN 1362-4962.
doi: 10.1093/nar/gkx1013.
URL [https://doi.org/10.1093/nar/gkx1013](https://doi.org/10.1093/nar/gkx1013).
- Hashimoto et al. [2016]
Tatsunori Hashimoto, David Gifford, and Tommi Jaakkola.
Learning population-level diffusions with generative rnns.
In Maria Florina Balcan and Kilian Q. Weinberger, editors, *Proceedings of The 33rd International Conference on Machine Learning*, volume 48 of *Proceedings of Machine Learning Research*, pages 2417–2426, New York, New York, USA, 20–22 Jun 2016. PMLR.
URL [https://proceedings.mlr.press/v48/hashimoto16.html](https://proceedings.mlr.press/v48/hashimoto16.html).
- He et al. [2015]
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.
Delving deep into rectifiers: Surpassing human-level performance on imagenet classification.
In *Proceedings of the IEEE International Conference on Computer Vision (ICCV)*, December 2015.
- Huguet et al. [2022]
Guillaume Huguet, D. S. Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy.
Manifold interpolating optimal-transport flows for trajectory inference, 2022.
URL [https://arxiv.org/abs/2206.14928](https://arxiv.org/abs/2206.14928).
- Huynh-Thu et al. [2010]
Vân Anh Huynh-Thu, Alexandre Irrthum, Louis Wehenkel, and Pierre Geurts.
Inferring regulatory networks from expression data using tree-based methods.
*PLOS ONE*, 5(9):1–10, 09 2010.
doi: 10.1371/journal.pone.0012776.
URL [https://doi.org/10.1371/journal.pone.0012776](https://doi.org/10.1371/journal.pone.0012776).
- Jinek et al. [2012]
Martin Jinek, Krzysztof Chylinski, Ines Fonfara, Michael Hauer, Jennifer A. Doudna, and Emmanuelle Charpentier.
A programmable dual-rna-guided dna endonuclease in adaptive bacterial immunity.
*Science*, 337(6096):816–821, 2012.
ISSN 0036-8075, 1095-9203.
doi: 10.1126/science.1225829.
- Kapusniak et al. [2024]
Kacper Kapusniak, Peter Potaptchik, Teodora Reu, Leo Zhang, Alexander Tong, Michael M. Bronstein, Joey Bose, and Francesco Di Giovanni.
Metric flow matching for smooth interpolations on the data manifold.
In *The Thirty-eighth Annual Conference on Neural Information Processing Systems*, 2024.
URL [https://openreview.net/forum?id=fE3RqiF4Nx](https://openreview.net/forum?id=fE3RqiF4Nx).
- Kharchenko [2021]
Peter V. Kharchenko.
The triumphs and limitations of computational methods for scrna-seq.
*Nature Methods*, 18(7):723–732, 2021.
ISSN 1548-7105.
doi: 10.1038/s41592-021-01171-x.
URL [https://doi.org/10.1038/s41592-021-01171-x](https://doi.org/10.1038/s41592-021-01171-x).
- Kingma and Ba [2017]
Diederik P. Kingma and Jimmy Ba.
Adam: A method for stochastic optimization, 2017.
URL [https://arxiv.org/abs/1412.6980](https://arxiv.org/abs/1412.6980).
- Koshizuka and Sato [2023]
Takeshi Koshizuka and Issei Sato.
Neural lagrangian schrödinger bridge: Diffusion modeling for population dynamics.
In *ICLR*, 2023.
URL [https://openreview.net/forum?id=d3QNWD_pcFv](https://openreview.net/forum?id=d3QNWD_pcFv).
- Kovachki et al. [2021]
Nikola B. Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew M. Stuart, and Anima Anandkumar.
Neural operator: Learning maps between function spaces.
*CoRR*, abs/2108.08481, 2021.
URL [https://arxiv.org/abs/2108.08481](https://arxiv.org/abs/2108.08481).
- Lavenant et al. [2023]
Hugo Lavenant, Stephen Zhang, Young-Heon Kim, and Geoffrey Schiebinger.
Towards a mathematical theory of trajectory inference, 2023.
URL [https://arxiv.org/abs/2102.09204](https://arxiv.org/abs/2102.09204).
- Li and Todorov [2004]
Weiwei Li and Emanuel Todorov.
Iterative linear quadratic regulator design for nonlinear biological movement systems.
In Helder Araújo, Alves Vieira, José Braz, Bruno Encarnação, and Marina Carvalho, editors, *ICINCO (1)*, pages 222–229. INSTICC Press, 2004.
ISBN 972-8865-12-0.
URL [http://dblp.uni-trier.de/db/conf/icinco/icinco2004.html#LiT04](http://dblp.uni-trier.de/db/conf/icinco/icinco2004.html#LiT04).
- Li et al. [2020]
Zongyi Li, Nikola B. Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew M. Stuart, and Anima Anandkumar.
Fourier neural operator for parametric partial differential equations.
*CoRR*, abs/2010.08895, 2020.
URL [https://arxiv.org/abs/2010.08895](https://arxiv.org/abs/2010.08895).
- Lin et al. [2025]
Zaikang Lin, Sei Chang, Aaron Zweig, Minseo Kang, Elham Azizi, and David A. Knowles.
Interpretable neural odes for gene regulatory network discovery under perturbations, 2025.
URL [https://arxiv.org/abs/2501.02409](https://arxiv.org/abs/2501.02409).
- Loshchilov and Hutter [2019]
Ilya Loshchilov and Frank Hutter.
Decoupled weight decay regularization.
In *International Conference on Learning Representations*, 2019.
URL [https://openreview.net/forum?id=Bkg6RiCqY7](https://openreview.net/forum?id=Bkg6RiCqY7).
- Lu et al. [2021]
Junjie Lu, Bianca Dumitrascu, Ian C. McDowell, Brian Jo, Alejandro Barrera, Lee K. Hong, Samuel M. Leichter, Timothy E. Reddy, and Barbara E. Engelhardt.
Causal network inference from gene transcriptional time-series response to glucocorticoids.
*PLoS Computational Biology*, 17(1):e1008223, 2021.
ISSN 1553-734X.
doi: 10.1371/journal.pcbi.1008223.
- Macosko et al. [2015]
Evan Z. Macosko, Anindita Basu, Rahul Satija, James Nemesh, Karthik Shekhar, Melissa Goldman, Itay Tirosh, Allison R. Bialas, Nolan Kamitaki, Emily M. Martersteck, John J. Trombetta, David A. Weitz, Joshua R. Sanes, Alex K. Shalek, Aviv Regev, and Steven A. McCarroll.
Highly parallel genome-wide expression profiling of individual cells using nanoliter droplets.
*Cell*, 161(5):1202–1214, 2015.
doi: 10.1016/j.cell.2015.05.002.
URL [https://doi.org/10.1016/j.cell.2015.05.002](https://doi.org/10.1016/j.cell.2015.05.002).
- Moerman et al. [2018]
Thomas Moerman, Sara Aibar Santos, Carmen Bravo González-Blas, Jaak Simm, Yves Moreau, Jan Aerts, and Stein Aerts.
Grnboost2 and arboreto: efficient and scalable inference of gene regulatory networks.
*Bioinformatics*, 35(12):2159–2161, 11 2018.
ISSN 1367-4803.
doi: 10.1093/bioinformatics/bty916.
URL [https://doi.org/10.1093/bioinformatics/bty916](https://doi.org/10.1093/bioinformatics/bty916).
- Moon et al. [2019]
Kevin R. Moon, David van Dijk, Zheng Wang, Scott Gigante, Daniel B. Burkhardt, William S. Chen, Kristina Yim, Antonia van den Elzen, Matthew J. Hirn, Ronald R. Coifman, Natalia B. Ivanova, Guy Wolf, and Smita Krishnaswamy.
Visualizing structure and transitions in high-dimensional biological data.
*Nature Biotechnology*, 37(12):1482–1492, 2019.
doi: 10.1038/s41587-019-0336-3.
- Neklyudov et al. [2023]
Kirill Neklyudov, Rob Brekelmans, Daniel Severo, and Alireza Makhzani.
Action matching: Learning stochastic dynamics from samples.
In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, *Proceedings of the 40th International Conference on Machine Learning*, volume 202 of *Proceedings of Machine Learning Research*, pages 25858–25889. PMLR, 23–29 Jul 2023.
URL [https://proceedings.mlr.press/v202/neklyudov23a.html](https://proceedings.mlr.press/v202/neklyudov23a.html).
- Neklyudov et al. [2024]
Kirill Neklyudov, Rob Brekelmans, Alexander Tong, Lazar Atanackovic, Qiang Liu, and Alireza Makhzani.
A computational framework for solving wasserstein lagrangian flows.
In *ICML*, 2024.
URL [https://openreview.net/forum?id=wwItuHdus6](https://openreview.net/forum?id=wwItuHdus6).
- Pearce et al. [2025]
James D Pearce, Sara E Simmonds, Gita Mahmoudabadi, Lakshmi Krishnan, Giovanni Palla, Ana-Maria Istrate, Alexander Tarashansky, Benjamin Nelson, Omar Valenzuela, Donghui Li, Stephen R Quake, and Theofanis Karaletsos.
A cross-species generative cell atlas across 1.5 billion years of evolution: The transcriptformer single-cell model.
*bioRxiv*, 2025.
doi: 10.1101/2025.04.25.650731.
URL [https://www.biorxiv.org/content/early/2025/04/29/2025.04.25.650731](https://www.biorxiv.org/content/early/2025/04/29/2025.04.25.650731).
- Pervez et al. [2024]
Adeel Pervez, Francesco Locatello, and Stratis Gavves.
Mechanistic neural networks for scientific machine learning.
In *Forty-first International Conference on Machine Learning*, 2024.
URL [https://openreview.net/forum?id=pLtuwhoQh7](https://openreview.net/forum?id=pLtuwhoQh7).
- Pratapa et al. [2020]
Aditya Pratapa, Amogh P. Jalihal, Jeffrey N. Law, Aditya Bharadwaj, and T. M. Murali.
Benchmarking algorithms for gene regulatory network inference from single-cell transcriptomic data.
*Nature Methods*, 17(2):147–154, 2020.
ISSN 1548-7105.
doi: 10.1038/s41592-019-0690-6.
URL [https://doi.org/10.1038/s41592-019-0690-6](https://doi.org/10.1038/s41592-019-0690-6).
- Richards et al. [2023]
Spencer M. Richards, Jean-Jacques Slotine, Navid Azizan, and Marco Pavone.
Learning control-oriented dynamical structure from data.
In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, *Proceedings of the 40th International Conference on Machine Learning*, volume 202 of *Proceedings of Machine Learning Research*, pages 29051–29062. PMLR, 23–29 Jul 2023.
URL [https://proceedings.mlr.press/v202/richards23a.html](https://proceedings.mlr.press/v202/richards23a.html).
- Rodrigues et al. [2019]
Melanie Rodrigues, Nina Kosaric, Clark A. Bonham, and Geoffrey C. Gurtner.
Wound healing: A cellular perspective.
*Physiological Reviews*, 99(1):665–706, 2019.
ISSN 1522-1210.
doi: 10.1152/physrev.00067.2017.
URL [https://doi.org/10.1152/physrev.00067.2017](https://doi.org/10.1152/physrev.00067.2017).
- Schiebinger et al. [2019]
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, Lia Lee, Jenny Chen, Justin Brumbaugh, Philippe Rigollet, Konrad Hochedlinger, Rudolf Jaenisch, Aviv Regev, and Eric S. Lander.
Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming.
*Cell*, 176(4):928–943.e22, 2019.
ISSN 0092-8674.
doi: https://doi.org/10.1016/j.cell.2019.01.006.
URL [https://www.sciencedirect.com/science/article/pii/S009286741930039X](https://www.sciencedirect.com/science/article/pii/S009286741930039X).
- Schmidt [1966]
STANLEY F. Schmidt.
Application of state-space methods to navigation problems.
volume 3 of *Advances in Control Systems*, pages 293–340. Elsevier, 1966.
doi: https://doi.org/10.1016/B978-1-4831-6716-9.50011-4.
URL [https://www.sciencedirect.com/science/article/pii/B9781483167169500114](https://www.sciencedirect.com/science/article/pii/B9781483167169500114).
- Szmuk et al. [2020]
Michael Szmuk, Taylor P. Reynolds, and Behçet Açıkmeşe.
Successive convexification for real-time six-degree-of-freedom powered descent guidance with state-triggered constraints.
*Journal of Guidance, Control, and Dynamics*, 43(8):1399–1413, 2020.
doi: 10.2514/1.G004549.
URL [https://doi.org/10.2514/1.G004549](https://doi.org/10.2514/1.G004549).
- Terpin et al. [2024]
Antonio Terpin, Nicolas Lanzetti, Martín Gadea, and Florian Dorfler.
Learning diffusion at lightspeed.
In *The Thirty-eighth Annual Conference on Neural Information Processing Systems*, 2024.
URL [https://openreview.net/forum?id=y10avdRFNK](https://openreview.net/forum?id=y10avdRFNK).
- Todorov and Li [2005]
E. Todorov and Weiwei Li.
A generalized iterative lqg method for locally-optimal feedback control of constrained nonlinear stochastic systems.
In *Proceedings of the 2005, American Control Conference, 2005.*, pages 300–306 vol. 1, 2005.
doi: 10.1109/ACC.2005.1469949.
- Tong et al. [2020]
Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy.
TrajectoryNet: A dynamic optimal transport network for modeling cellular dynamics.
In Hal Daumé III and Aarti Singh, editors, *Proceedings of the 37th International Conference on Machine Learning*, volume 119 of *Proceedings of Machine Learning Research*, pages 9526–9536. PMLR, 13–18 Jul 2020.
URL [https://proceedings.mlr.press/v119/tong20a.html](https://proceedings.mlr.press/v119/tong20a.html).
- Tong et al. [2023]
Alexander Tong, Manik Kuchroo, Shabarni Gupta, Aarthi Venkat, Beatriz P. San Juan, Laura Rangel, Brandon Zhu, John G. Lock, Christine L. Chaffer, and Smita Krishnaswamy.
Learning transcriptional and regulatory dynamics driving cancer cell plasticity using neural ode-based optimal transport.
*bioRxiv*, 2023.
doi: 10.1101/2023.03.28.534644.
URL [https://www.biorxiv.org/content/early/2023/03/29/2023.03.28.534644](https://www.biorxiv.org/content/early/2023/03/29/2023.03.28.534644).
- Tong et al. [2024a]
Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio.
Improving and generalizing flow-based generative models with minibatch optimal transport.
*Transactions on Machine Learning Research*, 2024a.
ISSN 2835-8856.
URL [https://openreview.net/forum?id=CD9Snc73AW](https://openreview.net/forum?id=CD9Snc73AW).
Expert Certification.
- Tong et al. [2024b]
Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio.
Simulation-free schrödinger bridges via score and flow matching.
In *AISTATS*, pages 1279–1287, 2024b.
URL [https://proceedings.mlr.press/v238/y-tong24a.html](https://proceedings.mlr.press/v238/y-tong24a.html).
- Wang et al. [2025]
Dongyi Wang, Yuanwei Jiang, Zhenyi Zhang, Xiang Gu, Peijie Zhou, and Jian Sun.
Joint velocity-growth flow matching for single-cell dynamics modeling, 2025.
URL [https://arxiv.org/abs/2505.13413](https://arxiv.org/abs/2505.13413).
- Yao et al. [2024]
Dingling Yao, Caroline Muller, and Francesco Locatello.
Marrying causal representation learning with dynamical systems for science.
In *The Thirty-eighth Annual Conference on Neural Information Processing Systems*, 2024.
URL [https://openreview.net/forum?id=MWHRxKz4mq](https://openreview.net/forum?id=MWHRxKz4mq).
- Zhang et al. [2023]
Jiaqi Zhang, Erica Larschan, Jeremy Bigness, and Ritambhara Singh.
scnode : Generative model for temporal single cell transcriptomic data prediction.
*bioRxiv*, 2023.
doi: 10.1101/2023.11.22.568346.
URL [https://www.biorxiv.org/content/early/2023/11/23/2023.11.22.568346](https://www.biorxiv.org/content/early/2023/11/23/2023.11.22.568346).
- Zhang et al. [2025]
Zhenyi Zhang, Tiejun Li, and Peijie Zhou.
Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport.
In *The Thirteenth International Conference on Learning Representations*, 2025.
URL [https://openreview.net/forum?id=gQlxd3Mtru](https://openreview.net/forum?id=gQlxd3Mtru).
- Zheng et al. [2017]
Grace X. Y. Zheng, Jessica M. Terry, Phillip Belgrader, Paul Ryvkin, Zachary W. Bent, Ryan Wilson, Solongo B. Ziraldo, Tobias D. Wheeler, Geoff P. McDermott, Junjie Zhu, Mark T. Gregory, Joe Shuga, Luz Montesclaros, Jason G. Underwood, Donald A. Masquelier, Stefanie Y. Nishimura, Michael Schnall-Levin, Paul W. Wyatt, Christopher M. Hindson, Rajiv Bharadwaj, Alexander Wong, Kevin D. Ness, Lan W. Beppu, H. Joachim Deeg, Christopher McFarland, Keith R. Loeb, William J. Valente, Nolan G. Ericson, Emily A. Stevens, Jerald P. Radich, Tarjei S. Mikkelsen, Benjamin J. Hindson, and Jason H. Bielas.
Massively parallel digital transcriptional profiling of single cells.
*Nature Communications*, 8(1):14049, 2017.
ISSN 2041-1723.
doi: 10.1038/ncomms14049.
URL [https://doi.org/10.1038/ncomms14049](https://doi.org/10.1038/ncomms14049).
- Äijö and Lähdesmäki [2009]
Tarmo Äijö and Harri Lähdesmäki.
Learning gene regulatory networks from gene expression measurements using non-parametric molecular kinetics.
*Bioinformatics*, 25(22):2937–2944, 08 2009.
ISSN 1367-4803.
doi: 10.1093/bioinformatics/btp511.
URL [https://doi.org/10.1093/bioinformatics/btp511](https://doi.org/10.1093/bioinformatics/btp511).
- Çimen [2010]
Tayfun Çimen.
Systematic and effective design of nonlinear feedback controllers via the state-dependent riccati equation (sdre) method.
*Annual Reviews in Control*, 34(1):32–51, 2010.
ISSN 1367-5788.
doi: https://doi.org/10.1016/j.arcontrol.2010.03.001.
URL [https://www.sciencedirect.com/science/article/pii/S1367578810000052](https://www.sciencedirect.com/science/article/pii/S1367578810000052).

## Appendix A Use of Large Language Models (LLMs)

In this work, we used LLMs for (i) coding assistance during the software development phase,
(ii) identifying relevant literature in response to specific research questions, and
(iii) polishing and improving the readability of the paper.
All substantive research contributions, analysis, and interpretations were carried out by the authors.

## Appendix B Definitions

Maximum Mean Discrepancy (MMD, Gretton et al. [16]): Given two distributions $p$ and $q$ over $\mathcal{X}$ and a positive-definite kernel function $k:\mathcal{X}\times\mathcal{X}\to\mathbb{R}$,
the squared Maximum Mean Discrepancy (MMD) is defined as

$$ $\mathrm{MMD}^{2}(p,q;k)=\mathbb{E}_{x,x^{\prime}\sim p}[k(x,x^{\prime})]+\mathbb{E}_{y,y^{\prime}\sim q}[k(y,y^{\prime})]-2\,\mathbb{E}_{x\sim p,y\sim q}[k(x,y)].$ $$

## Appendix C Architecture

Figure: Figure 6: The Cell-MNN architecture first applies the PCA projection matrix to map the gene expression state ${\bm{x}}$ to a latent representation ${\bm{z}}_{t}$. An MLP then predicts a locally linear approximation $\dot{{\bm{z}}}={\bm{A}}_{\theta}{\bm{z}}$ to the dynamics at the operating point $({\bm{z}}_{t},t)$. To decode, the analytical solution of this ODE is evaluated at a future time point and projected back into gene expression space.
Refer to caption: x8.png

**Table 2: Collection of training hyperparameters used in all experiments. All models were trained on a single NVIDIA RTX 2080 Ti GPU. Training on EB, Cite, and Multi separately, as well as on the amortization experiment, required about 1 hour per run, while training on inflated datasets took roughly 4 hours. No distributed training or large-scale compute resources were required.**
| Component | Hyperparameter |
| --- | --- |
| Data preprocessing | PCA projection to 5D |
| MLP (${\bm{A}}_{\theta}$) | Depth: 4 |
|  | Width: 96 (128 for amortization experiment) |
|  | Activation: Leaky ReLU |
|  | Initialization: Kaiming normal |
|  | Last layer scale: 0.01 |
| MMD kernel | Laplacian kernel $k(z,z^{\prime})=\exp[-\tfrac{\max(||z-z^{\prime}||_{1},\epsilon)}{\sigma d_{z}}]$ |
|  | $\sigma=1$, $\epsilon=10^{-8}$ |
| Optimization | Batch size per time point: 200 |
|  | Future discount factor $\gamma=0.1$ |
|  | Initialization scale: 0.01 |
|  | Regularization $\lambda_{\mathrm{kin}}=0.1$, $\lambda_{\mathrm{inv}}=1$ |
|  | Optimizer: AdamW |
|  | Learning rate: $2\times 10^{-4}$ |
|  | Weight decay: $1\times 10^{-5}$ |
| Validation | Frequency: every 10 steps |
|  | Patience: 40 validation checks |
| Training time | 60 minutes (240 minutes for inflated datasets) |
| Randomness | Seeds: 3 |
| Hardware | 1$\times$ NVIDIA GeForce RTX 2080 Ti (11 GB RAM) |

## Appendix D Proofs

###### Proposition 1 (Extension of Proposition 1 of Çimen [ 60 ] ) .

Let ${\bm{f}}:\mathbb{R}^{d_{z}}\times\mathbb{R}\to\mathbb{R}^{d_{z}}$ satisfy ${\bm{f}}(\mathbf{0},t)=\mathbf{0}$ for all $t\in\mathbb{R}$, and assume ${\bm{f}}\in\mathcal{C}^{k}(\mathbb{R}^{d_{z}}\times\mathbb{R})$ with $k\geq 1$.
Then there exists a matrix-valued map ${\bm{A}}:\mathbb{R}^{d_{z}}\times\mathbb{R}\to\mathbb{R}^{d_{z}\times d_{z}}$ such that ${\bm{f}}({\bm{z}},t)={\bm{A}}({\bm{z}},t)\,{\bm{z}}\,$ for all $({\bm{z}},t)\in\mathbb{R}^{d_{z}}\times\mathbb{R}$.

###### Proof.

Fix $({\bm{z}},t)\in\mathbb{R}^{d_{z}}\times\mathbb{R}$ and define $\gamma:[0,1]\to\mathbb{R}^{d_{z}}$ by
$\gamma(s):={\bm{f}}(s\cdot{\bm{z}},t)$. Since ${\bm{f}}\in\mathcal{C}^{k}$ and $k\geq 1$, the map
$s\mapsto\gamma(s)$ is differentiable and

$$ $\frac{d}{ds}\gamma(s)=D_{{\bm{z}}}{\bm{f}}(s\cdot{\bm{z}},t)\,{\bm{z}}.$ $$

By the fundamental theorem of calculus,

$$ ${\bm{f}}({\bm{z}},t)-{\bm{f}}(\mathbf{0},t)=\gamma(1)-\gamma(0)=\int_{0}^{1}D_{{\bm{z}}}{\bm{f}}(s\cdot{\bm{z}},t)\,{\bm{z}}\,ds=\Big(\int_{0}^{1}D_{{\bm{z}}}{\bm{f}}(s\cdot{\bm{z}},t)\,ds\Big){\bm{z}}.$ $$

Using ${\bm{f}}(\mathbf{0},t)=\mathbf{0}$ gives ${\bm{f}}({\bm{z}},t)={\bm{A}}({\bm{z}},t)\,{\bm{z}}$.
∎

###### Proposition 1 (Extension of Proposition 1 of Çimen [ 60 ] ) .

###### Proof.

## Appendix E Gene Regulatory Interaction Recovery

To quantitatively assess the learned gene interactions, we designed an unsupervised classification task based on the TRRUST database, which contains literature-curated gene regulatory interactions, many of which are annotated as activating or repressing. For evaluation of our model, we focus on the most dominant source genes predicted by Cell-MNN, i.e., those with the highest mean interaction strength with other genes. A source gene is included into the experiment if at least 10 of its interactions are listed in TRRUST. For each such gene, we classify the direction of its effect on downstream targets as activating or repressing. Since Cell-MNN produces cell-specific predictions of interaction weights, we average these over 10,000 cells to obtain a robust prediction for each interaction. Based on these predictions, we compute precision, recall, and F1 scores to quantify how well the model recovers known regulatory mechanisms and report them in Table [4](https://arxiv.org/html/2510.02903v1#A5.T4) and Table [5](https://arxiv.org/html/2510.02903v1#A5.T5).

**Table 3: Gene selection specifications for the TRRUST experiment. There are 16 genes that make up the top 10 predicted high-interaction source genes across the five time points, Of these, 14 are contained in TRRUST, and 6 have more than 10 interactions overlapping with the training gene set. This table provides shows how source genes were selected for downstream evaluation.**
| Top Source Gene | In<br>TRRUST | # Interactions<br>in TRRUST | # in Training<br>Gene Set | $>10$<br>Interactions |
| --- | --- | --- | --- | --- |
| HMGA1 | True | 18 | 10 | True |
| HMGB2 | True | 2 |  |  |
| JUNB | True | 15 | 4 |  |
| FOS | True | 63 | 25 | True |
| JUN | True | 173 | 65 | True |
| POU5F1 | True | 25 | 19 | True |
| HAND1 | False | 0 | 0 |  |
| ID2 | True | 2 |  |  |
| TERF1 | True | 1 |  |  |
| PITX2 | True | 11 | 4 |  |
| ID3 | True | 2 |  |  |
| HMGB1 | False | 0 | 0 |  |
| SOX2 | True | 23 | 16 | True |
| HMGA2 | True | 5 |  |  |
| YBX1 | True | 33 | 24 | True |
| ID1 | True | 1 |  |  |

**Table 4: Validation of predicted gene interactions on TRRUST: For each source gene $j$, we classify each TRRUST edge $j\!\rightarrow\!i$ as activating or repressing using the sign of the learned weight $w_{j\rightarrow i}$ from Cell-MNN with one eigenvalue set to zero (averaged over cells). For each source gene, we report the number of interactions in TRRUST and classification metrics (precision, recall, and F1) shown as mean $\pm$ std across ensemble models trained on three different seeds.**
| Source Gene | # Interactions $\downarrow$ | Precision | Recall | F1 |
| --- | --- | --- | --- | --- |
| JUN | 65 | $62\%\pm 8\%$ | $82\%\pm 3\%$ | $71\%\pm 6\%$ |
| FOS | 25 | $65\%\pm 10\%$ | $80\%\pm 10\%$ | $71\%\pm 10\%$ |
| YBX1 | 24 | $55\%\pm 10\%$ | $48\%\pm 3\%$ | $51\%\pm 6\%$ |
| POU5F1 | 19 | $82\%\pm 6\%$ | $58\%\pm 11\%$ | $67\%\pm 6\%$ |
| SOX2 | 16 | $73\%\pm 12\%$ | $69\%\pm 8\%$ | $71\%\pm 9\%$ |
| HMGA1 | 10 | $78\%\pm 9\%$ | $82\%\pm 2\%$ | $80\%\pm 6\%$ |

**Table 5: Ablation of models with no eigenvalue forced to zero on predicting gene interactions on TRRUST as in Table [4](https://arxiv.org/html/2510.02903v1#A5.T4).**
| Source Gene | # Interactions $\downarrow$ | Precision | Recall | F1 |
| --- | --- | --- | --- | --- |
| JUN | 65 | $69\%\pm 11\%$ | $86\%\pm 7\%$ | $76\%\pm 10\%$ |
| FOS | 25 | $66\%\pm 22\%$ | $79\%\pm 19\%$ | $72\%\pm 21\%$ |
| YBX1 | 24 | $58\%\pm 11\%$ | $45\%\pm 8\%$ | $50\%\pm 8\%$ |
| POU5F1 | 19 | $56\%\pm 39\%$ | $48\%\pm 26\%$ | $51\%\pm 32\%$ |
| SOX2 | 16 | $67\%\pm 12\%$ | $62\%\pm 4\%$ | $64\%\pm 8\%$ |
| HMGA1 | 10 | $47\%\pm 31\%$ | $47\%\pm 34\%$ | $46\%\pm 33\%$ |

**Table 6: Amortized model comparison across the Cite and Multi datasets. We report mean $\pm$ standard deviation of the EMD metric, along with the average across datasets. Lower values indicate better performance. Standard deviation is computed over left-out time points.**
| Model | Cite | Multi | Average $\downarrow$ |
| --- | --- | --- | --- |
| I-CFM [52] | 0.957 $\pm$ 0.211 | 0.892 $\pm$ 0.092 | 0.925 $\pm$ 0.047 |
| OT-CFM [52] | 0.849 $\pm$ 0.007 | 0.821 $\pm$ 0.013 | 0.835 $\pm$ 0.019 |
| Cell-MNN | 0.795 $\pm$ 0.022 | 0.741 $\pm$ 0.104 | 0.768 $\pm$ 0.038 |

**Table 7: Cell-MNN ablation study on single-cell interpolation benchmark when setting one eigenvalue (EV) of ${\bm{A}}_{\theta}$ to zero. Average predictive performance degrades by less than $1\%$.**
| Method | Cite | EB | Multi | Average $\downarrow$ |
| --- | --- | --- | --- | --- |
| Cell-MNN (One EV$=0$) | 0.795  $\pm$ 0.016 | 0.701  $\pm$ 0.076 | 0.746  $\pm$ 0.097 | 0.748  $\pm$ 0.049 |
| Cell-MNN (All EVs predicted) | 0.791  $\pm$ 0.022 | 0.690  $\pm$ 0.073 | 0.742  $\pm$ 0.100 | 0.741  $\pm$ 0.050 |

## Appendix F Additional Numerical Results

We provide further numerical results complementing the main experiments. For the single-cell interpolation task (Section [4.1](https://arxiv.org/html/2510.02903v1#S4.SS1)), Table [7](https://arxiv.org/html/2510.02903v1#A5.T7) reports an ablation in which the model is trained with one eigenvalue set to zero, as later used in the gene interaction discovery experiment. Table [6](https://arxiv.org/html/2510.02903v1#A5.T6) presents the results of the amortization experiment across datasets (Section [4.2](https://arxiv.org/html/2510.02903v1#S4.SS2)).

## Appendix G Additional Qualitative Results

In Figures [7](https://arxiv.org/html/2510.02903v1#A7.F7), [8](https://arxiv.org/html/2510.02903v1#A7.F8), [9](https://arxiv.org/html/2510.02903v1#A7.F9), [10](https://arxiv.org/html/2510.02903v1#A7.F10), we present UMAP projections of the learned operators for each time range, colored by all the cell types reported in the developmental graph of Moon et al. [37]. These correspond to the same UMAPs described in Section [2.1](https://arxiv.org/html/2510.02903v1#S2.SS1.SSS0.Px7), recolored by different cell type to highlight the cell-type dependence of the predicted dynamics. Cells are assigned to a type when the joint expression of the associated marker genes exceeds the 95th percentile. This analysis is enabled by having access to an explicit dynamics model conditioned on time and gene expression, which potentially allows inferences such as identifying when two cell types share similar dynamical laws within a given time range.

Figure: Figure 7:
Refer to caption: x9.png

Figure: Figure 8:
Refer to caption: x15.png

Figure: Figure 9:
Refer to caption: x21.png

Figure: Figure 10:
Refer to caption: x27.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 2
presentation: 3
questions: 1. How is your approach different from VeloVAE and LatentVelo?
2. Why does your less expressive approach outperform more expressive previous models? This seems surprising, because I would expect the local linearity constraint to make the predictions less accurate.
soundness: 3
strengths: Strengths: 
•	The idea is well-motivated, theoretically sound, and biologically meaningful.
•	The idea is clearly explained. The paper is well-written and I found it easy to read
•	The evaluations presented are overall sound and sensible.
summary: This paper introduces a framework for learning a locally linear differential equation in latent space that describes gene expression changes during cell differentiation. The latent space itself is also a linear projection from gene space. A key idea of the paper is to make the dynamics analytically tractable so that gene interactions can be explicitly modeled.
weaknesses: Weaknesses:
•	Highly similar concept to VeloVAE, ICML 2022. VeloVAE fits a mixture of linear ODEs and uses an encoder-decoder framework to perform Bayesian inference of cell times and ODE parameters. Another highly similar paper is LatentVelo, which fits a neural ODE end-to-end to learn dynamics in the latent space. Dynamo uses a linear ODE to estimate the Jacobian of gene regulation.
•	The method seems more incremental than revolutionary. There are closely related approaches for the same problem. The move to a locally linear ODE doesn't seem that impactful to me.
•	Baseline chosen for gene regulatory network evaluation is very weak. There are dozens of gene regulatory network inference algorithms that take single-cell RNA-seq data as input. Comparing performance against these would be more informative. Dynamo (Qiu et al. Cell 2022) seems particularly relevant, because a stated goal of the method is to recover gene-gene interactions by estimating the Jacobian.
•	I understand the motivation for interpretability, but it seems that this local linearity would have to come with some loss of predictive power. Approaches like VeloVAE don’t suffer from this limitation while retaining interpretability. I don't understand how this restricted model can outperform less interpretable but more expressive models, apart from scalability concerns.
•	Evaluation in terms of rate parameters seems important. The locally linear ODE can be interpreted in terms of gene expression changes (RNA velocity), and thus it's important to benchmark against the class of methods that aim to estimate these rates directly. Ground truth in the form of metabolic labeling data (see Dynamo paper) is available for an increasing number of datasets.

---

confidence: 5
contribution: 2
initial_rating: 2
presentation: 3
questions: Due to these concerns, I lean towards rejecting the paper in its current form. Significant theoretical justification, experimental clarification, and comparison to prior ODE-based methods would be needed for a more favorable evaluation. The authors should: 

1. Clarify justification for latent-space linearization. (See weakness 1)
2. Discuss the relation to prior GRN-from-ODE works. (See weakness 2)
3. Substantiate the scalability argument. (See weakness 3)
4. Provide complexity and latent-dimension analysis.  (See weakness 4)
5. Report detailed runtime and efficiency breakdown.(See weakness 5)
6. Revisit the invertibility regularization. (See weakness 6)
7. Clarify interpolation results and superiority to OT baselines. Could the authors explain why Cell-MNN achieves superior interpolation despite training toward similar objectives?  (See weakness 7)
8. Elaborate on amortized training implementation. (See weakness 8)
9. Enhance GRN validation comparisons (e.g., trajectorynet, TIGON). (See weakness 9) 
10. Include hyperparameter and kernel sensitivity studies. (See weakness 10)

References
1. Xiaojie Qiu, Yan Zhang, Jorge D Martin-Rufino, Chen Weng, Shayan Hosseinzadeh, Dian Yang, Angela N Pogson, Marco Y Hein, Kyung Hoi Joseph Min, Li Wang, et al. Mapping transcriptomic vector fields of single cells. Cell, 185(4):690–711, 2022.
2. Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning, pages 9526–9536. PMLR, 2020.
3.  Yutong Sha, Yuchi Qiu, Peijie Zhou, and Qing Nie. Reconstructing growth and dynamic trajectories from single-cell transcriptomics data. Nature Machine Intelligence, 6(1):25– 39, 2024.

This review was independently written by the reviewer. An LLM was employed solely for minor phrasing and grammar improvements.
soundness: 1
strengths: The paper is clearly written and tries to tackle an important problem in single-cell dynamic modeling — learning interpretable continuous-time dynamics.
summary: The paper proposes Cell-MNN, a ODE–based framework that learns explicit single-cell dynamics by modeling local linear ODEs in a PCA latent space. The approach avoids explicit OT preprocessing and claims to simultaneously provide interpretable gene regulatory interactions through the learned Jacobian matrices. Experiments on several benchmark single-cell datasets show competitive interpolation accuracy and advantages over OT-based baselines.
weaknesses: 1. Latent-space linearization is under-justified and fragile. The method hinges on representing dynamics as locally linear in a PCA latent space, $\dot z \approx A_\theta(z,t)z$, advanced via a matrix exponential. The paper itself acknowledges that evolving too far leaves the validity region and requires re-encoding to refresh the local ODE, but provides no analysis of linearization error growth, step-size control, or robustness to curvature in realistic trajectories. Moreover, the claimed efficiency advantage relies on a very low latent dimension; once d_z grows, computing and applying the operator (incl. eigendecomposition/inversion) becomes the dominant cost (see 4), which does not resolve the core limitations of Neural ODEs for stiff/strongly nonlinear flows— it only replaces numerical integration with piecewise-linear surrogates without guarantees.
2. GRN novelty is overstated; prior work is insufficiently acknowledged. Inferring gene regulatory relations from a learned ODE/vector field via its Jacobian has been proposed and used before (e.g., Dynamo [1]) and applied in multi-timepoint trajectory settings (e.g., TrajectoryNet [2], TIGON [3]). Cell-MNN can be viewed as the special case where the vector field is constrained to be piecewise linear in a low-dimensional PCA subspace. The “Gene Regulatory Network Discovery” related-work discussion omits or underplays these lines, which is misleading about the methodological novelty. 
3. Scalability claims (vs. OT) are overstated, and the mechanism is unclear. The paper attributes scalability improvements to replacing OT preprocessing with MMD-based distribution matching and reports out-of-memory (OOM) errors for standard OT on inflated datasets. However, many Neural ODE–based or flow-matching methods already alleviate memory issues by computing OT or matching losses in mini-batches, thereby avoiding O(n^2) coupling matrices in practice. From this perspective, it is unclear why Cell-MNN “solves” sample-size scalability more effectively than these existing approaches, given that both rely on batched stochastic training. Moreover, substituting OT with balanced MMD does not directly address unbalanced dynamics (cell proliferation/death), where Unbalanced-SB frameworks provide principled formulations. Thus, the claimed scalability advantage seems more an artifact of implementation choices than a fundamental methodological improvement.
4. Cubic scaling in the latent dimension threatens practicality beyond very low $d_z$. The paper states $O(T d_z^2)$ for applying the analytical solution, but also a one-time $O(d_z^3)$ per operating point to form the operator (e.g., eigendecomposition/inversion). For $d_z=50–100$, this rapidly dominates and can exceed the cost of Neural ODE field evaluations, undermining the claimed efficiency. The assertion that 5D PCA is “sufficient” is not substantiated by ablations or biological coverage analyses. And I also believe this assertion is not correct.
5. Wall-clock efficiency is not compelling. Despite a 5D latent space and no OT, the reported runs surprisingly high (one hour) for the stated computational simplicity, and notably slower than many baselines under comparable settings. The paper does not explain where time is spent (e.g., repeated eigendecompositions, kernel MMD computation) or provide profiling.
6. Redundant/ill-motivated invertibility regularization. The additional term $L_{\text{inv}}(\theta)$, encouraging $P_\theta$ invertibility, appears unnecessary: invertible matrices are dense; a randomly initialized square matrix is almost surely invertible. The paper does not justify why this regularizer is needed, nor analyze its numerical side-effects (e.g., unstable gradients via $\det(\cdot)$).
7. Interpolation results of OT baselines are unclear. In Table 1, several OT-based methods underperform the OT-Interpolate baseline, while Cell-MNN surpasses it. Since many methods eventually regress toward OT-derived supervision, it is unclear why Cell-MNN outperforms those indirect OT-fitting approaches. The paper should articulate the mechanism (e.g., bias/variance advantages of distribution-level MMD vs. velocity-level supervision, regularization effects) and provide controlled comparisons 
8. Amortized training procedure lacks essential details. Section 4.2 merges datasets in PCA space for amortized training and feeds a dataset index to the model, but does not explain the details.
9. Gene interaction validation lacks competitive baselines. Validation against TRRUST is useful, but the paper does not compare to other vector-field/Neural-ODE methods that could also extract Jacobian-based interactions under the same preprocessing. Without head-to-head comparisons, it remains unclear whether Cell-MNN offers any real advantage for GRN discovery beyond the convenience of an explicit linear operator.
10. Insufficient ablations/sensitivity analyses. Key design choices—kernel and bandwidth for MMD, discount factor $\gamma$, regularization weights $\lambda_{\text{kin}}$, $\lambda_{\text{inv}}$, latent dimension $d_z$, operator parameterization (e.g., fixing one eigenvalue to 0 vs. not ), and $\Delta t$ sampling—lack systematic ablations. Given that interpretability and stability hinge on these knobs, their impacts should be quantified.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Could the authors comment on the model's capability for extrapolation, that is, predicting a future time point beyond the final observed time in the training set?
soundness: 3
strengths: - The method eliminates the need for OT which is typically a computational bottleneck for many methods.
- The method achieves competitive performance on the cell interpolation task, outperforming multiple baselines on three different datasets.
- The authors perform a scaling experiment to benchmark the ability of their model to handle large datasets and show that it is scalable enough for typical contemporary large datasets.
- The ability of the model to learn explicit gene regulatory interactions is useful for interpretability and for learning the explicit gene networks governing differentiation.
summary: In this paper, the authors propose Cell-MNN, a new ODE-based method designed to model the single-cell dynamics of cellular differentiation. The method operates in PCA space, and uses a hypernetwork to learn the transition matrix A. The local dynamics are assumed to be linear. The model is trained by optimizing an MMD loss between the predicted and target marginals of gene expression in the latent space. They benchmark their model on single-cell interpolation, amortized training, and gene interaction directionality.
weaknesses: - The authors perform an unsupervised classification task to predict the direction of regulation for a known TF-gene link from the TRRUST database. The paper does not benchmark the method's ability to predict the links themselves, e.g., by evaluating if the interactions with the highest predicted strength are enriched for known links. Furthermore, analysis of the top source or target genes using gene enrichment analysis could shed light into any relevant pathways.
- The authors mention GRN inference as a complementary line of work but only reference static GRN models. A few works have experimented with learning temporal gene networks for cell differentiation or biological state progression, such as Scenic+, Marlene, Dictys [1-3]. A discussion on how the current work differs from these or a direct comparison is needed to strengthen the contribution.
- The authors project data into 5 dimensions using PCA. This is too limiting for large single-cell datasets with ~20k genes and likely ignores finer interactions between genes. The authors need to present additional work/analysis that uncovers the type of information lost by doing such a projection and the model's ability to learn fine-grained interactions.

[1] https://www.nature.com/articles/s41592-023-01938-4  
[2] https://academic.oup.com/bioinformatics/article/41/Supplement_1/i628/8199402  
[3] https://www.nature.com/articles/s41592-023-01971-3

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: How exactly is the model trained?  The locally linear ODE maps from z_t to z_{t + \Delta_t} but I cannot find an indication of how the authors turn this into an explicit loss.  They make no mention of differentiating through an ODE solver so is there a parameter \Delta_t set in the method and the authors explicitly differentiate through several linearization steps?
soundness: 3
strengths: The authors show good results compared to other interpolation methods on several common benchmarks.  The motivation for the method is straightforward and natural, and I appreciate the work in making simpler and more interpretable models competitive with flow-based trajectory learning.
summary: The authors propose to use a variant of Neural ODEs where the velocity is parameterized as a constant linear function of location at every timestep.  They use this parameterization to fit single-cell RNA data over several timepoints, evaluating with interpolation of a heldout timepoint.
weaknesses: Projecting to 5d PCA is a considerable limitation, as the method seems to not scale and the dynamics of single cell data is generally not so easily compressible.  Additionally, the actual training of the model is somewhat unclear (see questions).

I’m also not totally convinced by the gene analysis results, as it probably requires a baseline to convince a reader that these top gene interactions were not simply the ones encoded by PCA, and any model that learns a velocity (say, MFM-OT or TrajectoryNet) could linearize its velocity at some timepoint and find the same gene interactions.  The later claim about enforcing a zero in the spectrum seems a bit flimsy; it’s true that many individual genes will not be changing during the dynamics, but at only 5 PCs one would expect all PCs were relevant to the dynamics.


Decision: Accept (Poster)