Neural Continuous-Time Markov Models

Continuous-time Markov chains (CTMCs) are stochastic processes that model a variety of different systems, e.g., chemical reaction networks, population dynamics, gene regulatory networks, and infectious disease dynamics. CTMC sample paths are piecewise-constant functions that take values in the state space; for a given sample path, each jump discontinuity corresponds to a time at which the system transitions from one state to the next. On a discrete state space $\Omega$, a CTMC is completely characterized by its transition rates, an array $\Lambda=\{\lambda_{i,j}\}_{i,j\in\Omega}$ that describes the rate at which the system transitions from one state to another. Given an initial condition, which can be expressed as a distribution over the possible states, one can use the transition rates to generate sample paths of the system. However, the inverse problem is not as straightforward. Learning the transition rates from sample path data is an active field of research.

Our focus in this work is on a particular class of CTMCs that model the stochastic dynamics of reaction networks [1]. For this class of CTMCs, there has been a wealth of prior work on estimation and identification. In the structure identification problem, the goal is to estimate the reaction network itself. Here we see approaches that focus on graphical modeling [2], sparse identification [3, 4, 5], and subnetwork reconstruction [6].

Our work has more in common with the parameter estimation problem, where all reactions and corresponding transition rates are known up to a finite set of parameters, which one seeks to estimate from data. Various prior approaches can then be distinguished based on methodology and assumptions regarding the data. In situations where one assumes access only to steady-state observations, we find both Bayesian [7] and moment-based methods [8]. In the case where we assume access to inexact measurements of all chemical species, we see moment-based methods [8], variational Bayes methods [9], Gaussian adaptation methods [10], and expectation maximization methods, either for general systems [11], or for the particular case of birth-death processes [12]. There also exist filtering and Bayesian methods to estimate states and parameters from exact measurements of a subset of all chemical species in a given system [13], and adjoint methods for reaction-diffusion systems [14].

In this paper, we model CTMC transition rates $\Lambda$ with functions that we parameterize and estimate using neural networks. We assume that all reactions are known, and that trajectories are fully observed. Unlike prior parameter estimation studies we have seen, we assume nothing regarding the functional form of the transition rates (or, equivalently, propensity functions). The neural CTMC (N-CTMC) method we describe allows for transition rates with arbitrary kinetics and dependence on external covariates, such as temperature.

Recent work has introduced covariate dependence via generalized linear models (GLM): for Markov chains in discrete/continuous-time, transition probabilities/rates are modeled as a softmax/exponential function of a linear function of the covariates [15, 16]. In this paper, on tests with synthetic data, we show that the N-CTMC method can model kinetics that lie outside mass-action or Michaelis-Menten models. This goes beyond the capabilities of GLM models.

The likelihood function is often used to learn the transition rate matrix of a CTMC. The likelihood measures how probable a sample path is to occur given a transition rate matrix. To find CTMC transition rates that best fit one or more given sample paths, we solve for the transition rates that maximize the likelihood function. For stochastic reaction networks with even a small number of reactions, the state space $\Omega$ can be large or infinite, rendering the likelihood intractable. Strategies to deal with the intractable likelihood include truncating the state space [17], or using likelihood-free approaches such as approximate Bayesian computation [18, 19]. Truncating the state space introduces error, and does not allow for the full model to be learned.

Instead of attempting to estimate a large or infinite-dimensional matrix $\Lambda$, N-CTMC uses the current state $i$ as one of the inputs to the neural network. The neural network then outputs a vector of transition rates, one for each reaction. In a setting with $s$ species of interest, every unique integer $s$-tuple $(A_{1},\ldots,A_{s})$ denotes a state. If each $A_{i}\in[0,N]$, the state space has maximal dimension $N^{s}$. If we were to estimate each entry of $\Lambda$ individually, we would have to estimate $O(N^{2s})$ constant parameters. In contrast, the number of parameters in the neural network model is independent of $N$.

After deriving and explaining the method, we demonstrate its capabilities using tests with synthetic data. In each test, we start with sample paths generated from a known reaction network. We demonstrate that after training on these sample paths, N-CTMC succeeds in learning transition rates in close agreeement with the ground truth. We explain and show how N-CTMC significantly reduces the computational cost of learning transition rates. As we increase the sizes of our training sets, we show that the error of N-CTMC’s estimates decreases. Finally, we show that N-CTMC significantly outperforms a GLM of a similar construction.

