Predicting Time-Varying Flux and Balance in Metabolic Systems using Structured Neural ODE Processes

## Abstract

Abstract We develop a novel data-driven framework as an alternative to dynamic flux balance analysis, bypassing the demand for deep domain knowledge and manual efforts to formulate the optimization problem. The proposed framework is end-to-end, which trains a structured neural ODE process (SNODEP) model to estimate flux and balance samples using gene-expression time-series data. SNODEP is designed to circumvent the limitations of the standard neural ODE process model, including restricting the latent and decoder sampling distributions to be normal and lacking structure between context points for calculating the latent, thus more suitable for modeling the underlying dynamics of a metabolic system. Through comprehensive experiments ( 156 156 156 156 in total), we demonstrate that SNODEP not only predicts the unseen time points of real-world gene-expression data and the flux and balance estimates well but can even generalize to more challenging unseen knockout configurations and irregular data sampling scenarios, all essential for metabolic pathway analysis. We hope our work can serve as a catalyst for building more scalable and powerful models for genome-scale metabolic analysis. Our code is available at: https://github.com/TrustMLRG/SNODEP .

## 1 Introduction

A distinctive characteristic of deep neural networks is their capability to implicitly learn complicated features and dynamics from data, significantly saving human effort in composing those handcrafted features and devising complex models. Therefore, there has been a growing interest in using them in a variety of scientific contexts, such as quantum chemistry (von Glehn et al., ), tokamak controller design (Degrave et al., ), climate sciences (Lam et al., ; Nguyen et al., ), molecule generation (Hoogeboom et al., ) and drug discovery (Askr et al., ), to name a few. For drug discovery problems in particular, it is essential to answer the questions of where and how the drug should be targeted. The machine learning community has attracted increased attention in molecular design to address the latter question (Luo et al., ; Corso et al., ). On the other hand, metabolic pathway analysis techniques, such as flux balance analysis (FBA) (Orth et al., ) and dynamic FBA (Mahadevan et al., ), have been shown highly effective in finding drug targets (Sen & Orešič, ). These methods are widely used to study the effect of drugs or environmental stress simulated by gene knockouts on unwanted cells, such as cancer cells, by curbing their metabolism (Raškevičius et al., ).
Nevertheless, several key parameters, including the optimization objective and constraints for the reaction flux in their linear programming (LP) formulation, must be determined using domain expertise for each case, largely limiting their generality and scalability.
In this work, we aim to develop scalable data-driven methods that can directly predict the behavior of metabolic systems with time-varying flux, thus avoiding the manual effort required to build FBA models.

More specifically, we achieve this by leveraging single-cell RNA sequencing (scRNA-seq) time-series data (Chen et al., ) and using single-cell flux estimation analysis (scFEA) technique from Alghamdi et al. () to estimate flux and balance of the metabolic system, because scRNA-seq can churn out data in bulk, and getting time-series single-cell gene-expression data is much less labor intensive than getting actual flux-balance time-series data. The challenge, however, lies in that gene expression trajectories for individual cells cannot be tracked over time since cells die once their gene expression is read. Instead, we only have gene expression samples from different cells at each timestep, which can be viewed as samples from a time-varying distribution resembling a random process. In fact, it’s well known that gene transcription is stochastic, especially when considered at the single cell level (Thattai & Van Oudenaarden, ). Thus, the amounts of molecules produced, or the chemical concentration, from a collection of cells can be considered to be sampled from some distribution, with the amounts of mRNA molecules showing a Poisson-like behavior in a steady state as shown in Thattai & Van Oudenaarden ().

Since the time-varying metabolic concentrations are known to follow a non-linear ordinary differential equation (ODE), we propose a novel *Structured Neural ODE Process* (SNODEP) architecture that is built on top of the standard neural ODE processes (Norcliffe et al., ) to predict the underlying dynamics of the metabolic system.
We note that standard neural ODE processes have several design choices that might not help to model the ODE dynamics in our case, like lack of structure in the encoder to get the latent distribution from the context points and the use of Gaussian parametric family for latent posterior and decoder distributions. Consequently, we design the architecture of SNODEP to bypass these shortcomings, showing improved performance in tasks such as predicting gene-expression distributions on unseen timesteps, predicting metabolic-flux and metabolic-balance distribution on unseen timesteps, and predicting the corresponding distributions for gene-knockout cases, considering both regularly and irregularly sampled data, all for several metabolic pathways.

Contributions.
We formulate the prediction problem of metabolic flux and balance as a stochastic neural processing task, where the goal is to learn the underlying dynamics by predicting their time-varying distributions under different configurations (Section 2).
We propose an end-to-end training framework, which first defines the intermediary steps required to estimate metabolic flux and balance from scRNA-seq data and then learns a novel SNODEP model that can predict the unseen time points of flux and balance and their dynamics under gene-knockout configurations (Section 3.2). The proposed SNODEP architecture is designed by addressing a few limitations of the standard architecture of neural ODE processes (Section 3.1); thus, it is more suitable to model the time-varying distributions from metabolic systems.
Comprehensive experiments on real-world datasets and various metabolic pathways demonstrate that SNODEP is highly effective in modeling the dynamics of gene expressions and predicting metabolic flux and balance, consistently outperforming alternative models such as standard neural ODE processes (Sections 4.2-4.4).
We also showcase the superiority of SNODEP under gene-knockout variations and scenarios with irregularly sampled data (Section 4.5), suggesting its versatility and strong potential in solving challenges in biomedical domains.

### 1.1 Related Work

Metabolic Pathway Analysis.
Genome-scale metabolic models (GSMMs) have proven to be powerful tools in the design of therapeutic treatments. For instance, Raškevičius et al. () employed GSMMs to identify therapeutic windows for cancer treatment, while Larsson et al. () used them to simulate gene knockouts in a Glioblastoma cancer cell model, identifying potential therapeutic targets and predicting side effects in healthy brain tissue. Despite their importance, GSMMs are time-consuming and require significant domain expertise to build. Recent studies have explored integrating machine learning techniques with GSMMs, as reviewed in Sahu et al. (). From a dynamical standpoint, Costello & Martin () framed pathway dynamics prediction as a machine learning problem, using XGBoost models to predict such dynamics, but their framework is not end-to-end. More recently, Aghaee et al. () introduced a graph neural network model to simulate the dynamic behavior of metabolites in oxidative stress pathways in bacterial cell cultures for synthetic data.
In addition, RNA velocity (La Manno et al., ) estimates the time derivative of gene expressions but needs spliced and unspliced mRNA counts, usually not reported in the experiments. Similarly, Klumpe et al. () investigated single-cell time series prediction, albeit also using synthetic data with no specific focus on metabolic pathways. To the best of our knowledge, our work is the first to comprehensively study the dynamically varying flux and balance of metabolic pathways derived from real-world single-cell gene expression time-series data.

Neural ODE.
The neural ODE family of models has shown strong capabilities in modeling dynamic systems, particularly when the underlying dynamics are known to follow an ODE (Rubanova et al., ). While latent neural ODEs have been applied to interpolation and extrapolation tasks, they are not suitable for modeling random processes. In contrast, neural processes (NP) (Garnelo et al., ) can be used for modeling time-varying distributions, but they have no consideration for the underlying dynamics. These observations motivate us to explore models like neural ODE processes (NODEP) (Norcliffe et al., ), where the dynamics are defined over the parametric space of these distributions. Other models, such as those proposed in Kidger et al. (), assume a noisy evolution of dynamics, which does not align with our prediction problems of time-varying distributions in metabolic systems. Our work adapts standard neural ODE processes (Norcliffe et al., ) to better suit our specific settings, showing improvements across various tasks and metabolic pathways.

## 2 Problem Formulation

Figure: Figure 1: Overall pipeline of our framework for predicting time-varying distributions, such as gene expressions, flux, and balance, with (green) and without (orange) gene knockouts.
Refer to caption: x1.png

Classical methods like DFBA estimate time-varying metabolic flux and balance by solving an optimization problem to maximize the biomass at each timestep (see Appendix 7 for more details).
Our work proposes to directly train a model on scFEA-estimated flux-balance values until a certain timestep and then predict the distributions of gene expression, flux, and balance in future timesteps, expecting that the trained model will learn the underlying dynamics.
Figure 1 illustrates the overview of our pipeline.
Due to page limits, we defer more details on scFEA proposed by Alghamdi et al. () to Appendix B.2.
Below, we provide detailed descriptions of our problem setup. The key notations and their descriptions are provided in Appendix A.

