Learning Explicit Single-Cell Dynamics Using ODE Representations
Authors: Jan-Philipp von Bassewitz
Adeel Pervez
Marco Fumero
Matthew R Robinson
Theofanis Karaletsos
Francesco Locatello
Venue: ICLR 2026 Poster (2026)
Keywords: AI4Science
AI4Biology
gene interaction discovery
single-cell dynamics
dynamical systems
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We propose Cell-MNN, an encoder–decoder model with a locally linearized latent ODE representation that achieves competitive performance on single-cell dynamics benchmarks, while discovering biologically consistent gene interactions.

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 learns 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.

Published as a conference paper at ICLR 2026
LEARNING EXPLICIT SINGLE-CELL DYNAMICS USING ODE REPRESENTATIONS
Jan- Phn-Philipp Bassewitz! Adeell Pervez', Marco Fumero', Matthew Robinson', Theofanis Karaletsos², Francesco Locatello
Institute of Science and Technology Austria (ISTA),
Chan Zuckerberg Initiative (CZI)
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 partic- ularly promising and active domain for machine learning. Current state-of-the-ar dynamics models, however, rely on computationally expensive optimal transport preprocessing and multi-stage training, while also not directly learning explicit gene interactions. To address these challenges, we propose Mechansccaiiitic Neural Networks (Cell-MNN), an encoder-decoder architecture whose latent representation a locally linearized ODE governing the dynamies of cellular evolution from stem to tissue cells. Cell-MNN is fully end-to-end (besides a stan- dard PCA pre-processing) and its ODE representation learns interpretable gene interactions. Empirically, we show that Cell-MNN achieves competitive perfor- mance 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.
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 cancer (Chu et al. 2024), neurodegenerative diseases (Cuomo et al., (202)), and to improving wound healing (Ro- drigues et al. 2019). While all cells 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 due both the vast combinatorial search space, with 10th theoretically possible gene interactions, and the experimental effort required to validate specific mechanisms, However, recent advances single-cell sequencing technology (Macosko al. 2015 Zheng et al. have enabled high-throughput measurements that were previously prohibitively expensive, producing datasets that are growing at a pace exceeding Moore's law (Kharchenko) 202T). This rapid growth, coupled with the limitations of direct experimental approaches, presents a unique opportunity to apply machine learning methods to study single-cell dynamics.
In this work we propose Cell-MNN, method to jointly tackle the challenges of predicting cell fate and discovering gene regulatory interactions. Cell-MNN is an end-to-end encoder-decoder archi- tecture whose representation is locally linear ordinary differential equation (ODE) that governs the dynamics of cellular evolution from stem to tissue cells. The ODE representation of Cell-MNN leams explicit and interpretable gene interactions.
A key challenge in modeling single-cell dynamics is that cells are destroyed by measurement, re- sulting in datasets that contain single point along each cell's trajectory (Tong et al. 2020), i.e.,

---

Published as a conference paper at ICLR 2026
PCA-1(2)
Neural ODE predicts velocity = fo(2 t1)
Cell-MNN predicts a linear ODE
z = Ag(Zt, (1)
Zt2
21z
2
21
t
tval
t2
numerical solve
analytical solve
Validation:
EMD Ci Ptyaa)
A12 exp explAgl Ag(=tz f1) (12 ((2-+1)) st
(a) Single-cell interpolation problem
(b) Difference of Cell-MNN with respect to NODEs
Figure (a) Single-cell interpolation: trajectories are evaluated by the earth mover's distance (EMD) between predictions and the marginal distribution at a held-out time wva-- (b) Like hy- pernetwork, Cell-MNN predicts a linear operator As(z,t) t) that approximates the local dynamics explicitly, whereas Neural ODEs and Flow Matching models only output velocity.
snapshot obxervation. This motivated line of work reconstructing trajectories from snapshot data: The best-performing methods in this setting rely on optimal transport (OT) preprocessing to create label trajectories (Tong et al.) 2024a| Zhang al. 2025; Kapusniak et al. 2024 Wang al.) 2025), which becomes computational bottleneck for large datasets due to quadratic scaling of the Sinkhorn algorithm with the number of samples (Cuturi 2013). In contrast, CelI-MNN eliminates OT preprocessing entirely and is designed to be end-to-end. Another bottleneck of sate-of-the-an (SOTA) models such as OT-MFM (Kapusniak et al.) 2024) and DeepRUOT (Zhang et al., 2025) is that they involve multiple training stages and networks, making amortized training across datasets challenging. whereas CcIl-MNN is trained in single stage, enabling straightforward amortized training across multiple datasets. Furthermore, existing SOTA methods focus on accurate interpola- tion of empirical distributions and do not directly learn explicit gene interactions. By comparison, Cell-MNN's ' ODE representation exposes the gene interactions used to predict the cellular evolution. While there are dedicated methods for discovering gene interactions (Lin et al., (025), to the best of our knowledge, such method achieves SOTA predictive performance on single-cell interpola- tion 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 lincarized 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 (Han al. 2018 of gene interactions.
2 LEARNING THE DYNAMICS OF CELLS
Formalizing the Problem. We assume data-generating process consisting of cell state c(I) E C evolving over time in a high-dimensional state space 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 dg genes a, E Rd via an unknown, potentially noisy measurement process m + IRd., so that æ(t) = m(((()) 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 R to be continuous variable and denote an arbitrary time interval by At R. In practice, the lab schedules a discrete set of experimental time points 7 = {t,, tg..... tx } at which cell populations are sampled. We denote by p. the distribution of æ, time f. The dataset of snap-
2

---

Published as a conference paper at ICLR 2026
shot observations is D { with ((0) E T, and our goal is to learn a best-fit mechanistic model for the dynamics of the observable æ, that is consistent with the family of marginals (pr heT.
2.1 CELL--MNN
SOTA models on single-cell interpolation benchmarks face scalability issues from OT preprocess- ing 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 recent architecture that Pervez et al. (2024) showed to outperform Neural ODEs tasks such solar system dynamics and the r--ody 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 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 the resulting version Cell- MNN. Cell-MNN is an encoder-decoder model, learning mechanistic map
Ee+ A Ceel-MNN,, (x,,, At),
which maps gene expression vector 2; at time t to predicted state I++Ai after an arbitrary time interval At. We define the model-induced distribution at time t + At as L44A which is the distribution of Cell-MNNo t, At) when æ, drawn from Pt. As a core part of the architecture, CelI-MNN maps to a compressed representation z E Rd, with d₂ < dz. of the nigh-dimensional gene expression vector æ e Hd and learns the dynamics in the latent space. Following prior work (Tong et al., 2024a), we obtain this latent representation by applying principal component analysis (PCA), with projection matrix VICA E Rdd xds, so that VPCA2.
Locally Linearizing the Latent ODE, The latent vector R.ds in the PCA subspace assumed to follow non-autonomous, non-linear dynamics 1 = (). In practice, this ODE is often highly complex, and learning an explicis form that globally approximates it would be intractable due to the combinatorial search space of basis functions that grows with increasing latent space dimension d,.
To address this, we decompose the intractable global ODE discovery problem into smaller subproblems: the current state (z (2(+), ((0)), which we also call the operating point, we ap- proximate the dynamics by linear ODE in small neighborhood. The learning task is then predict these local dynamics models from the operating point (z((), ((1) using an encoder.
MLP
Subspace of linear operators
global operator
(2(1) (2()) ((1)
d
(1) dt
1 = =f(z,t)
A(2,t)2,
if f(0,t) V
Space of operators
~ A0(20000))
Figure 2: Visualization of the meta-learning task of Cell-MNN's encoder: Rather than directly pre- dicting the velocity at given operating point, as in the Neural ODE framework, the MIP of Cell- MNN maps to the space of linear operators. Con- ditioned the current system state, it predicts lo- cal linear approximations to the global dynamics.
We predict the linear operator Ag E A using multilayer perceptron MIP, Rd.+1 + A. Here A : C(Rd., RRdz) 전 Rd= represents the space of linear operators acting on dRd.. Note that, while the operator governing the lo- cal dynamics is linear, it is non-linear function of the current latent state z and time t(() In Appendix D we show that the reparametrization the right- hand side f(z,t) = A(z, t) z always exists under mild assumptions.
3

---

Published as a conference paper at ICLR 2026
This approach is conceptually orthogonal to Neural ODEs (Chen et al. 2018), which learn uncon- ditional black-box approximation to f. In the Cell-MNNN setting, the MLP functions more like a hy- pernetwork (Ha al. 2017). outputting a conditional white-box linear function go go f|2(6, ((9) = As( ((0 z that locally approximates f at the operating point (z(+) ((0).. Unlike most neural operators (Li et al., 2021 Kovachki et al. (021)) that learn single global operator, Cell-MNN pre- dicts 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 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 local closed-form solution. For fixed Ag at the operating point, the system & Ag is linear, time- ineeinnvriinnt ODE with solution
z 2tt + At) exp (AoAt) =,
Predictions in the gene expression space are obtained by projecting back (t++At) = (tCA=(t+
Parametrization of the Operator. For more fine-grained control over the parametrization of Aa. we let the MLP predict the matrix in an eigen-decomposed form A = Pa diag(Ao) Po- which is also beneficial to compute the matrix exponential. To ensure invertibility of Pe, we train with the additional regularizer L (0) = 1/(det(P) (o))((P))))))))))))) +c)), which is practical if the latent space is small. This also lets us introduce inductive bias by selectively fixing eigenvalues, for example zero, if needed.
Optimization. We train the MLP parameters 0 by minimizing the Maximum Mean Discrepancy (MMD, Gretton et al. (2012 between the model-induced marginals of and the empirical marginals μr, thereby fitting a mechanistic model whose dynamics align with the target marginals pz under future discounting factor 7. All discrepancies are computed in latent space via the pullback kernel
ka(x, = ks (=(VACA VPCAz'),
so that MMD* (gl Pti pt) MMD² př; k.). Here, vi and 0.= denote the distributions of the gene expression marginals in the latent space. The MMD loss is:
fx
MMD
(0) = E
Σ
7
"MD
(ge-pe; pe: k.)
Following Tong et al. (2020), we also regularize the kinetic energy to improve generalization:
(0) E, zy || 112] E,, z. (z,, 2.
which serves as a soft constraint encouraging trajectories close to optimal transport flows in the sense of the Benamou & Brenier (2000) formulation. Our final loss then becomes:
Lioaal
MMD (0) + >ain (0) + >in LinV (().
Computational Complexity. With Ag given in eigendecomposed form at operating point, eval- uating the analytical solution (Eq. 2.1 at T time points has time complexity O de) and space complexity Odd)) where de is the latent space dimensionality. This improves the time and space complexity over the Scalable Mechanistic Neural Network (S-MNN) (Chen et al.) (2025). Forming the full operator requires computing P incurring a one-time (1)) cost per operating point.
Limitations. The cubic time complexity in the latent dimension can become a challenge for very high-dimensional latent spaces but could be mitigated by imposing sparsity assumptions on Ag- In our application to single-cell dynamics, we follow the practice (Tong al. 20200; 2024a) of using a 5, 50 and 100-dimensional PCA space, which captures most of the variance of the data and, crucially, resolves cell-type information (see Appendix Eb. Note that OT preprocessing time points, when using the Sinkhorn algorithm, scales as (dg n²) with the number of samples n, which becomes bottleneck for large datasets, approximate batch approaches are also possible to address this (Tong et al.) 2024a).. A separate limitation of predicting the local dynamics is that evolving the system too far may cause it to leave the regime where the linear ODE is accurate, which would require a new forward pass through the encoder update the ODE, In our experiments, however, we did not encounter this issue.
two
is usually much larger than dz. However,
AA full definition of the MMD given in Appendix B
4

---

Published as a conference paper at ICLR 2026
Uncovering Local Gene Regulatory Interactions. Combining the linear projection to the PCA subspace z VIÇA x with locally linear dynamics around an operating point 2 Ag enables projecting the predicted local dynamics back into the gene expression space with:
d (ViCAz) = Aa VICA
di ===
Agz
x Ao VICA
di
di
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:
wi-z (2), 1) == [VPCA Ao (zz, 1) VPCAl i.j Ty
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 æ;, This makes our proposed approach fully interpretable, we can inspect the learned gene interactions directly.
RELATED WORKS
Single-cell Interpolation.
The single-cell trajectory inference problem. as formalized by Lavenant er al. (2023), entails reconstructing continuous dynamics from snapshot data. Early work based on recurrent neural networks (Hashimoto et al. (016) was followed by Neural ODE-based methods (Tong et al. 2020 2023 Zhang et al.]| [023 Koshizuka & Sato 2023 Huguet et al. 2022). in which a neural network outputs the velocity field governing the dynamics. 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 OT preprocessing to approximate cell tra- jectories (Schiebinger et al., 2019; Bunne et al. (02T). which were also used to train flow matching models such as by Tong al. (2024a); Kapusniak et al. (2024); ZZang al. (2025); Wang et al. 2025): Terpin et al. (2024). 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. (2024a) 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 standard PCA dimen- sionality reduction, which complicates amortized training across datasets. In contrast, CeIl-MNN involves only a single training stage while achieving competitive performance on single-cell bench- marks. Finally, Action Matching (Neklyudov et al. 2023) 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
graph, known gene regulatory network (GRN) (Davidson et al. 2002). Tong et al. (2024b) demonstrated that such GRNs can to some ex- tent be recovered from flow low-matching models in the setting of low-dimensional synthetic data as simulated by Pratapa et al. (2020). In contrast, we show that Cell-MNN learns biologically plausi- ble 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 (Huynh-) approaches (Chan et al. 2017). regression-basee tume-series models (Lu et al. 2021). Gaussian processes (Aijö & Lähdesmäki
Thu et al., 2010 Moerman al.) 2018). information-
2009 and -DE-b models such PerturbODE (Lin al. 2025) and SCODE (Matsumoto al,
2017 However, unlike Cell-MNN, these methods typically learn static GRN and they are either inapplicable to single-cell interpolation benchmarks or do not deliver competitive performance.
Orthogonal to these static GRN approaches, recent methods for time-resolved GRN discovery such as Dynamo (Qiu et al.] 2022), SCENIC+ (Bravo González-Blas et al., (023), Marlene (Hasanaj et al. 2025), and Dictys Waag et al. 2023) infer time-varying GRNs from RNA-velocity paired
scRNA-seg and ScATAC-seq data. In contrast, Cell-MNN operates directly on standard ScRNA- seq UMI counts from single-cell interpolation benchmarks and yields context-dependent signed interaction weights as a by-product of fitting the dynamics.
ODE Discovery. The idea to learn explicit ODE representation of the cell differentiation dy- namics as pursued by PerturbODE (Lin et al. 2025) and Cell-MNN relates directly to the broader problem of ODE discovery. A seminal method in this area is SINDy (Brunton al. 2016), which
5

---

Published as a conference paper at ICLR 2026
infers governing equations from data but requires access to full trajectories, making it unsuitable for the snapshot-basd baseed single-cell setting. Similar limitations apply to more recent approaches such as MNN and related methods such ODEFormer (Pervez al. 2024 Chen et al.) 2025 Yao et al. 2024 Ascoli et al. 2024), which extend ODE discovery to amortized settings by using neural net- works to predict the underlying dynamics from observed trajectories. In contrast, Cell-MNN is qual- itatively 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 naviga- tion filter (Schmidt) 1966), the control of a 2-link 6-muscle arm model (Li & Todorov 20044 Todorov & Li 2005) and rocket landing (Szmuk et al., 2020). The locally linear parameterization theoreti- cally imposes learning control-oriented structure and, in principle, supports the design of performant controllers as described by (Richards al. 2023). 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 commonly studied real single-cell datasets. Following Tong et al. (2020), we include the Embryoid Body (EB) dataset from Moon al. (2019), which after preprocessing contains ~ ~1 K human embryoid cells measured at five time points over 25 days. For EB, we model the time grid with Multiome (Multi) datasets from Burkhardt et al. (2022). repurposed by Tong et al. (2024a). Both consist of gene expression measurements at four time points of cells developing over seven days, with Cite containing ~31 K cells and Multi ~33 K cells after preprocessing. Here we model the time grid with the days of measurement, namely T {0, 1, 2, 3, 7). We use the datasets as preprocessed by Tong et al. (2020 2024a), which involves filtering for outliers and normalizing the data.
{0. 1..]. We also use the CITE-seq seq (Cite) and
Training. We use the same hyperparameters for all experiments unless stated otherwise. Fol- lowing Tong et al.] (2020), we project gene expression to PCA subspace before training. The MLP used to parameterize As has depth 4, width 96, leaky ReLU activations, and Kaiming normal initialization (He et al.] (0015). For stability, initialization so that predictions of Ao start near zero. For the MMD, we use the Laplacian kernel 10-$. We optimize the final
scale the MLP's last layer by 0.01 at
k(2,2') expl- max([a-='Illued with parameters
1 and loss (Eq. D with batch size per time point of 200, future discount factor y 0.1, initialization scale 0.0 and regularization weights Akin = 0.1 and Amn = 1. Optimization is performed using AdamW (Kingma & Ba 2017 Loshchilov & Hutter 2019) with learning rate of 2 x 10-4 and weight decay x 10- Hyperparameters are selected according to grid search and all experiments 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).
a.d,
4.1 SINGLE CELL INTERPOLATION
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.
Following Tong et al. 2020 2024a)
Metric. For easy comparison with SOTA methods, we follow Tong et al. (2020) and report results in terms of the 1-Wasserstein distance in the PCA subspace (W, or otherwise EMD). We use the exact linear programming EMD from the POT (Python Optimal Transport) package (Flamary al., 2021). 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 (Kapusniak) et al., 2024). OT-CFM (Tong et al.) 2024a) and DeepRUOT (Zhang et al. 2025). We found the pre- 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 indicated in the results Table

---

Published as a conference paper at ICLR 2026
Model comparison for single-cell interpolation across the Cite, EB, and Multi datasets, sorted by best average performance. We report the mean + 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.
Table 1:
Method
Cite (5D)
EB (SD)
Multi (SD)
Average +
TrajectoryNet (Tong et al. 2020)
0.848
WLF-UOT (Ncllyddv ct al.
2024
0.800 1.000
NLSB (Koshizuka & Sato
2023)
0.777
+ 0.021
SB-CFM (Tong et al.
2024a)
1.067 = 0.107
1.221
+ 0.380
1.129
0.363
1.139
M-Snkk (Tonn et 2024h
1.054 + 0.087
1.198
0.342
1.098
0.308
1.117
M-Geo Tong et al.
2024b
1.017 1.104
0.879
+ 0.148
1,255
0.117
1.050
I-CFM (Tong et al 20244)
0.965 10.111
0,872 + 0.087
1.085
+ 0.09
0,974
DSB (De Bortol et al. 2021
0.965 1.111
0.862 + 0.023
1.079 0.111
0.969
I-MFM (Kapusmak cf al. 2024
0,916 0.124
0.822 + 0.042
1.053 0.095
0,930
M-F|cM (Tong ct al.
2024b
0,920 1.04
0.793 + 0.066
0.933
0.054
0.882
OT-CFM [Tong al. 2024g
0.882 0.058
0.790
+ 0.068
0,937
0.054
0,870
DeepRUOT (Zhang et al.
2025
0,845 0.167
0.776
1 0.079
0,919
1 0.090
0,846
OT-Interpolate*
0,821 1: 0.004
0.749
1: 0.019
0,830
1 0.053
0,800
OT-MFM (Kapusniak et al. 2024)
0,724
1: 0.070
0.713
1 0.039
0.890
1 1.0.22
0.776
Cell-MNN (ours)*
0.791 1: 0.022
0.690 1 0.073
0.742 1. 0.100
0.741
As an intuitive bar 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 summarizes the results in 5-dimensional PCA subspace on all three datasets. This setting is the most realistic in terms of scientific application, as trajectory inference is typically done in low-dimensional PCA subspaces, which capture most of the variance (See Appendix E for more details). 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 the only method that outperforms our proposed OT- Onterropeee benchmark 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 OT-Interpolate as the ground truth. This also explains the strong performance of OT OTt-Interpolate. In Appendix G.T we provide additional evaluations in 50- and 100-dimensional PCA subspace showing that Cell-MNN performs competitively in these settings as well.
4.2
AMORTIZED TRAINING
I-CFM
OT-CFM
CaIl-MNN
Model
Cite (Inflated)
EB (Inflated)
Multi (Inflated)
LCFM
0.0390 0.0249
0.0403
0.0004
0.0482 1 0.0144
OT-CFM
0OM Rrror
EMD
0.8
DeepRUOT
0OM Bzroz
Batch-OTe CT
0.0232 0.0041
0.0243
0,0025
0.0302
0.0010
0.6
Cell-MNN
0.0225 0.0021
0.0240
0.0039
0.0252 0.0072
Average
(a) Scaling experiment
(b) Amortization experiment
Figure 3: (a) Model comparison across the synthetically inflated datasets. We report mean + stan- dard deviation of the MMD metric, along with the average across datasets. Lower values indicate better performance. Standard deviation is computed over left-out time points. (b) Comparison of models jointly trained on Cite and Multi datasets test potential for amortization. We report mean + standard deviation of the EMD metric, along with the average across datasets.
Foundation models have shown strong transfer learning capabilities across datasets in variety of domains (Bodnar et al. 2025 Pearce et al. 2025 Bodnar et al. 2025). However, current SOTA methods for single-cell interpolation, such as OT-MPM and DeepRUOT, rely multi-stage train- ing 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 datasets with the same time scale, namely Cite and Multi.
7

---

Published as a conference paper at ICLR 2026
Training. Our amortized training setup follows the single-cell interpolation experiment described in Section 4.1 with the only differences being that (i) we iteratively sample batches from Cite and Multi, (ii) we wider network with width 128, and (iii) pass 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 different set of genes, we the same PCA embeddings as in the previous experiment (Section 4.1) and merge datasets in the PCA subspace.
Baselines. We use OT-CFM 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. 2024a) and first reproduce the original results for the separate-datase setting 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 3b 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 train 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-IIterpollaee can become significant bottleneck due to the quadratic time and space complexity of the Sinkhorn algorithm (Cuturi 2013). To experimen- tally compare the scalability of different methods, 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 (00,, 0.1) to the PCA embeddings. Cell-MNN is trained with the same hyperparameters as our first experiment in Section 4.1
Baselines. We run OT-CFM and DeepRUOT on the inflated datasets and observe that both meth- ods 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 miti- gate this, Tong et al. (2024a) 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 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 also requires estimating the OT map. We therefore use the MMD metric with a Laplacian kernel compute the validation score. Since MMD can be computed in 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 forming dataset-wide OT preprocessing. as required by standard OT-CFM or DeepRUOT, leads to OOM eтrors due to the quadratic space complexity of the Sinkhorn algorithm. Batch-OT-CFI out- performs 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.
3a Per-
8

---

Published as a conference paper at ICLR 2026
4.4 DISCOVERING GENE INTERACTIONS
JUN
POUSF1
Interaction strength
Target Genes
HAND1 HMGAA
10.0
24-30
50%
100%
HMGB8
HIMGU2
JUNE
103
0.0
UMAP-2
5.0
7.5
18-23
FOS
YBX1
12-17
POUSE
TERF
S0X2
-1.0
6-11
HANIE
IMMN
SOX2
HMGA1
POU
0-5
-2.5
Cell-MNN EV=0)
Source Genes
10
Cell-MNN
UMAP-1
Guessing
(a) Gene interactions
(b) UMAP of operators
(c) FI Score on TRRUST
Figure 4: (a) Strongest predicted gene interactions by Cell-MNN for days 12-17 of the EB dataset, normalized to the range [-1, 1]. (b) UMAP projection of operators predicted by Cell-MNN on the EB dataset, showing that the model learns distinct dynamics different time points. (c) Validation of predicted gene interactions by two Cell-MNN versions: For each source gene j. we classify each TRRUST edge j > as activating or repressing using the sign of Cell-MNN's learned weight w-----
Days 0-5
Days 6-11
Days 12-17
Days 18-23
Days 24-30
10'
lng cape 1eeesii
UMAP-2
10
-5
10
Figure 5: UMAPs of the predicted operators by Cell-MNN across the five time ranges of EB. Points colored by whether joint expression of the EN-1 marker genes FOXA2 and SOXI7 is above the 95th percentile. Clustering indicates that Cell-MNN learns distinct dynamics for the EN-1 cell type.
CeII-MNN predicts the local dynam- ies of the cell differentiation process in gene expression space, thereby learn- ing interaction weights Wi-si from each gene to every other gene as described in Section 2.1 This corresponds to unsu- pervised learning of local gene interac- tions. To assess whether these learned interactions are biologically meaning- ful, validate them against the literature-curated TRRUST database (Han al. 2018), which contains 8,444 regulatory relationships synthe- sized from 11,237 PubMed articles.
Table 2: GRN discovery on the EB dataset (F1 score in %). For each source gene we report the Fl score of pre- dicting activation vs. repression for the N..t labeled TR- RUST interactions. Best method per row is highlighted in bold. Performance averaged over three seeds for OT- CFM and Cell-MNN.
Source gene
Cell-MNN (IEV=0)
Nest
SCODE
OT-CFM (J)
JUN
65
56.34%
64% 1.2%
71
1 168
FOS
25
38.10%
50%
6%
71% 10%
YBXI
54.55%
72%
+
51%
+
POUSFI
19
14.29%
13%
= 1%
67% + 6%
SOX2
16
84.21 %
60% 0%
71%
1: 1.9%
While TRRUST represents only small subset of all potential relationships, provides a valuable reference signal for evaluating Cell-MNN's predictions.
HMGA1
10
25.00%
31% %
9%
80%
1 6%
Average
46% 25%
48% + 22%
69% 10%
Unsupervised Classification. Using labels from the TRRUST database, we design an unsupervised classification task for a source gene j: for every interaction j - listed TRRUST, we predict repressing. Since Cell-MNN outputs wj-i per cell, whether the relationship is activating average these values over the dataset to obtain a single prediction for each interaction, classifying it activating if ({ ()->x(x f) 0 and repressing if smaller than zero. We report results for the genes j that in the top 10 most active ones (as by summing over all interaction weights

---

Published as a conference paper at ICLR 2026
per gene) for any time point and that have more than 10 matching TRRUST interactions within the EB gene set. To evaluate Cell-MNN''s 's capability of predicting the existence of interactions, we additionally performed enrichment analysis which is presented in Appendix F.I
Training and Inductive Bias. We train an ensemble of Cell-MNN models for single-cell interpo- lation on EB, each with different left-out marginal. The setup matches Section 4.1 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 Ag: in particular, we know that only some genes vary over time, implying that at least one eigenvalue of A 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% 1%, it significantly improves gene interaction discovery performance on TRRUST (see Appendix 8 for ablation results).
Baselines. To contextualize the results, we focus our comparison to approaches that predict signed GRNs, as this is required for our classification task. As an ODE-based baseline, we run SCODE Matssumoto et al. 2017) on the dataset. To compare with the idea of using the Jacobian of Neural ODEs as a proxy for the GRN (Qiu et al.) 2022). we compute the Jacobians of a fully trained OT-CFM model (OT-CFM (J)). This can be thought of as an alternative to the operators predicted by Cell-MNN. Similar to Cell-MNN, we use the sign of the Jacobian to classify an interaction as activating or repressing. Given different models, we only adapt the GRN prediction in our pipeline; the rest of the evaluation remains the same.
Results. We compute precision, recall, and F1 scores for the unsupervised classification task, with all numerical values in Table 5 & 66] As shown in Table 2 Cell-MNN outperforms both SCODE and OT-CFM(J) by twenty percentage points in terms of Fl score, indicating that for the tested gene interactions. the model can meaningfully discover activation or repression in fully unsupervised manner. Interestingly, the Cell-MNN variant with one eigenvalue set to zero improves classification performance from 60% to 69% in F1 score, demonstrating that the inductive bias introduced on the operator effectively constrains the solution space. Note that gene interactions are context-dependen (e.g., varying by cell type), 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 mean- ingful signal that Cell-MNN has the potential to discover biologically plausible gene interactions.
Visualizing Operators. To visualize the learned operators, we plot UMAP projections of Ag com- puted jointly across all time ranges (Figure 4b) and separately for each time range (Figure 55) 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 (Moon et al. 2019) are provided in Appendix [H]
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 model is able to perform trajectory inference, while explaining its predictions in terms of the underlying gene regulatory inter- actions used. Empirically, we show that Cell-MNN achieves competitive performance on single-cell benchmarks, as well as in scaling and amortization experiments. Importantly, the gene regulatory interactions learned from real single-cell data exhibit consistency with the literature-curated TR- RUST 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 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 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 (Jinek et al. 2012)
10

---

Published as a conference paper at ICLR 2026
ETHICS STATEMENT
All datasets used in this work (EB, Cite, Multi) are publicly available and were preprocessed by prior works Tong et al. (2024a) 2020 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. 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. We provide code and hyperparameters for reproducibility.
REPRODUCIBILITY
We use the datasets in the form preprocessed by prior work: Cite and Multi from Tong et al. 22224a 2 and EB from Tong et al. | For the gene discovery experiment, we use the less pre- processed version of EB provided by Tong al. (0024)). All experiments are run and hyperparameters are presented in Table 3 in the Appendix together with the hardware and train- ing time. We fix random seeds and report all results averaged over three runs with different seeds. The code is available on GitHub under giggtthubb.. com/ cz - /cel -mmnn containing environ- ment specifications in envire ronment. yml. Commands reproduce results are documented in the accompanying README md.
single GPU
ACKNOWLEDGEMENTS
This work supported by the Chan Zuckerberg Initiative (CZI) through the External Residency Program. We thank CZI for the opportunity to participate this program, which enabled close collaboration and access to critical resources. We furthermore thank Peter Kharchenko, Bianca Du- mitrascu, and Alicia Michael for very helpful discussions and guidance related to the single-cell application domain. We thank Alexander Tong and Lazar Atanackovic for sharing code and techni- cal clarifications, which were essential for reproducing their results and conducting our subsequent experiments. Similarly, we thank the first authors of DeepRUOT (Zhang et al. 2023) and VGFM (Wang et al.) (025). Dongyi Wang and Zhenyi Zhang, for providing their code and assistance with its use. Finally, we thank the Causal Learning and Artificial Intelligence group ISTA for their valuable feedback and discussions throughout the project.
REFERENCES
Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution the monge- kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000. URL
https://doi.... www.//rgg..... 1007 s0021 00000/0002000000222
Cristian Bodnar, Wessel Bruinsma, Ana Lucic, Megan Stanley, Anna Allen, Johannes Brandstet- ter, Patrick Garvan, Maik Riechert, Jonathan A. Weyn, Haiyu Dong. Jayesh K. Gupta, Kit Tham- biratnam, Alexander T. Archibald, Chun-Chich Wu, Elizabeth Heider, Max Welling, Richard E. Turner, and Paris Perdikaris. A foundation model for the earth system. Nature, 641(8065): 1187, 2025. URL http::ttt/:/dat wwwpiioorg1.... 10.00.103300 1038/ 010///44458600250090000
Carmen Bravo González-Blas, Seppe De Winter, Gert Hulselmans, Nikolai Hecker, Irina Matetovici, Valerie Christiaens, Suresh Poovathingal, Jasper Wouters, Sara Aibar, and Stein Aerts. Scenic+: single-cell multiomic inference of enhancers and gene regulatory networks. Nature Methods, 20 19):1355-1367, 2023.
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 Na-

---

Published as a conference paper at ICLR 2026
tional Academy of Sciences, 113((15)3932-3993 2016. URL https: hwtp
Charlotte Bunne, Laetitia Meng-Papaxanthos, Andreas Krause, and Marco Cuturi. Jkonet: Proximal optimal transport modeling of population dynamics. CoRR, 2021. URL https://arxiv.
org/abs/22 2206. 06345
Daniel Burkhardt, Malte Luecken, Andrew Benz, Peter Holderricth, Jonathan Bloom, Christopher Lance, Ashley Chow, and Ryan Holbrook. Open problems multimodal single-cell integration.
https://kaggle: com/ http:://gg......ooooo////o/////r mmttimmmdld 2022. Kag-
gle.
Thalia E. Chan, Michael P.H. Stumpf, and Ann C. Babtie. Gene regulatory network infer- ence from single-cell data using multivariate information measures. Cell Systems, 5(3):251- 267.e3, 2017. ISSN 2405-4712. URL ltttps://wwy sci www.:sscccecncirttcccm htt.:e//iieettccoooccccc.eeeeeee p
Jiale Chen, Dingling Yao, Adeel Pervez, Dan Alistarh, and Francesco Locatello. Scalable mecha- nistic neural networks. In The Thirteenth International Conference on Learning Representations, 2025. URL httttttt//ppeeeeeeeee
net wwt./ooummidd-arraaaaaeoee 8A244
Ricky T. Q Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordi- nary differential equations. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, vol- ume 31, 2018. URL httpp:://proocednggg..... neurips....errr...ooo.mrm cc/ppper faappprr_flllsepppprrr fi BLL11LLLL2SBSEbbaae6333226668833
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, 91)::170, 2024. ISSN 2059-3635. htLhttp:::/ 0t0ooloorg/1..101/1111133333333
024 01851-y
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. URL
s41576- 023- 00599-3
Marco Cuturi. Sinkhorn distances:
Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems, volume 26. 2013. URL https::/prrcceddii a 07700007662227111552225555353333 pdf
neur heurips. iss cc wttp:/parrralaeeee/eeeeeeeee http//ppprp../aaappppp//////////
Stéphane d' Ascoli, Sören Becker, Philippe Schwaller, Alexander Mathis, and Niki Kilbertus. ODE- Former: Symbolic regression of dynamical systems with transformers. In The Twelfth Interna- tional Conference on Learning Representations, 2024. URL
Corum21d=12oHLiGVW5
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 C. Clarke, Maria I. Arnone, Lee Rowen, R. Andrew Cameron, David R. McClay, Leroy Hood, and Hamid Bolouri. A genomic regulatory network for development. Sci- 2002.
ence,
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. In M. Ran- zato, A. Beygelzimer, Y. Dauphin, P.S. Liang. and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 17695-17709, 2021. URL ht ott.httsss//pooceeeeeeeraeoooooe
Fi Eiee:3000225552222ceec0223333333
12

---

Published as a conference paper at ICLR 2026
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aure@lie Boisbunon, Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, LÂ@o Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy. levgen Redko, An- toine 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.
wttppaapes/////2222200222 451. html
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:// owt.rroorrgppppppp //v.svvvvvv v133 /l3////et/o/oo//o/o/oooo/o 2a html
David Ha, Andrew M. Dai, and Quoe V. Le. Hypernetworks. International Conference on Learn ing Representations, 2017. URL
forum wttprummmddrkreeeeeee
Heonjong Han, Jae-Won Cho, Sangyoung L 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, (6(D1)) January 2018. ISSN 1362- 4962. URL http://///.o..../1..1111000
Euxhen Hasanaj, Barnabás Póczos, and Ziv Bar-Joseph. Recovering time-varying networks from single-cell data. Bioinformatics, 41 (1(Supplement.1) 1):6:tt:)::iii2-01 07 2025.
Tatsunori Hashimoto, David Gifford, and Tommi Jaakkola. Learning population-level diffusions with generative mns. Proceedings Machine Learning Research, pp. 2417-2426, New York, New York, USA, 20-22 Jun 2016. PMLR. URL https://pzoceeeinng/ www.inng.smmm....m..mmemm.. mressp v48/hashimote html]
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance imagenet classification. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), December 2015.
Guillaume Huguet, S. Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. Manifold interpolating optimal-transpor flows for trajectory inference, 2022. URL httttttss///xxiii> http://wwwwww.../...a.. abs/ 1122266.49288 14928
Vân Anh Huynh-Thu, Alexandre Irrthum, Louis Wehenkel, and Pierre Geurts. Inferring regulatory networks from expression data using tree-based methods. 5(9):1-10, 09 2010. URL https:
http:/110..11//ourneae. pone. 0012776
Martin Jinek, Krzysztof Chylinski, Ines Fonfara, Michael Hauer, Jennifer A. Doudna, and Em- manuelle Charpentier. A programmable dual-rna-guided dna endonuclease in adaptive bacterial immunity. Science, (376666))8166-888 2012.
Kacper Kapusniak, Peter Potaptchik, Teodora Reu, Leo Zhang, Alexander Tong, Michael M. Bron- stein, 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 htt.:httt///pneeeeeeee
ew net./
FFNN
Peter V. Kharchenko. The triumphs and limitations of computational methods for scrna- seq. Nature Methods, 18(7):723-732, 2021. URL
s4 841592-021-01172
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017. URL
https://arxiv. org/
wwgg/bs///ag abs/ 6980
Takeshi Koshizuka and Issei Sato. Neural lagrangian schrödinger bridge: Diffusion modeling for population dynamics. In ICLR, 2023. URL
ii=
d30NWD 13QNWD_pcFV
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, 2021. URL htl.tttts::/.. //apiiiw//////.//////a/aap htt.l//orga/abb//////e 08481
13

---

Published as a conference paper at ICLR 2026
Hugo Lavenant, Stephen Zhang, Young-Heon Kim, and Geoffrey Schiebinger. Towards mathemat- ical theory of trajectory inference, 2023. URL htRp:s/t::///// atwpi..iiiii./... httgg/bb///00g 09204
Weiwei Li and Emanuel Todorov. Iterative linear quadratic regulator design for nonlinear biological movement systems. In Proceedings of the 1st International Conference on Informatics in Control, Automation and Robotics, (ICINCO 2004), volume 1, pp. 222-2222 01 2004.
Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhat- tacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021. URL
https:/ htttss://openrevatt opwnreview.e..eee/ net http:/ooormmidieeee 2a
Zaikang Lin, Sei Chang, Aaron Zweig, Minseo Kang, Elham Azizi, and David A. Knowles. In- terpretable neural odes for gene regulatory network discovery under perturbations, 2025, URL
https://arxiv. org/ http://aaxx....rr/////22.2...... 02409
Iyya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019. URL httpps://ppenreview. wttp://wwww.e//oru.m.....
Bkg6RiCqY7
Jonathan Lu, Bianca Dumitrascu, Ian C. McDowell, Brian Jo, Alejandro Barrera, Linda K. Hong. Sarah 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):1-29, 01 2021.
1371 1371//juuraa pcbi 1008223
Evan Z. Macosko, Anindita Basu, Rahul Satija, James Nemesh, Karthik Shekhar, Melissa Gold- man, 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 gewomeenowiiieee widd expression profiling of individual cells using nanoliter droplets. Cell, 161 (5):1202-1114, 2015. Uttpttps:://////tttt http:///ii..rgg///// 1ttp:/...1111///111
1011/ www.eell.... 2015. 05.002)
Hirotaka Matsumoto, Hisanori Kiryu, Chikara Furusawa, Minoru $ H Ko, Shigeru B H Ko, Norio Gouda, Tetsutaro Hayashi, and Itoshi Nikaido. Scode: an efficient regulatory network inference algorithm from single-cell ma-seg during differentiation. Bioinformatics, 33(15):2314-232) 04 2017. URL
10 htt..0//1///aaaaaooataaetee////a
Thomas Moerman, Sara Aibar Santos, Carmen Bravo González-Blan Jaak Simm, Yves Moreau, Jan Aerts, and Stein Aerts. Graboost2 and arboreto: efficient and scalable inference of gene regulatory networks. Bioinformatics, 35112)1215--21111 11 2018. ISSN 1367-4803. URL //ttt::::/dd...
ca cs/bty915
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 Biotechmology. (111)))14821492, 2019. URL https:
//rature. com/ articles/ s41 1441587-119--11 0336-3
Kirill Neklyudov, Rob Brekelmans, Daniel Severo, and Alireza Makhzani. Action matching: Learn- ing stochastic dynamics from samples. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th Inter- national Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 25888925888, PMIR, 23-29 Jul 2023. URL. httpp:://procee Sress/v2022 http://nwkkyddov2aa.mmmmmm html
ngs .mlr.
Kirill Neklyudov, Rob Brekelmans, Alexander Tong, Lazar Atanackovic, Qiang Liu, and Alireza Makhzani. computational framework for solving wasserstein lagrangian flows. In ICML, 2024.
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- generative cell atlas across 1.5
billion years of evolution: The transcriptformeer cinl model. bioRxiv, 2025. URL https: biorx htwp://wwwwwi......ccoo////////p htwp://wwwaa...///oo..//////////
04. 650731
14