We now detail the neural CTMC (N-CTMC) method. For a stochastic reaction network ($RN$) as described in Section 2.1, assume we have complete observations: ${\mathcal{O}}=({\mathbf{S}}(t_{0}),{\mathbf{C}}_{0},t_{0}),\ldots({\mathbf{S}}(t_{N}),{\mathbf{C}}_{N},t_{N})$. Specifically, at time $t_{i}$, we have ${\mathbf{S}}(t_{i})\in{\mathcal{R}}^{s}$, a vector of counts (of all species), and ${\mathbf{C}}_{i}\in{\mathcal{R}}^{c}$, the associated covariates. Let ${\mathbf{x}}_{i}=[{\mathbf{S}}(t_{i})^{T},{\mathbf{C}}_{i}^{T}]\in{\mathcal{R}}^{s+c}$.

We look to model the propensity functions ${\boldsymbol{\alpha}}({\mathbf{x}})\in{\mathcal{R}}^{r}$, where $r$ is the number of reactions in the reaction network ($RN$). In this formulation, we assume no knowledge of the reaction rates $k$.

We now formulate the likelihood function for the sequence of 3-tuples ${\mathcal{O}}$. In this likelihood function, we replace the constant transition rate matrix $\Lambda$ with a vector of propensity function ${\boldsymbol{\alpha}}({\mathbf{x}})$. To model ${\boldsymbol{\alpha}}({\mathbf{x}})$ we use a neural network, a universal function approximator, with parameters $\theta$. The network takes as input ${\mathbf{x}}$ and outputs the propensity ${\boldsymbol{\alpha}}({\mathbf{x}};\theta)$. In contrast with Section 2.1, ${\boldsymbol{\alpha}}({\mathbf{x}};\theta)$ depends not only on the state ${\mathbf{S}}$ but also on the associated covariates ${\mathbf{C}}$. Taking the negative $\log$ of (1) gives

In the data ${\mathcal{O}}$, the transition $O_{i}\to O_{i+1}$ equates to the change in count vector ${\mathbf{S}}(t_{i})\rightarrow{\mathbf{S}}(t_{i+1})$; let $\rho_{i}\in RN$ denote the reaction to which this change corresponds. Then we can replace $\lambda_{O_{i}O_{i+1}}$ by the $\rho_{i}$-th element of the neural network’s output, which we denote by $\alpha({\mathbf{x}}_{i};\theta)_{\rho_{i}}$. Similarly, we can replace $\lambda_{O_{i}O_{i}}$, the rate of leaving state $O_{i}$, by $\sum_{j=1}^{r}\alpha({\mathbf{x}}_{i};\theta)_{j}$, the sum of the propensity functions for each reaction in ($RN$). Rewritten in terms of ${\mathbf{x}}$ and ${\boldsymbol{\alpha}}({\mathbf{x}};\theta)$, the negative log likelihood (3) becomes

Let $L^{\rho}$ be the number of times reaction $\rho$ occurs. Suppose we collect all row vectors ${\mathbf{x}}_{i}$ that correspond to states just before reaction $\rho$ occurs. We stack these vertically into an $L^{\rho}\times(s+c)$ matrix $X^{\rho}$. Similarly, let ${\mathbf{T}}^{\rho}\in{\mathcal{R}}^{\,L^{\rho}}$ be a vector of time spent in state before reaction $\rho$ occurs. With this, we rewrite the above negative log likelihood as

To train the N-CTMC model, we apply gradient descent to minimize this negative log likelihood. Note that the internal sum over $L^{\rho}$ does not require a for loop in our implementation. With neural network input vectorization, we can pass all of $X^{\rho}$ into ${\boldsymbol{\alpha}}$ at once. The $l$-th row of the output will be identical to the output we would have obtained had we passed in only $X^{\rho}_{l}$, i.e., $\alpha({X}^{\rho};\theta)_{lj}=\alpha({X}^{\rho}_{l};\theta)_{j}$. Using this, the internal sum can be accomplished entirely with matrix algebra, which is highly optimized in modern neural network frameworks. The only remaining loop corresponds to the sum over $\rho$ from $1$ to $r$; this should be contrasted with sums over $O(|\Omega|^{2})$ elements that would have resulted from naïve fusion of neural networks with the methods from Section 2.2. Together with automatic differentiation, our formulation leads to more efficient computation of $\nabla_{\theta}\log L(\theta)$ and resulting minimization of the loss (4).

Note that the above N-CTMC approach learns a model ${\boldsymbol{\alpha}}({\mathbf{x}};\theta)$ without direct observations of ground truth propensity functions. We observe the propensities only indirectly through their effect on the system, as evidenced in the data ${\mathcal{O}}$. This varies from standard neural network training approaches where one starts with ground truth observations of a response variable and then minimizes the mean squared error beween predicted and true values of the response.