Predicting Gene Expression, Flux and Balance. Suppose we have a gene count matrix of dimension $K\times N$, where $N$ is the total number of cells and $K$ is the total number of genes, with gene counts measured at each regular timestep $t$ and total $V$ timesteps. Let $\mathbb{B}_{t}$ be the index set representing the cells whose gene counts $\in\mathbb{R}^{K}$ are observed at time $t$. Then, we have $\sum_{t}|\mathbb{B}_{t}|=N$, indicating that all $N$ cells get their expressions counted over various timesteps.

For a metabolic pathway, we only extract the relevant $d$ genes from the total set of $K$ genes. Let $g_{i,t}\in\mathbb{R}^{d}$ be the gene-expression array for cell $i\in\mathbb{B}_{t}$ at time $t$, and $\mathbf{G}_{t}\in\mathbb{R}^{d\times|\mathbb{B}_{t}|}$ be the corresponding matrix.
For a certain metabolic pathway with $u$ modules and $v$ metabolites and each batch $\mathbb{B}_{t}$ of cells at time $t$, we estimate the flux $m^{f}$ and balance $m^{b}$ using the scFEA framework detailed in Appendix B.2. Specifically, we define:

- •
$S^{f}_{t}:\{g_{i,t}\}_{i\in\mathbb{B}_{t}}\rightarrow\{m^{f}_{i,t}\}_{i\in
\mathbb{B}_{t}}$ as the mapping that estimates the flux $m^{f}_{i,t}\in\mathbb{R}^{u}$ for each cell $i$ based on its gene expression. Let $\mathbf{M}^{f}_{t}\in\mathbb{R}^{u\times|\mathbb{B}_{t}|}$ be the matrix of the flux samples.
- •
$S^{b}_{t}:\{g_{i,t}\}_{i\in\mathbb{B}_{t}}\rightarrow\{m^{b}_{i,t}\}_{i\in
\mathbb{B}_{t}}$ as the analogous mapping for estimating the metabolic balance $m^{b}_{i,t}\in\mathbb{R}^{v}$. Let $\mathbf{M}^{b}_{t}\in\mathbb{R}^{v\times|\mathbb{B}_{t}|}$ be the matrix of the balance samples.

We note that scFEA was originally developed for static-FBA, but since the static-DFBA formulation (Equation 4) can be interpreted as solving the static-FBA for different timesteps, we use scFEA to estimate flux-balance values for different timesteps.

Gene-knockout.
Gene knockout is a way to understand how a gene influences the metabolic network, for example, in understanding how essential genes in pathogens affect metabolic pathways to design drugs to inhibit those pathways (Larsson et al., ); it’s also widely used in synthetic biology Dalvie et al. ().
In the gene-knockout simulations in FBA models, the constraints of the reaction fluxes affected by essential genes are usually modified (Maranas & Zomorrodi, ).
In contrast, we train a model on certain gene-knockout configurations and then predict the distribution on unseen configurations and timesteps.
To simulate gene-knockout conditions, we randomly sample $S$ subsets of $k$ most expressed genes, set the gene-expression of genes from those subsets to zero (See Algorithm 1 for more details), and estimate the flux-balance values again based on the scFEA techniques. For $s\in\{1,2,\ldots,S\}$, we analogously define $\{\tilde{m}^{f}_{i,t}\}_{s,i\in\mathbb{B}_{t}}$ and $\{\tilde{m}^{b}_{i,t}\}_{s,i\in\mathbb{B}_{t}}$ as gene-knockout flux and balance estimates, respectively, where we use $\tilde{\mathbf{M}}^{f}_{s,t}\in\mathbb{R}^{u\times|\mathbb{B}_{s,t}|}$ and $\tilde{\mathbf{M}}^{b}_{s,t}\in\mathbb{R}^{v\times|\mathbb{B}_{s,t}|}$ to denote the corresponding matrix of samples.

Essentially, our framework assumes that metabolic flux and balance from scRNA-seq data can be estimated using scFEA techniques, that knocking out a subset of genes does not change the expression levels of the rest of the genes, and that gene essentiality and gene expression levels are correlated.

Learning Objective.
For each timestep $t\in\{t_{1},t_{2},\ldots,t_{V}\}$, we collect samples of gene expression $\{g_{i,t}\}_{i\in\mathbb{B}_{t}}$, flux $\{m^{f}_{i,t}\}_{i\in\mathbb{B}_{t}}$ and balance $\{m^{b}_{i,t}\}_{i\in\mathbb{B}_{t}}$ and their gene-knockout samples $\{\{\tilde{m}^{f}_{i,t}\}_{j,i\in\mathbb{B}_{t}},\{\tilde{m}^{b}_{i,t}\}_{j,i
\in\mathbb{B}_{t}}\}$ with cells $\mathbb{B}_{t}$ using previous steps. We assume these samples are drawn from some underlying distributions corresponding to gene expression $G(\theta_{g}(t))$, flux $M^{f}(\theta_{f}(t))$, balance $M^{b}(\theta_{b}(t))$ and their gene knockout versions $\{\tilde{M}^{f}(\theta_{f}(t)),\tilde{M}^{b}(\theta_{b}(t))\}$, respectively.
The goal is to learn a model $F:t\rightarrow Y(\theta_{t})$ that can predict the underlying dynamics of time-varying distributions, which depend on some latent distribution $L$. In our setup, $F$ is considered as an encoder-decoder neural network, with a different network for each distribution in $\{G,M^{f},M^{b},\tilde{M}^{f},\tilde{M}^{b}\}$.

Let $C<T<V$ be the length of context, target, and total available time points, respectively.
Given a distribution $Y$, let $y_{i}\sim Y(\theta_{t})$ for any $i\in\{1,\ldots,V\}$. When we say $y_{i}\sim Y(\theta_{t})$, it means a random sample from the set $\{y_{i,t}\}_{i\in\mathbb{B}_{t}}$.
During training, our model’s encoder takes as input the context data, which includes samples from context points $\mathcal{C}=\{(t_{1},y_{1}),\ldots,(t_{C},y_{C})\}$. The decoder then predicts samples from the target points $\mathcal{T}=\{(t_{1},y_{1}),\ldots,(t_{T},y_{T})\}$.
During inference, the model is used to predict every timestep available, including hitherto unseen timesteps $\mathcal{V}$ = $\{(t_{1},y_{1}),\ldots,(t_{V},y_{V})\}$.
In the following discussions, we denote $\mathbb{I}_{\mathcal{C}}=\{1,\ldots,C\}$ and $\mathbb{I}_{\mathcal{T}}=\{1,\ldots,T\}$ for simplicity.

## 3 Methodology

### 3.1 Issues with Standard Neural ODE Process

The standard neural ODE process (NODEP) model (Norcliffe et al., ) employs an encoder-decoder model architecture, where the context points $\{(t_{i},y_{i})\}_{i\in\mathbb{I}_{\mathcal{C}}}$ are used to calculate the latent distributions $L_{0}(\theta_{l_{0}})\text{ and }D(\theta_{d})$, and the latent $l_{0}\sim L_{0}$ evolves over target timesteps $\{t_{i}\}_{i\in\mathbb{I}_{\mathcal{T}}}$, according to an ODE that is modeled by a neural network $\mathbf{f}_{w}$ as follows:

$$ $l(t_{i})=l_{0}+\int_{t_{0}}^{t_{i}}\mathbf{f}_{w}(l(t),d,t)dt.$ (1) $$

The time-evolving latent distributions are then fed into a decoder to obtain the target distributions: $\{N_{i}(y_{i}|\mu_{w_{1}}(l(t_{i})),\sigma_{w_{2}}(l(t_{i})))\}_{i\in\mathbb{I
}_{\mathcal{T}}}$.
Although NODEP has been shown effective in modeling ODE dynamics for scientific discovery, there are a few limitations with NODEP if applied to our settings:

- 1.
The latent and decoding distributions are treated as Normal. This is not the best choice of distributions to model gene-expression data, which is usually discrete and Poisson-like (Thattai & Van Oudenaarden, ) and confirmed by Figures 9 and 9 in Appendix E.
- 2.
The encoded representation $r_{i}=f_{e}(\{t_{i}^{\mathcal{C}},y_{i}^{\mathcal{C}}\})$ is calculated using context points without any particular structure in NODEP. These $r_{i}$’s are then order-invariantly aggregated to give $r$, and finally $D\sim q_{D}(d|\mathcal{C})=\mathcal{N}(d|\mu_{D}(r),\text{diag}(\sigma_{D}(r)))$, similarly for $L_{0}$. The order between the context points and their sequential dependence on each other is not efficiently utilized. Enforcing this sequential dependence can be highly useful for guiding the ODE decoder because otherwise, it might lead to unintended attention to certain context points.

This sequential dependence of context points is even more important for irregularly sampled data, where an order-invariant encoder might lead to different representations for different timesteps sampled, even though the underlying condition is the same. This further motivates us to employ a GRU-ODE encoder to capture the underlying dynamics and thus not be sensitive to irregularity.