---

Published as a conference paper at ICLR 2026
Adeel Pervez, Francesco Locatello, and Stratis Gavves. Mechanistic neural networks for scientific machine learning. In Forty-first International Conference on Machine Learning. 2024. URL btttps://openrevien net/ wtt.:/oooummiitttuiie
Aditya Pratapa, Amogh P. Jalihal, Jeffrey N. Law, Aditya Bharadwaj, and T. M. Murali. Benchmarking algorithms for gene regulatory network inference from single-cell transcrip- tomic data. Nature Methods, 17(2):177-154 2020. URL. https://doi wwweiii..o.r//gooo .rgg/0/ggg0. /00. 1038/
841592-019-069C
Xiaojie Qiu, Yan Zhang, Jorge D. Martin- Chen Weng, Shayan Hosseinzadeh, Dian Yang, Angela N. Pogson. Marco Y. Hein. Kyung Hoi (Joseph) Min, Li Wang, Emanuelle Grody, Matthew J. Shurtleff, Ruoshi Yuan, Song Xu, Yian Ma, Joseph M. Replogle, Eric S. Lander, Spy- ros Darmanis, Ivet Bahar, Vijay G. Sankaran, Jianhua Xing, and Jonathan S. Weissman. Mapping transcriptomic vector fields of single cells. Cell, 185(4):690-) 7111990--1100 710--111.446 1.4,, 2022.
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 (eds.), Proceedings of the 40th Inter- national Conference Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 29051-29062 PMLR, 23-29 Jul 2023. URL https: /ttpas///pcceedingss.. .mlr. html
press/v202 http:0//rricaardss//o
Melanie Rodrigues, Nina Kosaric, Clark A. Bonham, and Geoffrey C. Gurtner. Wound healing: A cellular perspective. Physiological Reviews, 99(1):665-706 2019. |RL /doi .org/
10. wtw.1122ppyyseee 00067.200 .2017
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-transpon analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928 923 2019. URL https: wwwp://www//.oo///////////wt wcw.wssseieenedrec...... com/
Stanley F. Schmidt. Application of state-space methods to navigation problems. volume of Advances in Control Systems, pp. 293-340. Elsevier, 1966. URL https: /t::://wwwe //www
sci scienneddrreci www.eeett...c.eeeeeeeeeeeeeeeeee htt.:///.....aaaaaa.a....//////.
Michael Szmuk, Taylor P. Reynolds, and Behçet Açıkmeşe. Successive convexification for real- freedom powered descent guidance with state-triggered constraints. Journal
time dix-degree-of-freeee
of Guidance, Control, and Dynamics, 43(8):1399-1411 1413, 2020. htLLhttss///:occco
0554/1..G04504
Antonio Terpin, Nicolas Lanzetti, Martín Gadea, and Florian Dorfler. Learning diffusion light- speed. In The eight Annual Conference on Neural Information Processing Systems, 2024.
URL httpss:://pppeeeeeaaaeteee www.://wwwweeeett.iiddeedaaaaaae
Emanuel Todorov and Weiwei Li. A generalized iterative lgg method for locally-optimal feedback control of constrained nonlinear stochastic systems. In Proceedings of the 2005, American Control Conference, 2005., pp. 300-306 vol. 1, 2005.
Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajecto- ryNet: A dynamic optimal transport network for modeling cellular dynamics. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 9526-9536. PMLR, 13-18 Jul 2020. URL httph::t:///e httt:::://proceedingg .ww.ediiinssmm........mmm.... mlr pttp.......eee//11ttt//00000
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 transcrip- tional and regulatory dynamics driving cancer cell plasticity using neural bde-based optimal trans- port. BioRxiv, 2023. URL https:://// wwwp://www.// bi wwwp:/iioiiiiiiii org http///ottettt http:////www.....ooooooooooooooo
29/2023. 03. 28 644
15