Here we briefly sketch how one might prove that neural networks approximations can converge to true propensity functions. For brevity, let ${\boldsymbol{\alpha}}={\boldsymbol{\alpha}}({\mathbf{x}})$. Let $\widehat{\boldsymbol{\alpha}}^{\theta}$ and $\widehat{\boldsymbol{\alpha}}^{N}$ denote propensity functions estimated by, respectively, a neural network and maximum likelihood estimation as in (2). More specifically, we begin by binning the covariate space ${\mathcal{R}}^{c}$. For a given bin ${\mathcal{B}}$, using transitions ${\mathbf{S}}(t_{i})\rightarrow{\mathbf{S}}(t_{i+1})$ from ${\mathcal{O}}$ such that ${\mathbf{C}}_{i}\in{\mathcal{B}}$, we form MLE estimates as in (2). Now given a state ${\mathbf{S}}$ and covariates ${\mathbf{C}}$, we find the bin corresponding to ${\mathbf{C}}$ and evaluate the MLE estimates (2) in that bin associated with all transitions of the form ${\mathbf{S}}\rightarrow{\mathbf{S}}^{\prime}$. Organizing these estimates by reaction number, we form $\widehat{\boldsymbol{\alpha}}^{N}$. Let us assume that MLE consistency results for covariate-free CTMCs [20] can be extended to include covariates, assuming conditions on the distribution of covariates ${\mathbf{C}}_{i}$ in the data ${\mathcal{O}}$ and on the ergodicity of the CTMC, such that as $N\to\infty$, we will have $\|\widehat{\boldsymbol{\alpha}}^{N}-{\boldsymbol{\alpha}}\|_{\infty}\to 0$ almost surely (a.s.). Next, treating $\widehat{\boldsymbol{\alpha}}^{N}$ as ground truth, we train a neural network with suitable activation functions, width $W$ and depth $D$; as $D\to\infty$, we have $\|\widehat{{\boldsymbol{\alpha}}}^{\theta}-\widehat{{\boldsymbol{\alpha}}}^{N}\|_{\infty}\to 0$ [21]. Putting everything together, we obtain $\|{\boldsymbol{\alpha}}^{\theta}-{\boldsymbol{\alpha}}\|_{\infty}\to 0$ a.s. as both $N,D\to\infty$.

To evaluate the N-CTMC method, we will use several metrics: the mean absolute error (MAE) and mean squared error (MSE) across each unique state in ${\mathcal{O}}$, together with the weighted mean absolute error (W-MAE) and weighted mean squared error (W-MSE). The W-MAE and W-MSE examines the prevalence of each unique state in ${\mathcal{O}}$ and weights the contribution to the overall absolute error accordingly. This means that a state that is rarely visited in ${\mathcal{O}}$ will have less effect on the W-MAE and W-MSE than on the MAE and MSE.

We examine three different applications of the N-CTMC: a birth-death process, Lotka-Volterra population dynamics, and a temperature-dependent chemical reaction network. Each application is chosen to test the N-CTMC model in a specific way.

In the birth-death process, we use a single covariate and let propensity functions depend solely on this covariate and not on the current population count. While this model lacks realism, it allows us to compare the N-CTMC to the counting-based MLE (2) for calculating rate parameters described in Section 2.2. To test the N-CTMC’s ability to capture arbitrary kinetics, we turn to a Lotka-Volterra population model. Here there are no covariates, but the propensity functions do not follow the law of mass action. The final reaction network we consider tests the N-CTMC’s ability to model arbitrary dependence on covariates. Here the system of reactions follows the law of mass action, but the propensities depend on temperature nonlinearly.

For each system, we generate trajectories using the Gillespie algorithm [22] and then use N-CTMC to learn propensity functions (equivalently, transition rates). We model all systems with an increasing amount of training data to show empirical model convergence.

We have assumed access to complete observations of trajectories together with knowledge of the number of reactions and associated stoichiometric coefficients. Subject to these assumptions, we have demonstrated that the N-CTMC method learns consistent transition rates for stochastic reaction networks with both mass action and non-mass action kinetics, including for one system with more than 13,000 states.

Unlike other methods where the likelihood is truncated, N-CTMC uses the entire state space in the likelihood calculation. The N-CTMC estimated propensity function can be evaluated on any state in $\Omega$, including but not limited to states visited in the training set. We can only guarantee accurate estimates of propensity functions on states that have been visited sufficiently often in the training data. On rarely visited states, however, N-CTMC outperformed a GLM approach in which propensity functions depend linearly on inputs (count vectors and/or covariates); such GLM approaches have been explored in prior work [15, 16]. The N-CTMC also models transition rates which vary as arbitrary functions of covariates, allowing for study of systems where propensity functions depend on external stimuli.

In the current work, we focused on systems with exponentially distributed sojourn times, an unreasonable assumption for some systems [25]. Using the framework developed here, the negative log likelihood (NLL) can be easily modified to accommodate a semi-Markov process [26]. For the birth-death example, we showed that N-CTMC can also be used to control—in an open-loop fashion—covariate-dependent CTMC systems. We hope to explore both semi-Markov versions and control applications of N-CTMC in future work. This will expand the class of problems to which we can apply our methods.