### 3.2 Structured Neural ODE Process (SNODEP)

Encoder with Regularly Sampled Data.
To address the above issues, we propose a modified architecture where the encoder leverages Long Short-Term Memory (LSTM) (Hochreiter, ). The LSTM encoder is designed to capture dependencies between context points across time, allowing for a more informed and contextually-aware calculation of latent distributions $L_{0}(\theta_{l_{0}})$ and $D(\theta_{d})$.
We run the LSTM backward since we want the initial value of the latent variable $l_{0}$.
Formally, the encoder takes the context sequence $\{(t_{i},y_{i})\}_{i\in\mathbb{I}_{\mathcal{C}}}$ and computes hidden representations $\{h_{i}\}_{i\in\mathbb{I}_{\mathcal{C}}}$:

$$ $\displaystyle h_{i}^{\mathrm{bwd}}$ $\displaystyle=\mathrm{LSTM}_{\mathrm{bwd}}(y_{i},h_{i+1}^{\mathrm{bwd}}),\>\> \text{for }i\in\mathbb{I}_{\mathcal{C}}.$ $$

Encoder with Irregularly Sampled Data.
Recurrent networks assume inputs to be regularly spaced and have no consideration for the actual time the input was sampled, not applicable to irregularly sampled data (Rubanova et al., ). Thus, our hidden state varies according to a GRU-ODE:

$$ $\displaystyle\hat{h}^{\mathrm{bwd}}_{i-1}=h^{\mathrm{bwd}}_{i}+\int_{t_{i}}^{t _{i-1}}\textbf{g}_{\phi}(h^{\mathrm{bwd}}(t))\,dt,\>\>h^{\mathrm{bwd}}_{i-1}= \mathrm{GRU}(y_{i},\hat{h}^{\mathrm{bwd}}_{i-1}),\>\>\text{for }i\in\mathbb{I} _{\mathcal{C}},$ $$

where $\textbf{g}_{\phi}$ is the network supposed to capture the time-dependent underlying dynamics of the hidden state, and GRU stands for the Gated Recurrent Unit (Cho, ), a gating mechanism typically employed in recurrent neural networks.
For irregular data, our encoder uses the final hidden state $h^{\mathrm{bwd}}_{0}$ to calculate the parameters of the initial latent $l_{0}$ and control $d$, which then evolves to give us the time-varying probability distributions. But in Rubanova et al. (), the $h^{\mathrm{bwd}}_{0}$ is used to get the initial latent, $l_{0}$ which then evolves directly, giving us quantities of interest and there’s no time-varying distribution involved.
For both regular and irregular scenarios, the final hidden state from the backward pass gives us the representation $r=[h_{0}^{\mathrm{bwd}}]$, which is then used to parameterize the latent distributions $L_{0}$ and $D$, via a feed-forward layer (FFW in Figure 2).

Figure: Figure 2: Illustration of the overall pipeline of the proposed SNODEP architecture.
Refer to caption: x2.png

Latent distributions.
The latent distributions, $L_{0}(\theta_{l_{0}})$ and $D(\theta_{d})$, are chosen based on the dataset. For gene-expression data, we model the latent distribution as a LogNormal distribution, to resemble the Poisson-like nature of the data:

$$ $l_{0}\sim\mathrm{LogNormal}(\mu_{L_{0}}(r),\sigma_{L_{0}}(r)),\quad d\sim \mathrm{LogNormal}(\mu_{D}(r),\sigma_{D}(r)),$ $$

whereas for metabolic-flux and balance data, we use a Gaussian distribution:

$$ $l_{0}\sim\mathcal{N}(\mu_{L_{0}}(r),\mathrm{diag}(\sigma_{L_{0}}(r))),\quad d \sim\mathcal{N}(\mu_{D}(r),\mathrm{diag}(\sigma_{D}(r))).$ $$

where $\mu_{L_{0}}$, $\sigma_{L_{0}}$, $\mu_{D}$ and $\sigma_{D}$ are learned functions. Using LogNormal ensures that we can resemble the Poisson-like nature of gene-expression data while still being able to use the re-parametrization trick (Kingma & Welling, ).

Decoder.
The decoder relies on evolving the latent variable $l(t)$ over time based on a neural ODE. For a given latent state at time $t_{0}$, the evolution is governed by:

$$ $l(t_{i})=l_{0}+\int_{t_{0}}^{t_{i}}\textbf{f}_{\theta}(l(t),d,t)\,dt,$ $$

where $\textbf{f}_{\theta}$ represents the dynamics defined by the Neural ODE, and $d$ is used for tuning the trajectory. At each target time $\{t_{i}\}_{i\in\mathbb{I}_{\mathcal{T}}}$, the latent state $l(t_{i})$ is used to determine the target distributions. For gene expressions, we model the predicted distributions as a Poisson distribution:

$$ $y_{i}\sim\text{Poisson}(\lambda_{y}(l(t)))\quad\text{for }i\in\mathbb{I}_{ \mathcal{T}}.$ $$

Whereas for metabolic flux and balances, we model the predicted distributions as a Gaussian:

$$ $y_{i}\sim\mathcal{N}(\mu_{y}(l(t)),\sigma_{y}(l(t)))\quad\text{for }i\in \mathbb{I}_{\mathcal{T}},$ $$

where $\lambda_{y}$, $\mu_{y}$ and $\sigma_{y}$ are again learned functions. The decoding distributions are meant to capture the nature of the corresponding data. The output distribution is motivated by the nature of distribution that we observe in the datasets as seen in Figure 9. During inference, we use the learned $\textbf{f}_{\theta}$ to give latent values over unseen timesteps, from $\mathcal{V}$, as well.

### 3.3 Optimization objective

Since the generative process is highly nonlinear, the true posterior is intractable. Thus, the model is trained using the amortized variational inference method using the evidence lower bound (ELBO):

$$ $\mathbb{E}_{q\left(l_{0},d\mid\mathcal{T}\right)}\left[\sum_{i\in\mathbb{I}_{ \mathcal{T}}}\log Y\left(y_{i}\mid l_{0},d,t_{i}\right)+\log\bigg{(}\frac{L_{0 }\left(l_{0}\mid\mathcal{C}\right)}{L_{0}\left(l_{0}\mid\mathcal{T}\right)} \bigg{)}+\log\bigg{(}\frac{D\left(d\mid\mathcal{C}\right)}{D\left(d\mid \mathcal{T}\right)}\bigg{)}\right],$ (2) $$

where the expectation is taken over joint latent distribution $q(l_{0},d)=L_{0}(\theta_{l_{0}})\times D(\theta_{d})$.

## 4 Experiments

### 4.1 Experimental Settings

**Table 1: Illustration of considered pathways with the number of genes, metabolites, and modules.**
| Pathway | Num Genes | Num Metabolites | Num Modules |
| --- | --- | --- | --- |
| M171 | $623$ | $70$ | $168$ |
| MHC-i | $281$ | $6$ | $9$ |
| Iron Ions | $136$ | $8$ | $15$ |
| Glucose-TCACycle | $84$ | $11$ | $15$ |

Datasets.
We use the gene-expression time-series dataset from Ori et al. (), which investigates the differentiation of human pluripotent stem cells into lung and hepatocyte progenitors using single-cell RNA sequencing to map the transcriptional changes during this process. The gene-count matrix has dimensions $10667\times 26936$, with $10667$ cells and $26936$ genes. The gene expression is counted regularly across $16$ days in batches with $\mathbb{B}_{t}$ being the index set of cells being counted on day $t$ and $|\mathbb{B}_{0}|+|\mathbb{B}_{1}|+\ldots+|\mathbb{B}_{15}|=10667$. For each cell batch $\mathbb{B}_{t}$ and a given metabolic pathway, we only consider genes responsible for encoding the metabolites from the pathway.
Table 1 summarizes the four metabolic pathways from Alghamdi et al. () we considered in this study.
Alghamdi et al. () considered the metabolic reactions from the KEGG database (Kanehisa & Goto, ), import and export reactions, and reorganized them into modules based on the topological structure. This reorganization is, in essence, the simplification of the system of reactions by coercing connected reactions into a module. Thus, when we say flux or balance, we mean it with regard to a module.

Methods.
We compare performances on model architectures, including neural process (NP) (Garnelo et al., ), neural ODE process (NODEP) (Norcliffe et al., ), and our structured neural ODE process (SNODEP). We treat NP architecture as a baseline model to get insights on modeling our problems as a differentiable random process without considering the underlying dynamics. We also compare performances between NODEP and SNODEP, which have a neural-ODE decoder, with the latter exploiting sequential relationships between the context points via its encoder.