---

Published as a conference paper at ICLR 2026
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 Machine Learning Research, 2024a. URL net forumrrrmiio
w (tt.uummmmmcccnnnnen3333 Expert Certification.
Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation- schrödinger bridges via score and flow matching. In AISTATS, pp. 2779-1287, 2024b. URL https::/proceeding/ mllr.
preass//333/y-toog22aa.__. htmml
Dongyi Wang, Yuanwei Jiang, Zhenyi Zhang, Xiang Gu, Peijie Zhou, and Jian Sun. Joint velocity- growth flow matching for single-cell dynamics modeling, 2025. https::/aaxxxi... org/
Lingfei Wang, Nikolaos Trasanidis, Ting Wu, Guanlan Dong, Michael Hu, Daniel E. Bauer, and Luca Pinello. Dictys: dynamic gene regulatory network dissects developmental continuum with single-cell multiomics. Nature Methods, 20(9):1368- 1378, 2023.
Dingling Yao, Caroline Muller, and Francesco Locatello. Marrying causal representation learn- ing with dynamical systems for science. The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: /oppn//v///eereeeew net wwt.:/rretrrrrpa../// forum?id
MWHRxX2 4mq
Jiaqi Zhang, Erica Larschan, Jeremy Bigness, and Ritambhara Singh. scnode Generative model for temporal single cell transcriptomic data prediction. bioRxiv, 2023. URL https://www.
biorxiv. org/
Zhenyi Zhang, Tiejun Li, and Peijie Zhou. Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport. In The Thirteenth International Conference on Learn- ing Representations, 2025. / http:///ppreeeeeeeeeemeeeee new. forum ? 2id=g0lxd3Mttrr
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 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 Hindson, and Jason H. Bielas. Massively parallel digital transcriptional profiling of single cells. Nature Communications, 8(1):14049, 2017. URL https::/doi.
1038 ncomms 4049
Tarmo Äijö and Harri Lähdesmäki. Learning gene regulatory networks from gene expression mea- surements using par-meraiiiiic molecular kinetics. Bioinformatics, 25(22): 08 2009.
URL httppp/:/.aaaaaao....00///aaaaaa
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.
URL http:p///ww...eeeeeeeeeee.....oo
$136757881003005
16

---

Published as a conference paper at ICLR 2026
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 im- proving the readability of the paper. All substantive research contributions, analysis, and interpreta- tions were carried out by the authors.
B DEFINITIONS
Maximum Mean Diserepancy (MMD, Gretton al. (2012)): ): Given distributions p and q over X' and a positive-definite kernel function k X' x X R, the squared Maximum Mean Discrepancy (MMD) defined as
MMD kk)
E Bz, (k(1, (kp,,=) By [k(,,0)) -2 y. [k v)]
C ARCHITECTURE
Encoder
Decoder
Ar
S2
Vida
MLP
Representation Agz
zt
Vyca
Figure 6: The Cell-MNN architecture first applies the PCA projection matrix map the gene pression state æ to a latent representation z, An MLP then predicts a locally linear approximation Agz the dynamics at the operating point (zz, (). To decode, the analytical solution of this ODE is evaluated at future time point and projected back into gene expression space.
D PROOFS
Proposition (Extension of Proposition 1 of Çimen (010 Let f Rd, x R + Rdd satisfy
((, 0 for all RR and assume f E CK Rd: R) with k 1. Then there exists matrix-
valued map Rds x TR + Rdd xdz such that f(z, t) = A(z, 1) z for all (z, E TRds x RR
Proof. Fix (z, 1) E Rd and define y : [0, + Rd= by (s)) f(s 2,1). Since f C and k > 1, the map + (s) is differentiable and
d (ss) = D
ds
By the fundamental theorem of calculus,
f(z,t) f(0,t) (1(11)-00) 2 =(f Dzf(s D.f(s.2.t)ds)
Using f(0,t) 0 gives f(z,t) A(z,t)z
17

---

Published as a conference paper at ICLR 2026
Table 3: Collection of training hyperparameters used in all experiments. All models were trained on single NVIDIA RTX 2080 Ti GPU. Training on EB, Cite, and Multi separately, as well as on the amortization experiment, required about 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
MLP (Ag)
Depth:
Width: 96 (128 for amortization experiment)
Activation: Leaky ReLU
Initialization: Kaiming normal
Last layer scale: 0.01
Laplacian kernel kzz, =') expl- max [2 [h.e)
MMD kernel
ddz
σ 10-8
Optimization
Batch size per time point: 200 Future discount factor 7 = 0.1
Initialization scale: 0.01
Regularization Akin 0.1, Ainv
Optimizer: AdamW
Learning rate: 2 x 10
Weight decay: 1 x 10-5
Validation
Frequency: every 10 steps
Patience: 40 validation checks 60 minutes (240 minutes for inflated datasets)
Training time Randomness
Seeds: 3
Hardware
1x NVIDIA GeForce RTX 2080 Ti (11 GB RAM)
E ON THE USE OF LOW DIMENSIONAL PCA EMBEDDINGS
Low-dimensional PCA as standard practive. In line with prior work on single-cell trajectory inference, we operate in a low-dimensional latent space obtained by PCA on the gene expression matrix. Concretely, we project each dataset to the first five principal components and train both Cell-MNN and all baselines in this common representation. This choice follows the prevailing assumption in computational biology that scRNA-seq data lie near a low-dimensional manifold, and it is consistent with the implementation of nine previous works reported in Table
Empirical validation of the 5D PCA representation. Figure 7 provides a quantitative and qual- itative assessment of the 5D PCA space used in our experiments. Figures 7a 7cl and 7 show that the first few principal components account for the large part of the variance in each dataset: using five components yields cumulative explained variance above 60% for Cite and Multi and above 40% for Embryoid. Importantly, this div-dimensional representation preserves cell-type structure.
We compute k-nearest neighbor classification in the 5D PCA space with & 15 and obtain accura- cies of 87% (Multi) and 90% (Cite); for Embryoid, cell-type labels are not available. Consistently, the UMAPs computed from the 5D PCA embedding (Figures 7b and 7 exhibit clear clustering by cell type, indicating that the embedding retains the information required to distinguish cellular lineages. To show flexibility with respect the number of principle components used, we also train Cell-MNN and OT-CFM in 10 dimensional PCA subspace with no tuning and find that they perform similarly, see Table TO
Scope of our contribution. Our focus in this work is orthogonal to learning the optimal represen- tation: We focus on improving the dynamical model given standard low-dimensional embedding. Within the 5D PCA space, Cell-MNN achieves SOTA average performance single-cell interpo- lation benchmarks (Table T). while removing OT preprocessing and learning explicit, local ODEs that can be interpreted as gene interactions. For scientific questions centered lineage bifurcations and fate decisions, the key requirement is that the representation preserves cell-type structure. The experiments in Figure 7 show that this requirement is met. Exploring richer or alternative represen-
18

---

Published as a conference paper at ICLR 2026
Map
Mep
MouP
Pirdppal component
Lmap, Sprct
(a) Multi: Explained variance
(b) Multi: colored by cell-type
Percene Venance Eaglenee 141
DP
FryP
Mup
Nea?
Prircipal Comporent
(c) Cite: Explained variance
(d) Cite: UMAP colored by cell-
Parcartage teronce Exelaieed
Princigal
(e) Embryoid: Explained variance
Figure 7: (a, e) Explained variance of each of the principle components plotted for the datasets Multi, Cite and Embryoid. The first few components capture the majority of the variance. (b, d) UMAPs computed on the 5 dimensional PCA embedding colored by cell-type. Clustering by cell- type shows that 5 dimensional PCA embedding retains cell-type information.
tations is an interesting orthogonal direction, and Cell-MNN can in principle be applied on top of such alternatives without changing the core methodology.
19

---

Published as a conference paper at ICLR 2026
Table 4: Gene selection specifications for the TRRUST experiment. There are 16 genes that make up the top 10 predicted high-intenn 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.
In
Interactions in TRRUST
in Training
> 10
Top Source Gene
TRRUST
Gene Set
Interactions
HMGA1
True
18
10
True
HMGB2
True
2
JUNB
True
15
4
FOS
True
True True
JUN
True
173
65
POUSFI
True
25
19
True
HANDI
False
0
0
D2
True
TERFI
True
PITX2
True
11
4
ID3
True
2
HMGBI
False
0
0
SOX2
True
True
HMGA2
True
5
YBXI
True
33
24 True
IDI
True
Table 5: Validation of predicted gene interactions TRRUST: For each source gene j, classify each TRRUST edge j-i as activating or repressing using the sign of the learned weight wi from (averaged
Cell-MNN with eigenvalue the number interactions in TRRUST and classification metrics (precision, recall, and F1) shown as mean + std across ensemble models trained on three different seeds.
cells). For each source gene, report
Source Gene
# Interactions
Precision
Recall
F1 71% 6%
JUN
65
62% + 8%
82% 3%
FOS
25
65% + 10%
80% + 10%
71% = 10%
YBXI
24
55% + 10%
48% 3%
51% + 6%
POUSFI
19
82% + 6%
58% + 11%
67% + 6%
SOX2
16
73% + 12%
69% 8%
71% + 9%
HMGA1
10
78% + 9%
82% + 2%
80% 6%
F GENE REGULATORY INTERACTION RECOVERY
To quantitatively assess the learned gene interactions, designed unsupervised classification task based on the TRRUST database, which contains literature-curated gene regulatory interactions, many of which annotated activating or repressing. For evaluation of our model, focus on the most dominant source genes predicted by Cell-MNN, i.e., those with the highest mean in- teraction strength with other genes. A source gene included into the experiment 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 robust prediction for each interaction. Based on these predictions, we compute precision, recall, and Fl scores to quantify how well the model recovers known regulatory mechanisms and report them in Table 5 and Table
least 10 of
20

---

Published as a conference paper at ICLR 2026
Table 6: Ablation of models with no eigenvalue forced to zero on predicting gene interactions on TRRUST as in Table
Source Gene
# Interactions
Precision
Recall
F1
JUN
65
69% + 11%
86% + 7%
76% + 10%
FOS
25
66% + 22%
79% + 19%
72% 21%
YBXI
24
58% + 11%
45% + 8%
50% 8%
POUSFI
19
56% + 39%
48% + 26%
51% 32%
SOX2
16
67% + 12%
62% = 4%
64% 8%
HMGAI
10
47% + 31%
47% + 34%
46% = 33%
Table 7: Amortized model comparison across the Cite and Multi datasets. We report mean + stan- dard 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 +
I-CFM (Tong et al. 2024a)
0.957 + 0.211
0.892 + 0.092
0.925 + 0.047
OT-CFM (Togg et al. 2024a)
0.849 = 0.007
0.821 = 0.013
0.835 + 0.019
Cell-MNN
0.795 + 0.022
0.741 +- 0.104
0.768 + 0.038
Table 8: Cell-MNN ablation study on single-cell interpolation benchmark when setting one eigen- value (EV) of Ag to zero. Average predictive performance degrades by less than 1%.
Method
Cite
Multi
Average +
CeIl-MNN (One EV= 0)
0.795 + 0.016
0.701
+ 0,076
0.746
+ 0.097
0.748 1:0.4
Cell-MNN (All EVs predicted)
0.791
0.022
0.690
+ 0.073
0.742 +: 0.100
0.741 0.050
F.1 PREDICTING THE EXISTENCE OF INTERACTIONS
As an additional validation of the gene interac- tions predicted by Cell-MNN, we evaluate its per- formance predicting the existence of regulatory links. To this end, we first rank all predicted in- teractions by their inferred strength and then as- sess this ranking against the TRRUST database. To restrict TRRUST to interactions that are plau- sibly involved in the differentiation dynamics of the EB dataset, we subset the database to inter- actions whose transcription factor (TF) regulator is mentioned as relevant in the original analysis of Moon et (2019) (Fig. 6d), This yields 447 interactions regulated by 70 TFs, which we treat
Table 9: Comparison of GRN discovery meth- ods on the EB dataset. We report preci- sion @ 500 and AUROC for predicting existence of TRRUST interactions. Higher is better. All metrics were computed over three seeds.
Method
Enrichment@
AUROC
GRNBoos12 SCODE
20.429
0.633
16.714
0.686
OT-CFM (J)
20.429
0,661
Cell-MNN (ours)
18.572
0.659
as our ground -trnthhh-thhhh-e signal. This restriction is also
necessary to keep the evaluation of the baselines computationally tractable. We define the candidate interaction set as all directed TF-tagget pairs where the TF is among the 70 EB regulators and the target is any gene in the EB dataset.
Baselines. As this experiment only requires a ranking of interactions, we can compare against methods that predict unsigned GRNs. We therefore include the widely used GRN discovery methods GRNBoost2 and SCODE as baselines. We also compute the performance of OT-CFM (J) on this task. Because the absolute interaction strength, ignoring the sign of the effect.
only interested in the presence of regulatory link, evaluate all methods
Evaluation and Metrics. For each method, obtain scalar interaction score for every TF- target pair in the candidate set, and assemble these into a TF-target score matrix in the common gene space. For Cell-MNN, interaction scores derived from average over 100 operators from
21

---

Published as a conference paper at ICLR 2026
each time point in the dataset. We use ensemble of three Cell-MNN models each trained with different left-out time point. We evaluate all methods using AUROC and Precision@ K. The AUROC measures how well method separates interacting from -on-interacting pairs actoss all possible thresholds, whereas Precision @ K measures the proportion of true interactions among the top-K ranked edges. To make Precision @ K more interpretable, uniform ranking over all candidate interactions and refer to the resulting quantity as Enrichment @ K. Intuitively, Enrichment @ K measures by which factor a method outperforms random guessing.
normalize it by the base precision of
Results and Discussion. The results of this experiment are summarized in Table 9 We find that Cell-MNN performs competitively with the baselines terms of both AUROC and Enrichment @K. We note that we restrict GRNBoost2 to learning interactions only for the 70 EB regulators, which ef- fectively provides them with additional prior knowledge. For both CeII-MNN and OT-CFM, we also substantially parse-grain their outputs by averaging interaction scores across cells and discarding sign information. Consequently, both methods, which in principle can produce context-dependent signed predictions, are evaluated here in much more restricted, global setting.
G ADDITIONAL INTERPOLATION RESULTS
Table 10: Model comparison a 10-dimensional PCA embedding. We report the mean + standard deviation of the EMD metric across the Cite, EB, and Multi datasets, along with the average across datasets. Lower values indicate better performance.
Method
Cite (10D)
EB (IOD)
Multi (10D)
Average (10D)
OT-CFM
1.491 0001
1.607
0.074
1.678 = 0.248
1.592 = 12
Cell-MNN (ours)
1.502 00012
1.587 0.113
1.709 0017
1.599 = 0.101
We provide further numerical results complementing the main experiments. For the single-cell interpolation task (Section 4.1). Table [8 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 7 presents the results of the amortization experiment across datasets (Section
G.1 HIGH DIMENSIONAL EXPERIMENTS
In this section, evaluate the performance of Cell-MNN in higher-dimensional latent spaces. To this end, we train the same model, with slightly modified hyperparameters, in 50- and 100- dimensional PCA subspaces. Following Neklyudov Following/Nellludov et al. (2024), we do not whiten the data in PCA space to preserve the empirical variance structure. To keep feature magnitudes in a numerically well-conditioned range for MLP training while preserving their relative variance, we rescale all components by the standard deviation of the first principal component. We find that this improves the stability of training Cell-MNN. Since the EMD is homogeneous under global rescaling of both distributions (EMD (AX, AY) = À EMD(X. X, Y), À > 0), we multiply the EMD scores computed on the rescaled PCA coordinates by the standard deviation of the first principal component so that they are comparable to those obtained on the original (unscaled) PCA space.
Due to increased RAM requirements, we use a different GPU, namely an NVIDIA L40S (48 GB RAM). This also allows us to train with larger batch size of 1028. We set the learning rate to x 10-3 (50D) and 5 10-5 (100D), patience for early stopping to 10 evaluation steps and keep the remaining hyperparameters unchanged. We train on the Cite and Multi and find that the runtimes range from 4m 25s to 31m 29s, depending on the seed and left-out time point. Memory usage remains below 25 GB of RAM in this setup.
We report the results of the experiments in Table Without tuning the hyperparameters further, find that Cell-MNN performs within error bars of SOTA approaches for the Multi dataset.
We remark that, due to the analytical solution of the ODE, one can choose to decode the trajectories at fewer time discretization points without impacting the accuracy of the predicted trajectories. This can be used to reduce the RAM requirements of the method and is unique when compared to Neural ODEs, whose accuracy depends on the step size due numerical solving.
22

---

Published as a conference paper at ICLR 2026
Table 1I: Single-cell interpolation on 50- and 100-dimensional PCA embeddings across the Cite and Multi datasets. We report the mean + standard deviation of the EMD metric computed over three seeds. Values marked * are computed by us.
Method
Cite (50D)
Multi (50D)
Cite (100D)
Multi (100D)
I-CFM
41.834 3.284
49.779
4.400
48.276
3.281
57.262
3.85
WLF-SB
39,695 1,935
47.828
6.382
46.131
55.065
5.49
WLF-OT
38.352
0.203
47,890
6,492
44.821
0.126
55.416
6,997
OT-CFM
38.756
+ 0.398
47.576
=: 6.622
45.393
= 0.416
54.814
1 5.860
SFF M-Exctt
40.009
0.783
45.337
2.833
46.530
0.426
52.888
1,986
M-
38.524
= 0.293
44,795
= 1.911
44.498
0416
52.203
= 1.957
WLF-UOT
37.007 + 1,200
46.286
= 5.841
43.731
1.375
54.222
= 5.827
Cell-MNN (ours)*
38.803 = 0.635
43,926
2,590
46,020 = 1.177
52.698
+ 2,341
OT-CFM*
38,576 = 0.429
43.141
= 3.918
45.368 0.433
51.399
3,922
OT-MEM
36.394 1.886
45.160 4.960
1.020 41.784
50,906
4.627
H ADDITIONAL QUALITATIVE RESULTS
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. (2019). These correspond to the same UMAPs described in Section 2. 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 identifying when two cell types share similar dynamical laws within a given time range.
In Figures 10
23

---

Published as a conference paper at ICLR 2026
coloured by TNNT2 expression (top-95 for Cardiac Precursors
Days 0-5
Days 6-11
Days 12-17
Days 18-23
Days 24-30
expresiion
UMAP-2
Aix_t coloured by T,CER GATA4 expression (top-95 %6 Cardiac-1
Days
Days
Days
Days 16-23
Days
UMMP22
10-
20
10
A/x_t. coloured by NEK2-5, expression Itop-95 for Cardiac-2
Days
Days
Days 12-17
Days 18-23
Days
UMAP-2
10
-5
UMAP-1
ADx_t t coloured by I FOXA2, 5CX17 expression (top-95 %)) for EN-1
Days 0-5
Days 6-11
Days 12-17
Days 18-23
Days 24-30
10
UMAP-2
10
10
10
10
UMAR-1
Afx. f. coloured by GATA3,SATB1, 110.35AT1111,XXOX1 15 expression (top-95 %) for EN-2
Days 0-5
Days 6-11
Days 12-17
Days 18-23
Days 24-30
UMAP2
10
UMAP-1
ADCC. tl coloured by I ARIDJA expression (top-95 %) - for EN
Days 0-5
Days 6-11
Days 12-17
Days 18-23
Days 24-30
10'
10
100 exprssiion
UMM-22
20
20
20a
UMAP-1
Figure 8
24

---

Published as a conference paper at ICLR 2026
Alx coloured NANOG,POUSFI expression (top-95 ESC
Days 0-5
Days
Days 12-17
18-23
Days 24-30
10
10
UMAP-2
high
no high
30
A(x_t. t) by HOXBA, CD3,1, expression (top-95 Hemangidalast
Days
Days
Days
Days 16-23
Days
201
UMMP2
10
-5
A(x_t. coloured by TBXS, MYC expression (top-95 %6 for Lateral Plate ME
Days
Days
Days 12-17
Days 18-23
Days
UMAP-2
20
10
UMAP-1
Atxtt U) coloured by expression (top-95 %) - for
Days 0-5
Days 6-11
Days 12-17
Days 18-23
Days 24-30
201
10
UMAP-
nO high
10
10
10
10
UMAR-1
SOX9 Afx I. t) coloured by E FAX3,FOX33, FOX9 50X110 expression (top-95 %) for NC
Days 0-5
Days 6-11
Days 12-17
Days 18-23
Days 24-30
10
UMMP2
10
UMAP-1
Abx_t U) coloured by I HOXA2) JOX01 expression (top-95 %) - for NE/NC
Days 0-5
Days 6-11
Days 12-17
Days 18-23
Days 24-30
10g expression
UMMP22
20
20a
UMAP-1
Figure 9
25

---

Published as a conference paper at ICLR 2026
GBX2.OLIG3.HOKD expression (top-95 %)
Days 0-5
Days 6-11
Days 12-17
Days 18-23
24-30
UMAP-2
10
20
10
coloured by KLF7,ISL DNCCC L.NECCT expression Itop-95 Neuron Subtype
Days 0-5
Days
Days
Days 16-23
Days
201
UMMP2
10
-3
-5
Alx_t. coloured by NUGLLNAASLLLAI NR2F1 expression Itop-95 for Neuron Subtype
Days
Days
Days 12-17
Days 18-23
Days
10
UMAP22
t0
10
-5
-5
UMAP-1
Alx_t. coloured by NUNLLLLLL.O LMX1A expression (top-95 %) for Neuron Subtype
Days 0-5
Days 6-11
Days 12-17
Days 18-23
Days 24-30
201 corezsoon
UMAP2
10
10
UMAP-1
Abx.t, coloured by NKOC-8,EN2, E,2,,5X expression (top-95 %) for Neuron Subtype
Days 0-5
Days 6-11
Days 12-17
Days 18-23
Days 24-30
UMAP2
10
UMAP-1
Aix_t U) coloured by PAX6,Z8T016 expression (top-95 %) for Neuron Subtype
Days 0-5
Days 6-11
Days 12-17
Days 18-23
Days 24-30
1.0 opresiion
UMA-22
20
20a
UMAP-1
Figure 10
26

---

Published as a conference paper at ICLR 2026
coloured by NES,MAP2 expression (top-95 %)
Days 0-5
Days
Days 12-17
18-23
Days 24-30
10
expression
20
UMAP-2
30
Abx.t. coloured by CDX2,ASC12,KLF5, NKX2-1 expression (top-95
Days
Days
Days
Days 18-23
Days
UMMP22
10
10
Atx_t tl coloured by POUSF1,0TX2 expression Itop-95 %6 Pre-NE
Days
Days
Days 12-17
Days 18-23
Days
UMAP-2
20m4
10
20
-5
UMAP-1
ADK_t. U coloured by EOMES,MIXLI SAT.L expression (top-95 %)
Days 0-5
Days 6-11
Days 12-17
Days 18-23
Days 24-30
10'
UMAP2
10
10
UMAP-1
AOx.t. t) coloured by TEX1 B.SIX2,TBXXI 5, PDGFRA expression (top-95 %6) Smooth Muscle Precursors
Days 0-5
Days 6-11
Days 12-17
Days 18-23
Days 24-30
101
UMMP2
UMAP-1
Atx_t. l) coloured by WT1 expression (top-95 %) for Epicardial Precursors
Days 0-5
Days 6-11
Days 12-17
Days 18-23
Days 24-30
10'
10
100 cpressioon
UMAP22
10
20
20a
UMAP-1
Figure 11
27


Reviews:
[{"review_id": "stpct1FlC9", "reviewer": "Reviewer_4yFL", "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.", "questions": "1. How is your approach different from VeloVAE and LatentVelo?\n2. 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.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Strengths: \n•\tThe idea is well-motivated, theoretically sound, and biologically meaningful.\n•\tThe idea is clearly explained. The paper is well-written and I found it easy to read\n•\tThe evaluations presented are overall sound and sensible.", "weaknesses": "Weaknesses:\n•\tHighly 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.\n•\tThe 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.\n•\tBaseline 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.\n•\tI 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.\n•\tEvaluation 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VjFprwbgot", "reviewer": "Reviewer_51Jg", "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.", "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: \n\n1. Clarify justification for latent-space linearization. (See weakness 1)\n2. Discuss the relation to prior GRN-from-ODE works. (See weakness 2)\n3. Substantiate the scalability argument. (See weakness 3)\n4. Provide complexity and latent-dimension analysis.  (See weakness 4)\n5. Report detailed runtime and efficiency breakdown.(See weakness 5)\n6. Revisit the invertibility regularization. (See weakness 6)\n7. 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)\n8. Elaborate on amortized training implementation. (See weakness 8)\n9. Enhance GRN validation comparisons (e.g., trajectorynet, TIGON). (See weakness 9) \n10. Include hyperparameter and kernel sensitivity studies. (See weakness 10)\n\nReferences\n1. 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.\n2. 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.\n3.  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.\n\nThis review was independently written by the reviewer. An LLM was employed solely for minor phrasing and grammar improvements.", "limitations": "", "rating": 2, "confidence": 5, "soundness": 1, "presentation": 3, "contribution": 2, "strengths": "The paper is clearly written and tries to tackle an important problem in single-cell dynamic modeling — learning interpretable continuous-time dynamics.", "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.\n2. 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. \n3. 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.\n4. 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.\n5. 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.\n6. 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)$).\n7. 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 \n8. 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.\n9. 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.\n10. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iqZMODWuTP", "reviewer": "Reviewer_otWS", "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.", "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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The method eliminates the need for OT which is typically a computational bottleneck for many methods.\n- The method achieves competitive performance on the cell interpolation task, outperforming multiple baselines on three different datasets.\n- 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.\n- The ability of the model to learn explicit gene regulatory interactions is useful for interpretability and for learning the explicit gene networks governing differentiation.", "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.\n- 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.\n- 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.\n\n[1] https://www.nature.com/articles/s41592-023-01938-4  \n[2] https://academic.oup.com/bioinformatics/article/41/Supplement_1/i628/8199402  \n[3] https://www.nature.com/articles/s41592-023-01971-3", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "okxDo4gO6Q", "reviewer": "Reviewer_JgxX", "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.", "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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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.", "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).\n\nI’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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)