Hyperparameters.
We vary the context length on the largest metabolic pathway, M171, to specify the hyperparameter setup for context length and train-test splits (see Appendix D).
We observe that setting the context length as $8$ had a small test-MSE, corresponding to a $80/20$ split for train and test timesteps available. Thus for the experiments below, we set our context as $\mathbb{I}_{\mathcal{C}}=\{0,1,\ldots,8\}$ and our target as $\mathbb{I}_{\mathcal{T}}=\{0,1,\cdots,12\}$, while at inference, we predict for all the timesteps $\mathbb{I}_{\mathcal{V}}=\{0,1,\ldots,15\}$. Our training input is a sample $y\sim\Pi_{t=0}^{|\mathcal{T}|}Y(\theta_{y}(t))$ with context being $y[0:|\mathcal{C}|]$ and target being $y[0:|\mathcal{T}|]$. And during testing, we sample $y\sim\Pi_{t=0}^{|\mathcal{V}|}Y(\theta_{y}(t))$.
For gene-knockout experiments, we set the train-test split of gene-knockout subsets as $80/20$.

Figure: Figure 3: Comparison of test-MSE in log-scale between NP, NODEP, and SNODEP across different metabolic pathways on ground-truth gene-expression time-series data.
Refer to caption: extracted/5937526/pictures/test_mse_comparison_gene_expression.png

Evaluation Metric.
We use the MSE loss to measure model performance in predicting unseen timesteps of time-varying distributions. Let $s\sim Y_{s}$ and $s_{*}\sim Y_{s_{*}}$, where $Y_{s}$ is the learned distribution, and $Y_{s_{*}}$ is the ground-truth distribution, where $s,s_{*}\in\mathbb{R}$. Let $\mu_{r}$ and $\mathrm{Var}(r)$ be the mean and variance for any random variable $r$. Then the Mean Squared Error (MSE) is given by:

$$ $\displaystyle\mathbb{E}\left[(\mu_{s}-s_{*})^{2}\right]=\mathbb{E}\left[\mu_{s }^{2}+{s_{*}}^{2}-2\mu_{s}s_{*}\right]=\mathrm{Var}(s_{*})+(\mu_{s}-\mu_{s_{*} })^{2}.$ $$

Assuming independence between dimensions, $\mathrm{MSE}=\sum_{i=1}^{d}\mathrm{Var}(s_{*},i)+(\mu_{s,i}-\mu_{s_{*},i})^{2}$ for $s\in\mathbb{R}^{d}$.
For gene-expression data, assume $Y_{s}=\mathrm{Poisson}(\lambda)$ and $Y_{s_{*}}=\mathrm{Poisson}(\lambda_{*})$. We thus get $\mathrm{MSE}=\sum_{i=1}^{d}\left[\lambda_{*,i}+(\lambda_{i}-\lambda_{*,i})^{2}\right]$.
Note that gene-expression data is usually very sparse (Figures 9 and 9), and hence $\lambda_{g}$ is usually very low. So in this case, minimizing MSE essentially boils down to getting as close to the Poisson approximation as possible.
For metabolic flux and balance data, suppose $Y_{s}=\mathcal{N}(\mu,{\sigma}^{2})$ and $Y_{s_{*}}=\mathcal{N}({\mu_{*}},{\sigma_{*}}^{2})$. Then, we have $\mathrm{MSE}=\sum_{i=1}^{d}\left[\sigma^{2}_{*,i}+(\mu_{i}-\mu_{*,i})^{2}\right]$.
Since we observe the estimated flux and balance are of low variances (Figures 9-9), minimizing MSE essentially boils down to bringing the model mean $\mu$ closer to ground-truth mean $\mu_{*}$.

Figure: (a) Predicting Metabolic-flux without Gene-knockout
Refer to caption: extracted/5937526/pictures/test_mse_comparison_flux.png

### 4.2 Gene-expression

Ideally, we would like to collect the ground-truth metabolic flux and balance at an individual cell or tissue level. However, this is difficult because there is very little data on them. Gene-expression counts can be considered as a rough approximation for the concentration of proteins, metabolites, and enzymes they encode since they are highly correlated. Specifically, mRNA molecules are transcribed at a certain rate from the template DNA strand, which are then translated into proteins at some rate.
Thus, we explore the timestep prediction task on log-normalized and scaled gene-expression time-series data.
Here, $Y=G(\theta_{g}(t))$, which is defined in Section 2.

From Figure 3, we can clearly see that SNODEP achieves much lower MSE across different pathways, showing the efficacy of our proposed SNODEP. Both setting the sampling distribution as Poisson and using the contextual information for the latent variables, in conjunction, help in obtaining better performance. Even though we are working with ground-truth gene expressions, this result should encourage further study on ground-truth flux datasets.

### 4.3 Metabolic-flux

Applying techniques from single-cell flux estimation analysis on the gene-expression data, we obtain samples of metabolic fluxes for the metabolic pathways. Specifically for each time $t$, given a metabolic pathway with $d$ genes and $u$ modules with gene-expression matrix $\textbf{G}_{t}\in\mathbb{R}^{d\times|\mathbb{B}_{t}|}$, we get $S^{f}_{t}(\textbf{G}_{t})=\textbf{M}^{f}_{t}$, where $\textbf{M}^{f}_{t}\in\mathbb{R}^{u\times|\mathbb{B}_{t}|}$ is the flux values for cell batch $\mathbb{B}_{t}$.
From Figure 4, we can observe that SNODEP performs generally better than the other two methods across different metabolic pathways, though the difference is not visually significant in some of them. We hypothesize that this is due to the nature of distributions as seen for some modules in Figures 9 and 9, they have low variances, and if the mean of the distributions varies in an uncomplicated manner like linear or Markovian, incorporating the context in the latent is expected not to help much.

Gene-knockout.
Gene-knockout experiments are meant to simulate the effect of disturbances in the pathway, such as the effect of any drug or environmental stress. Algorithm 1 in Appendix C describes the algorithmic form for our creation of knockout data. We model this by assuming that the gene expression level is correlated with how sensitive the metabolic pathway is with respect to the enzymes/proteins encoded by the gene. We consider $k$ most-expressed genes in the dataset and sample random subsets of these $k$ genes with the maximum cardinality of $k//2$. We call these random subsets as knockout sets where the gene expression for the genes contained is set to zero. We again calculate flux samples using scFEA (Appendix B.2) corresponding to each of knockout set, with train and test containing data corresponding to different knockout sets.
In our experiments, we set $k=20$ and the number of subsets $S=5$ for all pathways.
Figure 4 shows that our methodology is robust to gene knockout predictions, and overall, SNODEP performs better than NP and NODEP. This validates that we can use our model to predict behaviors of unseen gene knockout configurations experiments and unseen timesteps.

Figure: (a) Predicting Metabolic-balance without Gene-knockout
Refer to caption: extracted/5937526/pictures/test_mse_comparison_metabolites.png

### 4.4 Metabolic-balance

Once we get the flux values for all the modules, we can immediately obtain the change in concentration of a particular metabolite, known as the balance in flux balance analysis, by multiplying the flux with the stoichiometric matrix. We thus perform analogous experiments as in Section 4.3, where for each time $t$ for a metabolic pathway with $d$ genes and $v$ metabolites with gene-expression matrix $\textbf{G}_{t}\in\mathbb{R}^{d\times|\mathbb{B}_{t}|}$ we get $S_{t}^{b}(\textbf{G}_{t})=\textbf{M}^{b}_{t}$, with $\textbf{M}^{b}_{t}\in\mathbb{R}^{v\times|\mathbb{B}_{t}|}$ as defined in Section 2.
Figure 5 shows SNODEP generally outperforms NODEP, especially for the Iron Ions pathway. We believe that the performance is similar in pathways like MHC-i and M171 due to the simplistic nature of distributions (Figures 9 and 9), akin to what we have mentioned in Section 4.3. Similar to the previous section, we follow the steps mentioned in Algorithm 1 to get the metabolic balance samples corresponding to gene knockout, and the test MSE is shown in Figure 5. We can observe that the overall performance of SNODEP is better than that of NP and NODEP for all pathways.

### 4.5 IRREGULARLY SAMPLED TIMESTEPS

Data collection in experiments involving temporal profiling of gene expression is often performed irregularly (Rade et al., ; Nouri et al., ). Therefore, we also performed experiments where the points are irregularly sampled. To tackle the irregularity, we use GRU-ODE (Rubanova et al., ) to calculate latent distributions 3. Our context $\mathbb{I}_{\mathcal{C}}$ and target $\mathbb{I}_{\mathcal{T}}$ are similarly chosen to earlier sections, and we randomly set data from a fraction of timesteps to zero for each batch. During test time, we predict the remaining unseen timesteps. Figure 6 depicts heatmap visualizations of the difference between MSE of NODEP and SNODEP with GRU-ODE encoder. Entries in the heatmap with a positive value indicate that our SNODEP outperforms NODEP, and the higher the value is, the better the performance is. The negative values, where NODEP has a smaller test-MSE, are very low. We clearly see that SNODEP outperforms NODEP most of the time, especially towards lower frequencies, confirming the value of our model on irregularly sampled data.

Figure: Figure 6: Heatmap of test-MSE difference ($\times 10^{-2}$) between NODEP and SNODEP with GRU-ODE encoder across metabolic pathways for flux, balance, and their knockout versions. Frequency refers to the fraction of the timesteps present. In Appendix F, we provide the corresponding tables.
Refer to caption: extracted/5937526/pictures/heatmap_all_together.png

## 5 Conclusion and Future Work

In this work, we have shown how to get the time-varying metabolic flux of a system using genomics data rather than metabolomics data, which is much harder to procure. Through our framework, we intend to use the learned dynamics to generate quantities from future time steps and unseen gene-knockout configurations without any particular domain expertise. Nevertheless, we want to point out that our results with respect to flux and balance and their corresponding gene-knockout results are on data estimated via scFEA. Ideally, we would’ve preferred a gene-expression time-series that was sampled keeping metabolic pathways in mind, meaning time-series for normal conditions and several metabolic stresses, along with ground-truth metabolic flux and balance measurement. Such an experiment should also have the alternative DFBA formulation available so that we can benchmark our method with it. However, we could not find such an open-sourced dataset, so we provided our evaluations on scFEA estimated values instead of an ideal real-world dataset. Apart from such an evaluation, several future directions could be taken, like making the scFEA methods differentiable, enabling a single end-to-end differentiable pipeline, incorporating hypergraph structure into them, modifying the loss and distribution appropriately for the sparsity of gene-expression data, and exploring non-parametric probability estimations for the decoder, to name a few. We believe our work can also be helpful for integrating genomic and metabolomic data by using our pre-trained framework to fine-tune metabolomic data, for example. In conclusion, we believe our work can serve as a starting point for several interesting directions in making metabolic analysis more scalable.

## Appendix A Terminology

The key notations and terms used throughout this paper are defined below for clarity.

### A.1 Indices and Sets

- •
$t\in\{t_{1},t_{2},\ldots,t_{V}\}$: Discrete timesteps where measurements are collected, with $V$ representing the total number of timesteps.
- •
$\mathbb{B}_{t}$: The set of cells whose gene counts are measured at timestep $t$, such that the total number of cells across all timesteps is $N=\sum_{t}|\mathbb{B}_{t}|$.
- •
$\mathbb{I}_{\mathcal{C}}=\{1,\ldots,C\}$: Index set representing context points (i.e., the known timesteps used during model training).
- •
$\mathbb{I}_{\mathcal{T}}=\{1,\ldots,T\}$: Index set representing target points (i.e., the timesteps over which the model makes predictions).
- •
$\mathcal{V}$: The set of including unseen timesteps for which we aim to make predictions.

### A.2 Gene Expression Data

- •
$K$: Total number of genes measured.
- •
$N$: Total number of cells.
- •
$g_{i,t}\in\mathbb{R}^{d}$: Gene-expression vector for cell $i$ at time $t$, with $d$ representing the number of genes relevant to a particular metabolic pathway.
- •
$\textbf{G}_{t}\in\mathbb{R}^{d\times|\mathbb{B}_{t}|}$: Gene-expression matrix at time $t$ for the set of cells $\mathbb{B}_{t}$.

### A.3 Metabolic Quantities

- •
$m^{f}_{i,t}\in\mathbb{R}^{u}$: Metabolic-flux vector for cell $i$ at time $t$, with $u$ representing the number of modules.
- •
$m^{b}_{i,t}\in\mathbb{R}^{v}$: Metabolic-balance vector for cell $i$ at time $t$, with $v$ representing the number of metabolites.
- •
$\textbf{M}^{f}_{t}\in\mathbb{R}^{u\times|\mathbb{B}_{t}|}$: Matrix of metabolic-fluxes for the set of cells $\mathbb{B}_{t}$ at time $t$.
- •
$\textbf{M}^{b}_{t}\in\mathbb{R}^{v\times|\mathbb{B}_{t}|}$: Matrix of metabolic-balances for the set of cells $\mathbb{B}_{t}$ at time $t$.

### A.4 Gene-Knockout Quantities

- •
$S$: The subset of gene-knockout configurations sampled.
- •
$k$: The number of most expressed genes selected for knockout.
- •
$\tilde{g}_{i,t,s}\in\mathbb{R}^{d}$: Gene expression vector for cell $i$ at time $t$ under the $s$-th knockout configuration (where $s\in\{1,\ldots,S\}$).
- •
$\tilde{\textbf{G}}_{t,s}\in\mathbb{R}^{d\times|\mathbb{B}_{t}|}$: Gene-expression matrix at time $t$ for the set of cells $\mathbb{B}_{t}$ under the $s$-th knockout configuration.
- •
$\tilde{m}^{f}_{i,t,s}\in\mathbb{R}^{u}$: Metabolic-flux vector for cell $i$ at time $t$ under the $s$-th knockout configuration.
- •
$\tilde{m}^{b}_{i,t,s}\in\mathbb{R}^{v}$: Metabolic-balance vector for cell $i$ at time $t$ under the $s$-th knockout configuration.
- •
$\tilde{\textbf{M}}^{f}_{s,t}\in\mathbb{R}^{u\times|\mathbb{B}_{s,t}|}$: Matrix of metabolic-fluxes for the set of cells $\mathbb{B}_{t}$ at time $t$ under the $s$-th knockout configuration.
- •
$\tilde{\textbf{M}}^{b}_{s,t}\in\mathbb{R}^{v\times|\mathbb{B}_{s,t}|}$: Matrix of metabolic-balances for the set of cells $\mathbb{B}_{t}$ at time $t$ under the $s$-th knockout configuration.

### A.5 Latent Variables and Neural ODE Process

- •
$L_{0}(\theta_{l_{0}})$: Latent distribution of the initial state of the hidden representation $l_{0}$ for the cells, parameterized by $\theta_{l_{0}}$.
- •
$D(\theta_{d})$: Latent distribution of an auxiliary variable $d$, parameterized by $\theta_{d}$, used to control the trajectory of the latent variables.
- •
$l(t_{i})\in\mathbb{R}^{z}$: Latent state at time $t_{i}$ evolved from $l_{0}$ over time, where $z$ denotes the latent dimension.

### A.6 Neural Networks and Functions

- •
$\mathbf{f}_{\theta}$: Neural network modeling the evolution of the latent state $l(t)$ through a Neural ODE.
- •
$\mathbf{g}_{\phi}$: Neural network modeling the evolution of the hidden state in the irregularly sampled case.
- •
$\mu_{L_{0}}(r),\sigma_{L_{0}}(r)$: Mean and standard deviation functions for the latent distribution $L_{0}$, parameterized by the hidden representation $r$.
- •
$\mu_{D}(r),\sigma_{D}(r)$: Mean and standard deviation functions for the latent distribution $D$.
- •
$\mu_{y}(l(t)),\sigma_{y}(l(t))$: Functions parameterizing the distribution of target outputs (gene expression, metabolic flux, or balance) at time $t$, based on the evolved latent state $l(t)$.
- •
$\lambda_{y}(l(t))$: Rate parameter for the Poisson distribution used to model gene-expression data.

### A.7 Distributions and Loss Function

- •
$G(\theta_{g}(t))$: Distribution of gene expression at time $t$, parameterized by $\theta_{g}(t)$.
- •
$M^{f}(\theta_{f}(t))$: Distribution of flux at time $t$, parameterized by $\theta_{f}(t)$.
- •
$M^{b}(\theta_{b}(t))$: Distribution of balance at time $t$, parameterized by $\theta_{b}(t)$.
- •
$\tilde{M}^{f}(\theta_{f}(t))$: Distribution of flux under gene-knockout at time $t$, parameterized by $\theta_{f}(t)$.
- •
$\tilde{M}^{b}(\theta_{b}(t))$: Distribution of balance under gene-knockout at time $t$, parameterized by $\theta_{b}(t)$.
- •
$Y(\theta_{t})$: General distribution (i.e., gene expression, flux and balance, and their gene-knockout variations) at time $t$, parameterized by $\theta_{t}$.
- •
ELBO: Evidence lower bound, the objective function used for model optimization, combining the log-likelihood of observed data and the Kullback-Leibler (KL) divergence between the true and approximate posterior distributions.

## Appendix B Background

### B.1 Flux Balance Analysis

Flux balance analysis (FBA) is by far the most widely used computational method for analyzing stoichiometric-based genome-scale metabolic models. In this section, we introduce the optimization formulation of FBA and the meaning of the terms in the formulation. We refer readers to Maranas & Zomorrodi () for a more detailed description.

Static FBA.
In its most general form, FBA is formulated as an LP problem (Dantzig, ) maximizing or minimizing a linear combination of reaction fluxes subject to the conservation of mass, thermodynamic, and capacity constraints. The most widely used objective function in the FBA of metabolic networks is the maximization of the biomass reaction flux built upon the assumption that a cell is striving to maximally allocate all its available resources towards growth or maximizing the biomass, which is a predefined reaction that includes all the components required for cell growth like amino acids, nucleotides, lipids, and cofactors, etc. The formulation is as follows:

$$ $\displaystyle\mathrm{maximize}$ $\displaystyle z=\sum_{j\in\{\mathrm{biomass}\}}c_{j}v_{j}$ (3) subject to $\displaystyle\sum_{j\in J}S_{ij}v_{j}=0,\quad\forall i\in I$ $\displaystyle\mathrm{LB}_{j}\leq v_{j}\leq\mathrm{UB}_{j},\quad\forall j\in J.$ $$

Equation 3 assumes that the time constants for metabolic reactions are very small, hence a pseudo steady state constraint ${\bm{S}}\cdot{\bm{v}}=0$. Here, $z$ is the combination of reaction fluxes involved in the biomass, $v_{j}$ is the $j$-th reaction flux, $I$ is the set of all reactants, $J$ is the set of the reactions, $S_{ij}$ is the stoichiometric coefficient of reactant $i$ in reaction $j$. $\mathrm{LB}_{j}$ and $\mathrm{UB}_{j}$ are the lower and upper bounds on the rate of flux for reaction $j$, which depend on several factors such as reversibility based on Gibbs Free Energy $(\Delta G)$, type of reaction, etc. These constraint values typically are hard to determine since they need to be meticulously determined experimentally for each case.

Figure: Figure 7: Different types of GPR relationships. G, E and R denote genes, enzymes and reactions, respectively. Solid arrows are enzyme production, while dashed arrows denote catalyzing a reaction.
Refer to caption: extracted/5937526/pictures/gpr.png

Depending on the use case, the constraints and the objective change. For example, in the case of simulating gene knockouts, we set $v_{j}=0,\forall j\in J^{\mathrm{KO}}$, and gauge the gene essentiality via the GPR relationships (see Figure 7), by targeting the biomass to be less than a threshold (no cell growth) when we want to kill cells, e.g., cancer cells. In metabolic engineering, we want overproduction of a target metabolic so our objective changes accordingly, and we can try different substrates, growth conditions, and gene knockout/knockin combinations towards overproduction.

Dynamic FBA.
There are, however, processes where the time constants of the reactions are higher, like in the case of transcriptional regulation (minutes) or cellular growth (several minutes or hours) (Maranas & Zomorrodi, ). In such cases, the steady-state assumption ${\bm{S}}\cdot{\bm{v}}=0$ isn’t valid.
Mahadevan et al. () forego this steady-state assumption and study dynamic flux balance analysis in the context of Diauxic growth in E. Coli. Their study explores two separate formulations of dynamic FBA, namely (a) Dynamic Optimization-based DFBA and (b) Static Optimization-based DFBA. In (a), the optimization objective is integrated over the entire time duration via a dynamic function, while in (b), the batch is divided into intervals, and the LP is solved at the starting timestep of each interval. More precisely, (b) has the following formulation:

$$ $\displaystyle\mathrm{maximize}$ $\displaystyle\quad z(t)=\sum_{j\in\{\mathrm{biomass}\}}c_{j}v_{j}(t)$ (4) subject to $\displaystyle\quad x_{i}(t+\Delta T)\geq 0,\quad\forall i\in I$ $\displaystyle\quad v_{j}(t)\geq 0,\quad\forall j\in J$ $\displaystyle\quad\bm{\hat{c}}(\bm{z}(t))\bm{v}(t)\leq 0,\quad\forall t\in[t_{ 0},t_{f}]$ $\displaystyle\quad|v_{i}(t)-v_{i}(t-\Delta T)|\leq\dot{v_{i}}_{\mathrm{max}} \Delta T,\quad\forall t\in[t_{0},t_{f}],\quad\forall i\in I$ $\displaystyle\quad x_{i}(t+\Delta T)=x_{i}(t)+\sum_{j\in J}S_{ij}v_{j}\Delta T ,\quad\forall i\in I.$ $$

Here, $z$ is the combination of reaction fluxes involved in the biomass, $x_{i}$ is $i$-th reactant balance, $v_{j}$ is $j$-th reaction flux, $I$ is the set of all reactants, $J$ is the set of the reactions, and $S_{ij}$ is the stoichiometric coefficient of reactant $i$ in reaction $j$. In addition, $\hat{c}$ is a function representing nonlinear constraints that could arise due to consideration of kinetic expressions for fluxes, and $t_{0}$ and $t_{f}$ denote the initial and the final timestamps, respectively.

The static-DFBA formulation has much fewer parameters to solve for and is, therefore, more scalable. We would like to highlight the fact that static DFBA can be treated as a series of static FBAs that are solved locally for each timestep. In our work, for estimating metabolic flux from gene expression using techniques from Alghamdi et al. (), we thus solve for flux values for each timestep at the beginning of an interval.

### B.2 single-cell Flux Estimation Analysis

Singe-cell flux estimation analysis (scFEA) (Alghamdi et al., ) is a computational method to infer single-cell fluxome from single-cell RNA-sequencing (scRNA-seq) data. And we use it to estimate metabolic flux for the gene-expression data considered in our study.
In scFEA, they reorganize the metabolic maps extracted from the KEGG database (Kanehisa & Goto, ), transporter classification database (Saier et al., ), biosynthesis pathways, etc., into factor graphs of metabolic modules and metabolites. We use the provided genes for metabolic pathway mappings in scFEA for several organisms to estimate metabolic fluxes.

Flux estimation is a neural network-based optimization problem where the likelihood of tissue-level flux is minimized. In particular, the network iteratively minimizes the flux balance, $\mathcal{L}_{k}^{*}$ with respect to each intermediate metabolite $C_{k}$:

$$ $\displaystyle\mathcal{L}_{k}^{*}=\sum_{j=1}^{N}\bigg{(}\sum_{m\in\mathbb{F}_{ \mathrm{in}}^{C_{k}}}\mathrm{Flux}_{m}^{(j)}$ $\displaystyle-\sum_{m^{\prime}\in\mathbb{F}_{\mathrm{out}}^{C_{k}}}\mathrm{ Flux}_{m^{\prime}}^{(j)}\bigg{)}^{2}$ (5) $\displaystyle\quad+\sum_{k^{\prime}}W_{k^{\prime}}\sum_{j=1}^{N}\bigg{(}\sum_{ m\in\mathbb{F}_{\mathrm{in}}^{C_{k^{\prime}}}}\mathrm{Flux}_{m}^{(j)}-\sum_{m^ {\prime}\in\mathbb{F}_{\mathrm{out}}^{C_{k^{\prime}}}}\mathrm{Flux}_{m^{\prime }}^{(j)}\bigg{)}^{2}.$ $$

Let $\mathbb{G}^{m}_{j}=\{G^{m}_{i_{1},j},\ldots,G^{m}_{i_{m},j}\}$ be the set of genes associated with metabolic module $F_{m}$, then here
$\mathrm{Flux}_{m}^{(j)}=f_{\mathrm{nn}}^{m}(\mathbb{G}^{m}_{j}|\theta_{m})$, flux of $F_{m}$ for $j$-th cell, is modeled as a multi-layer fully connected neural network with input $\mathbb{G}^{m}_{j}$ and $\theta_{m}$ being the parameters. Here, $C_{k^{{}^{\prime}}}$ denotes the Hop-2 neighbors of $C_{k}$ in the factor graph. $\mathbb{F}_{\mathrm{in}}^{C_{k}}$ and $\mathbb{F}_{\mathrm{out}}^{C_{k}}$ are the set of modules involved in production and consumption of $C_{k}$ respectively. The optimization problem of scFEA can be thought of as finding the optimal neural network configuration that gives us reaction fluxes from gene expressions, such that the total flux regarding metabolites is minimized when considered across all tissues.

## Appendix C Gene Knockout Algorithm

Figure: Algorithm 1 Gene-knockout Flux and Balance Calculation

## Appendix D Effect of varying context length

To perform our experiments in Section 4, we need to know how many context-target timesteps our model needs to be able to predict the remaining time steps properly. For this, we used the gene-expression data of the M171 pathway since it is the largest, and for a context length $C$, our extra target length is $C//2$. In Figure 8, we see that our model is able to learn optimally after $C=6$ context steps or $C+C//2=9$ total steps. With little context, the model essentially learns next-step prediction (e.g., for $C=2\text{ or }3,C//2=1$), which does not perform well. This experiment thus validates that with enough context, our model can learn the latent dynamics and can be used to generate data from future unseen timesteps.

Figure: Figure 8: Curve plot of final test-MSE (in log scale) vs. context length.
Refer to caption: extracted/5937526/pictures/test_mse_history_context_variation_log_scale.png

## Appendix E Dataset Visualization

Figure: (a) Gene-expression: ALDOA
Refer to caption: extracted/5937526/pictures/kde_plot_ALDOA_m171.png

## Appendix F Irregularly Sampled Data

We provide the detailed results in table format of the heatmap visualization for our experiments with irregularly sampled data (Figure 6)
For each setting, the lowest value is highlighted in bold.

**Table 2: Test-MSE ($\times 10^{-2}$) across different cases and frequencies for M171 pathway.**
| Case | Frequency | NODEP | SNODEP | Diff |
| --- | --- | --- | --- | --- |
| (GRU-ODE) |  |  |  |  |
| flux (knockout) | 0.8 | 0.9434 | 0.5451 | 0.3983 |
| flux (knockout) | 0.6 | 1.6365 | 0.4743 | 1.1622 |
| flux (knockout) | 0.4 | 8.5318 | 0.3745 | 8.1573 |
| balance (knockout) | 0.8 | 0.9452 | 0.6174 | 0.3278 |
| balance (knockout) | 0.6 | 1.5794 | 0.5397 | 1.0397 |
| balance (knockout) | 0.4 | 8.5747 | 0.3985 | 8.1762 |
| flux | 0.8 | 0.0232 | 0.0285 | -0.0053 |
| flux | 0.6 | 0.0185 | 0.0245 | -0.0060 |
| flux | 0.4 | 0.0138 | 0.0188 | -0.0050 |
| balance | 0.8 | 0.0196 | 0.0255 | -0.0059 |
| balance | 0.6 | 0.0163 | 0.0202 | -0.0039 |
| balance | 0.4 | 0.0113 | 0.0160 | -0.0047 |

**Table 3: Test-MSE ($\times 10^{-2}$) across different cases and frequencies for MHC-i pathway.**
| Case | Frequency | NODEP | SNODEP | Diff |
| --- | --- | --- | --- | --- |
| (GRU-ODE) |  |  |  |  |
| flux (knockout) | 0.8 | 1.5262 | 1.4450 | 0.0812 |
| flux (knockout) | 0.6 | 1.7510 | 1.2789 | 0.4721 |
| flux (knockout) | 0.4 | 5.8514 | 0.9251 | 4.9263 |
| balance (knockout) | 0.8 | 1.7747 | 1.4393 | 0.3354 |
| balance (knockout) | 0.6 | 2.2862 | 1.2553 | 1.0309 |
| balance (knockout) | 0.4 | 3.8700 | 0.9410 | 2.9290 |
| flux | 0.8 | 1.0332 | 0.9609 | 0.0723 |
| flux | 0.6 | 0.8438 | 0.7881 | 0.0557 |
| flux | 0.4 | 0.5904 | 0.5262 | 0.0642 |
| balance | 0.8 | 0.1697 | 0.1722 | -0.0025 |
| balance | 0.6 | 0.1484 | 0.1494 | -0.0010 |
| balance | 0.4 | 0.1179 | 0.1123 | 0.0056 |

**Table 4: Test-MSE ($\times 10^{-2}$) across different cases and frequencies for Iron Ions pathway.**
| Case | Frequency | NODEP | SNODEP | Diff |
| --- | --- | --- | --- | --- |
| (GRU-ODE) |  |  |  |  |
| flux (knockout) | 0.8 | 3.0829 | 3.2716 | -0.1887 |
| flux (knockout) | 0.6 | 3.0149 | 2.6521 | 0.3628 |
| flux (knockout) | 0.4 | 2.2513 | 1.9470 | 0.3043 |
| balance (knockout) | 0.8 | 3.3664 | 3.2149 | 0.1515 |
| balance (knockout) | 0.6 | 3.0049 | 2.6651 | 0.3398 |
| balance (knockout) | 0.4 | 2.6629 | 2.0024 | 0.6605 |
| flux | 0.8 | 0.8134 | 0.7822 | 0.0312 |
| flux | 0.6 | 0.6596 | 0.6310 | 0.0286 |
| flux | 0.4 | 0.4470 | 0.4137 | 0.0333 |
| balance | 0.8 | 0.6903 | 0.6788 | 0.0115 |
| balance | 0.6 | 0.5861 | 0.5693 | 0.0168 |
| balance | 0.4 | 0.4098 | 0.3900 | 0.0198 |

**Table 5: Test-MSE ($\times 10^{-2}$) across different cases and frequencies for Glucose-TCACycle pathway.**
| Case | Frequency | NODEP | SNODEP | Diff |
| --- | --- | --- | --- | --- |
| (GRU-ODE) |  |  |  |  |
| flux (knockout) | 0.8 | 5.4157 | 5.4325 | -0.0168 |
| flux (knockout) | 0.6 | 4.6533 | 4.6676 | -0.0143 |
| flux (knockout) | 0.4 | 3.4707 | 3.5078 | -0.0371 |
| balance (knockout) | 0.8 | 5.6183 | 5.4531 | 0.1652 |
| balance (knockout) | 0.6 | 4.8713 | 4.8077 | 0.0636 |
| balance (knockout) | 0.4 | 3.5659 | 3.5082 | 0.0577 |
| flux | 0.8 | 0.6731 | 0.6308 | 0.0423 |
| flux | 0.6 | 0.5547 | 0.5210 | 0.0337 |
| flux | 0.4 | 0.3964 | 0.3741 | 0.0223 |
| balance | 0.8 | 0.3372 | 0.3329 | 0.0043 |
| balance | 0.6 | 0.2996 | 0.2881 | 0.0115 |
| balance | 0.4 | 0.2090 | 0.2113 | -0.0023 |

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: - Line 181: "the goal is to learn a model $F: t \rightarrow Y(\theta) ...$". Does the model only take time as input? My understanding is that $F$ is the decoder. Is using just time sufficient?
- What does the $Y$ variables represent? The distributions $\{G, M^f, M^b, \tilde{M}^f, \tilde{M}^b \}$? This is not entirely clear.
- If the objective is to predict entire distributions, why only consider the mean squared error (MSE) metric? Why not consider other metrics which can be specifically used to quantify the distance between distributions? For example, the Wasserstein distance and/or maximum mean discrepancy (MMD) distance. 
- What is the setup for predicting metabolic flux and metabolic balance with gene knockouts, i.e. is this for *test* gene knockouts that are seen during training, or gene-knockouts *unseen* by the model during training? This is not entirely clear.
soundness: 2
strengths: The proposed method, SNODEP, presents several key contributions that help address the challenging problem of metabolic flux and balance modelling and prediction: 

- SNODEP is able to effectively and efficiently model the dynamics of metabolic flux and balance, normally a challenging endeavour, by formulating metabolic flux and balance prediction as a data-driven problem. 
- SNODEP is evaluated over an extensive set of different pathways and experiments to showcase its utility for predicting metabolic flux and balance.
- The proposed method can learn the dynamics of metabolic flux and metabolic balance in regularly and irregularly sampled time-series settings as well as under various gene-knockout conditions.
summary: This paper addresses the challenging problem of predicting metabolic flux and metabolic balance from single-cell data. The authors formulate this prediction problem as learning the underlying dynamics of metabolic flux and metabolic balance and propose a novel method, structured neural ODE process (SNODEP), based on neural ordinary differential equations (Neural ODEs). Through extensive empirical experiments, the authors demonstrate that their novel method, SNODEP, yields improved performance on metabolic flux and metabolic balance prediction relative to baseline methods.
weaknesses: Although this work presents a novel perspective for tackling the problem of predicting metabolic flux and metabolic balance under a data-driven framework, it contains several unaddressed items that limit the overall quality and contribution of this work. For example:

- At times, the presentation and exposition of text in the manuscript is difficult to follow. For example, the learning objective in the problem formulation section, which appears to be a key section of the paper, is not very clear. Moreover, equations throughout the text don't have labels -- it would be helpful to numerically label all equations and reference them in the text. See the questions below for further details.
- Evaluation metrics are limited, only the mean squared error (MSE) metrics is considered. This provides a narrow view and comparison of model performance relative to baselines. The task appears to be to predicting over distributions of cells, in which case it is worth considering a wider range of metrics that quantify prediction performance over distributions (see questions below). 
- It feels that the novelty and methodological contribution is limited. SNODEP incorporates a variety of data-driven learning tools, including Neural ODEs, variational inference, and recurrent architectures. Given this, the primary contributions and novelty of the proposed method appear to be the application to predicting metabolic flux and metabolic balance in a data-driven way. In contrast, there is limited advancement in the individual methodological components used in this work.

---

confidence: 3
contribution: 1
initial_rating: 3
presentation: 2
questions: * I did not fully understand the inputs and outputs to the method. For example, are the gene expression values and metabolic flux and balance values treated as separate time series or does the method process them jointly?
* Could the problem be set up as a regular prediction task, with a simple LSTM or GRU architecture to predict, or are there specific problem characteristics that force the use of ODE based architectures?
soundness: 2
strengths: The specific proposed changes to the NODEP architecture look novel to me.
summary: The paper proposes a forecasting architecture based on Neural Ordinary Different Equation Processes (NODEP) which they call Strutrude Neural ODE Process (SNODEP) and applies it to the problem of forecasting gene expression and metabolic flux and balance values across time. More specifically, given a time series of single-cell gene expression data, the authors first generate metabolic flux and balance values from this data using the previously published [scFEA algorithm](https://genome.cshlp.org/content/31/10/1867.short). They do this both for the actual data as well as for synthetic genetic knockout data produced by overwriting some genes' expression value to zero and then applying the scFEA algorithm to this semi-synthetic dataset. Given a partially observed time course of this data, the task then is to predict the value at future time points.

In order to solve the task, the authors modify the original NODEP architecture. NODEP parametrizes distributions over ODEs through an encoder-decoder framework. In this framework, the authors change the latent distribution from Gaussian to Log-normal and the sampling distribution/loss function from Gaussian to Poisson to better match the nature of the gene expression data. Moreover, they replace the order-invariant encoder architecture of the original NODEP framework with time-structured networks such as Long Short-Term Memory (LSTM) and Gated Recurrent Unit (GRU).

They apply this framework to one dataset of human pluripotent stem cell differentiation and observe improved MSE performance compared to Neural Process (NP) and NODEP baselines.
weaknesses: * **Limited overall contribution:** The paper makes two contributions, introducing the metabolic system forecasting problem as a test case of interest and developing a novel architecture to solve it. In the following, I will outline why I think it does not hit the mark in either of those areas.
* **Poorly motivated metabolic system forecasting problem**
    * While understanding metabolic flux and balances is useful, the authors don't make a strong point for why the forecasting problem is specifically impactful. They consider an 80%/20% data split, so if they could show near-perfect prediction of future values, this would correspond to a 20% saving in experimental budget. Is that impactful enough in this domain for people to adopt this? Moreover, they make no attempt in characterizing how good is "good enough". The only evaluation shown is mean squared error. Are there specific scientific questions that can be solved or cannot be solved at the attained accuracy as compared to the baselines?
    * ODE systems often come with the advantage of understanding inherent dynamics of the underlying system. It seems to me like the NODEP formulation foregoes this in favor of an abstract latent space. However, even latent spaces could lend themselves to interpretations.  Unfortunately, the authors do not show any interpretation of the learned systems.
* **Narrow scope of ML contributions**
    * While SNODEP definitely seems superior to NODEP in the presented experiments, the paper leaves me wondering which of the changes (different neural network architecture, different loss function, different latent distribution) is responsible for the performance boost. This could easily be explored in ablation studies.
    * To meet the bar of a significant machine learning contribution, I would have liked to see comparisons of SNODEP on problems in the original NODEP paper and related subsequent publications, not only one very bespoke biological problem setup.
* **Limited choice of baselines:** The considered baselines are verry narrow in scope, essentially copied over from the NODEP paper. If I understand the metabolic systems prediction task correctly as a time series prediction task, what about trying simpler methods (like LSTM or GRU directly applied to the data) or time series foundation models?
* **Poor presentation of problem & results:**
    * I could not fully understand the inputs and outputs to the method (see questions below).
    * I am confused as to why all the MSE results are shown as MSE curves over training epochs. This suggests no replicate experiments were performed and thus that no uncertainty quantification on the results can be performed. It would have been simpler to summarize the final MSE values as bar plots or tables with error bars/intervals.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: - Can you please address the concern regarding the destructive measurement process of scRNAseq ? It seems different cells at different time points are directly taken as a time series, which seems inaccurate.

- I didn't fully understand why the distribution of the latent is that important, given that you end up modulating it with a non-linear function for the final predictions. Could you please motivate this choice more rigorously ? Line 260 - Latent distributions.
soundness: 1
strengths: Predicting flux and balance from metabolic systems in a temporal way is a very interesting topic and the Neural ODE formalism seems well adapted for this task. Therefore, the motivations of the paper are clear and the intended impact is evident.

The results suggest than SNODEP outperforms the considered baselines, highlighting its potential for metabolic predictions.
summary: This papers proposes a new method (SNODEP) to predict temporal flux and balance in metabolic systems. The architecture relies on a Neural ODE approach where they model the joint temporal evolution of scRNAseq, flux, and balance. The model is then trained to predict the future of the trajectories based on a contextual window. The authors state that their machine learning contributions includes modifying the emission distribution, as well as a new structured encoder approach for encoding context. The experiment section shows comparison with neural processes and neural ode processes, and shows that SNODEP performs better on scRNAseq forecasting, as well as flux and balance prediction.
weaknesses: A. The contributions of this paper are insufficient from a machine learning perspective.

1. changing the emission distribution of a neural ODE model is typically considered a design choice in all existing NODE architectures. Such that choosing a Poisson emission can hardly qualify as a contribution. 
2. I am afraid that the identified limitation of the lack of structure for the conditioning is artificial. The authors state that : "The encoded representation is calculated using context points without any particular structure in NODEP. [...] The order between the context points and their sequential dependence on each other is not efficiently utilized." This may be the case for the specific NODEP model the authors refer too but this statement is inaccurate for most NODE models such as RNN-ODE (Rubanova et al.) or GRU-ODE (De Brouwer et al.)

Given that these are the two contributions laid out by the authors, the machine learning impact is minimal. That said, that does not take away anything from their impact for biological modeling (metabolic).

B. Modeling problems.

1. The authors appear to consider that cells taken at different time points form a time series. However, they recognize that "The challenge, however, lies in that gene expression trajectories for individual cells cannot be tracked over time since cells die once their gene expression is read.". I could not find anything in the present paper that directly addresses that challenge. Instead, it seems different cells at different time points are directly taken as a time series, which seems inaccurate. At the very least, the authors should make this assumption clear and motivate it.

C. Results.

1. Most of the results are given in the form of training curves which make rapid evaluation of the performance difficult. The scales of some subplot is sometimes off, such that the different methods cannot be distinguished.
2. I am pretty sure that RNN-ODE and GRU-ODE could also be used for this problem. I would encourage the authors to expand their list of baselines. I would also include very simple baselines like RNNs - or even one-step-ahead linear regressors.

---

confidence: 2
contribution: 1
initial_rating: 3
presentation: 2
questions: None
soundness: 2
strengths: - The paper is well-organized and well-motivated by a biological question.

- The empirical assessment clearly demonstrate the superiority of the proposed method.
summary: In this submission, the author propose a data-driven solution to Dynamic Flux Balance Analysis (DFBA), a classical tool from Systems Biology to analyze genome-scale metabolic network. At each timestep, DFBA solves a time-varying constrained optimization problem where the goal is to determine the value of the fluxes that lead to maximal biomass production, under time-varying biological constraints, which typically can only be brought by a  deep understanding of the system at hand.

Leveraging gene expression data in both control and knockout experiments thanks to single-cell flux estimation analysis, the authors aim at predicting time-resolved unseen gene expression, flux, balance, in both control and knockout experiments. This is achieved by means of Neural ODE processes (NODEP). Particularly, the main contribution of the paper is to cast DFBA in the framework of NODEP, and to highlight the need for a modification of NODEP to better handle the 1) irregularly sampled data and 2) non-gaussianity of certain distributions involved. This leads to Structural NODEP (SNODEP).

Finally, the authors carry a number of experiments on different metabolic pathways, demonstrating the clear improvement provided by SNODEP, particularly for sparsely sampled data across time.
weaknesses: - There is little to no technical novelty. The paper feels like the application of Neural ODE processes to DFBA, with minor tweaks to make it work better. Although the results are convincing, and there is definitely value in framing the DFBA problem in terms of Neural ODE processes, I am not sure that ICLR is a relevant venue for this work.

Meta-review: The reviewers and the authors engaged in a good discussion. Although the reviewers appreciate the problem tackled in the paper (neural ODEs for metabolic networks), the reviewers were not convinced that the paper is ready for ICLR.

The authors are encouraged to continue their work and submit the updated version of the paper elsewhere.

Decision: Reject