Bayesian inference of Stochastic reaction networks using Multifidelity Sequential Tempered Markov Chain Monte Carlo

A reaction network consists of $N$ different chemical species $S_{1},\ldots,S_{N}$ that are interacting via the following $M$ chemical reactions

We are interested in keeping track of the integral vectors $\bm{x}\equiv(x_{1},\dots,x_{N})^{T}$, where $x_{i}$ is the population of the $i$th species. Assuming constant temperature and volume, the time-evolution of this system can be modeled by a continuous-time, discrete-space Markov process [9]. The $j$th reaction channel is associated with a stoichiometric vector $\bm{\nu}_{j}=(\nu_{1,j}^{prod}-\nu_{1,j}^{react},\ldots,\nu_{N,j}^{prod}-\nu_{N,j}^{react})^{T}$ ($j=1,\ldots,M$) such that, if the system is in state $\bm{x}$ and reaction $j$ occurs, the system transitions to state $\bm{x}+\bm{\nu}_{j}$. Given $\bm{x}(t)=\bm{x}$, the propensity $\alpha_{j}(\bm{x};\theta)dt$ determines the probability that reaction $j$ occurs in the next infinitesimal time interval $[t,t+dt)$, where $\theta$ is the vector of model parameters. In other words,

An important case of reaction networks are those that follow mass-action kinetics, whose propensity functions take the form

In this formulation, $c_{j}(\theta)$ is the probability per unit time that a combination of molecules can react via reaction $j$, and the remaining factor is the number of ways the existing molecules can be combined to form the left side of the chemical equation (1).

The time-evolution of the probability distribution of this Markov process is the solution of the linear system of differential equations known as the chemical master equation (CME)

where $\bm{p}(t)$ is the time-dependent probability distribution of all states, $p(t,\bm{x})=\textrm{Prob}\{\bm{x}(t)=\bm{x}_{i}|\bm{x}(0)\}$. The initial distribution $\bm{p}_{0}$ is assumed to be given, and $\bm{A}(\theta)$ is the infinitesimal generator of the Markov process, defined entry-wise as

Here, we have made explicit the dependence of $\bm{A}$ on the model parameter vector $\theta$, which we need to infer from experimental data.

Typically, reaction networks model open biochemical systems, where the set of all possible molecular states is unbounded. This makes the CME an infinite-dimensional linear system of ODEs. The finite state projection (FSP) is a well-known strategy to systematically reduce this linear system into a finite surrogate model with a strict error bound.

The FSP can be thought of as a special class of projection-based model reduction applied to the CME. Specifically, let $\Omega$ be a finite subset of the CME state space. The projection of the CME operator $\bm{A}$ onto the subspace spanned by the point-mass measures $\{\delta_{\bm{x}}|\bm{x}\in\Omega\}$ is given by

We can then define a reduced model of the dynamical system (3) based on this projection as

Clearly, in solving (6) we only need to keep track of the equations corresponding to states in the finite set $\Omega$, which is amenable to numerical treatments.

In contrast to generic projection methods, the gap between the reduced-order model and the true CME can be computed for the FSP. Indeed, Munsky and Khammash  [42, Theorem 2] proved that the truncation error can be quantified in $\ell_{1}$ norm as

Clearly, the right hand side can be readily computed from the solution of the reduced system (6). From this precise error quantification, we have effective iterative method for solving the CME. Choosing an error tolerance $\varepsilon>0$, starting from any initial set $\Omega:=\Omega_{0}$, we solve system (6) and check that the right hand side of (7) is less than $\varepsilon$. If this fails, we add more states to $\Omega_{0}$ to get a strictly larger set $\Omega_{1}$ and repeat the procedure until we find an approximation that satisfies our error tolerance.

As the sequence of subsets $\Omega_{i}$ grows until it eventually covers the whole state space, we might expect that the finite-time solution of (6) will likewise converge to the true solution. This is indeed the case for all models in practice, with only a few theoretical counterexamples in which the Markov chain is explosive [10]. Sufficient conditions for the convergence of the FSP can be checked based on the form of the propensity function [43, 44]. In practice, reaction networks tend to have only reactions between two or fewer molecules, with propensity functions in mass-action form (2), and these are guaranteed to be approximable with the FSP [43].

For the rest of the paper, we only concern ourselves with non-explosive SRNs where the FSP converges. Given such models, any exhaustive sequence of subsets $\{\Omega_{j}\}$ suffices to guarantee that the FSP solution eventually satisfies any prespecified error tolerance. However, some choices are more effecient than other, and it is also more advantageous to partition the time interval into smaller timesteps and use a smaller $\Omega_{j}$ on each step [45]. We will return to these observations in section 4.3.

A final point to make about the FSP is that, if the sequence of reduced sets $\Omega_{j}$ are increasing, that is, $\Omega_{j}\subset\Omega_{j+1},j=1,2,\ldots$, the resulting truncation error has been shown to decrease monotonically [42, Theorem 1] (also see [46, Theorem 2.5]). This gives us a natural way to form a hierarchy of reduced models for the CME, and we will return to this in section 4.1. With the computed distribution of the SRN in our hands, we can now match them directly to experimental data using a straightforward likelihood function.

We are interested in inferring the parameters for the reaction network from discrete, single-cell datasets [47, 6, 3, 7, 8] that consist of several snapshots of many independent cells taken at discrete times $t_{1},\ldots,t_{T}$. The snapshot at time $t_{i}$ records gene expression in $n_{i}$ cells, each of which can be collected in the data vector $\bm{c}_{j,i},\;j=1,\ldots,n_{i}$ of molecular populations in cell $j$ at time $t_{i}$.

Assume that a model class $\mathcal{M}=\left\{M(\theta)\right\}_{\theta\in\Theta}$ of stochastic reaction networks has been chosen to model the data consisting of a fixed set of reactions with unknown reaction parameters $\theta$. Let $p(t,\bm{x}|M(\theta))$ denote the entry of the CME solution corresponding to state $\bm{x}$ at time $t$, given by SRN model $M(\theta)$. The log-likelihood of the dataset $\mathcal{D}$ given $M(\theta)$ is given by

A common approach to fitting this model is find model parameters $\theta$ that maximize the log-likelihood function. These parameters would be those that best fit the data, but this process does not capture any uncertainty about those parameters. There are potentially many other models that could fit the data approximately as well as the maximum likelihood model. In contrast, with a Bayesian approach we seek to quantify this parameter uncertainty so that we can be confident in the inference results, understand the influence of this uncertainty on future predictions, and design future experiments to reduce the parameter uncertainty.

Bayesian inference is rooted in the Bayesian philosophy of probability in which our uncertainty about the world is modeled using probability distributions [48]. As information becomes available, we update these distributions to reflect our new state of understanding about the world. Therefore, for Bayesian inference we begin with a prior distribution, $p\left(\theta\right)$, that captures our initial beliefs about model parameters. Then after data has been observed, the likelihood of the data given a model class and associated parameters can be found as $p\left(\mathcal{D}\mid\theta\right)=\exp({L}(\mathcal{D}|\theta))$. By applying Bayes’ Theorem, we can combine the prior and likelihood information together to construct the posterior distribution on model parameters that reflects our updated beliefs after data has been observed:

Here, $p\left(\mathcal{D}\right)$ is a normalization constant known as the model evidence. This is relevant for Bayesian model selection problems but is not required for parameter calibration. When $p\left(\theta\right)$ is a constant, the parameters that maximize the posterior density are equivalent to the maximum likelihood estimator. Solving the Bayesian inference problem to identify parameters allows us to quantify our uncertainty regarding the accuracy of the parameter fit by sampling from the posterior distribution. However, it can be a computationally challenging problem as discussed in the next subsection.

The Bayesian framework also provides a criteria to select or weigh different model classes as data become available. Suppose, instead of a single model class, we are given $K$ possible network structures that could potentially explain the observations. Let $\mathcal{M}^{k}=\{M^{k}(\theta^{k})\}_{\theta^{k}\in\Theta^{k}}$ denote the $k$-th class, where the parameter domains $\Theta^{k}$ need not have the same dimensionality. Each model class is associated with a prior weight $P(\mathcal{M}^{k})$ that represents the prior level of belief in each class. If $\mathcal{D}$ denotes the dataset as before, we can compute the model evidence of $\mathcal{M}^{k}$ as

These probabilities then reflect the posterior weighting of the different model classes that can be used to make average predictions over the models. Similarly, the Bayes factors

allow us to rank and compare the models based on how strongly they are supported by the current data for Bayesian model selection. The drawback of Bayesian model selection in the context of stochastic gene expression is the computational cost of computing the model evidences. We will return to this issue in section 3 where we show how the Mutlfidelity ST-MCMC framework provides an efficient way to estimate these evidences.

Markov Chain Monte Carlo algorithms are widely used for sampling the posterior distribution of a Bayesian inference problem. These methods design a Markov chain whose stationary distribution is the target posterior distribution, $p\left(\theta\mid\mathcal{D}\right)$. Therefore, by simulating the evolution of samples, $\theta$, according to the Markov chain, samples are asymptotically drawn from the posterior distribution. However, unlike Monte Carlo sampling, these samples are correlated. A common MCMC method is the Metropolis-Hastings algorithm. The algorithm begins by initializing the parameter state to some $\theta_{0}$. Then at a step $i+1$ in the evolution, a candidate sample $\theta^{\prime}$ is drawn according to a proposal distribution $Q\left(\theta^{\prime}\mid\theta_{i}\right)$. This candidate is then accepted or rejected with probability $\alpha\left(\theta^{\prime}\mid\theta_{i}\right)$ given by

A common method to judge whether sufficient posterior samples have been generated using the MCMC sampler is the notion of effective sample size (ESS). The ESS of a $N$ sample correlated population $\theta_{i=1...N}$ corresponds to the number of independent samples which would estimate a quantity of interest, $\hat{q}=E\left[q\left(\theta\right)\right]$, with the same variance as the correlated sample population. Therefore, the ESS of a quantify of interest can be computed as

Based upon the Markov chain Central Limit Theorem, for correlated samples generated according to the stationary distribution of a Markov process, the ESS can be approximated as

To overcome many of the challenges associated with a standard Metropolis-Hastings based MCMC method, parallel methods, like Sequential Monte Carlo (SMC), have been introduced to better leverage high performance computing resources. SMC methods for Bayesian inference transport a sample population, initially distributed so that it can approximate expectations with respect to the prior, to one which can approximate posterior expectations [30, 28, 31, 32]. Typically, this means a population of samples initially distributed according to the prior being transformed into a population of samples approximately distributed according to the posterior. For Sequential Tempered MCMC (ST-MCMC) [32], we break down the inference problem into a series of annealing levels $i$ defined by an annealing factor $\beta_{i}\in\left[0,1\right]$. Each level defines an intermediate distribution, $\pi_{\beta_{i}}\left(\theta\right)$, which we would like to use to generate samples and to compute expectations. These intermediate distributions take the form of

This annealing approach is common to many SMC methods used for Bayesian inference and can be thought of as gradually integrating the influence of the data into the solution. To simplify the problem of transporting samples from the prior, $\pi_{0}\left(\theta\right)$, to posterior, $\pi_{1}\left(\theta\right)$, we transport samples sequentially through each level in the sequence, i.e. $\pi_{\beta_{i}}\left(\theta\right)$ to $\pi_{\beta_{i+1}}\left(\theta\right)$ with $\beta_{i+1}>\beta_{i}$. Because we can control the size of the jump $\Delta\beta=\beta_{i+1}-\beta_{i}$, we can ensure that this change is not too drastic as to cause poor approximation of the true distribution i.e. too drastic a decrease in the ESS. Transporting samples is done in three steps:

1. 1.

   Re-weight the previous sample population, distributed according to $\pi_{\beta_{i}}\left(\theta\right)$, with unnormalized weights $w_{i}=p\left(\mathcal{D}\mid\theta_{i}\right)^{\Delta\beta}$ to reflect expectations with respect to the new distribution $\pi_{\beta_{i+1}}\left(\theta\right)$.

2. 2.

   Re-sample the population according to the weights so that the samples now reflect $\pi_{\beta_{i+1}}\left(\theta\right)$.

3. 3.

   Seed a Markov chain starting at each sample, and then use MCMC to explore $\pi_{\beta_{i+1}}\left(\theta\right)$.

The MCMC step is essential to ensure the sample population does not degenerate since the re-weighting and re-sampling steps reduce the ESS of the population. MCMC increases the ESS because it decorrelates the seeds and explores the target distribution, causing samples to better reflect it. Typically, $\Delta\beta$ is chosen adaptively to not decrease the ESS too much during the update. This is achieved by finding a $\Delta\beta$ such that the coefficient of variation (COV) of the sample weights equals a target $\kappa$. The COV approximates the ESS by $N_{ESS}\approx\frac{N}{1+\kappa^{2}}$. Therefore, we find a $\Delta\beta>0$ that solves the equation

One approach to choosing the appropriate annealing factor and model fidelity is a combined likelihood tempering and model bridging scheme discussed in Latz et al. [38]. This scheme is based upon the ESS statistic discussed in Section 2.5. Within Multifidelity ST-MCMC, at every level $l$ of Algorithm 1, we choose whether to temper by changing $\beta_{l}=\beta_{I-1}+\Delta\beta$ or to bridge by changing the model fidelity, $m_{l}=m_{l-1}+1$. This choice is made by measuring the ESS of the sample population with respect to the next change in the model fidelity by computing the unnormalized weights as if we were to bridge:

We can then compute the coefficient of variation of the weights, $w_{i}$, to determine if it exceeds a target $\kappa$. If it does, we choose to bridge to the next model fidelity because the sample population is beginning to degenerate, so it no longer has sufficient ESS to approximate the next level intermediate posterior. If the COV is less than $\kappa$, we choose to keep the current model but instead temper $\beta$. The next beta is chosen using the same strategy as before by solving Equation 17.

We introduce a new criteria for model fidelity selection based on information theory. This criteria is motivated by the fact that the ESS-based strategy described above only decides to change fidelity based upon the next model in the hierarchy. This means the sampler may continue to use a low fidelity model because it still meets the ESS criteria with respect to the next model, even when it is drifting away from the high fidelity posterior. Instead, we introduce a method that utilizes a limited number of full fidelity model evaluations to help us better decide when to bridge model fidelity. Depending on the computational cost of the full fidelity simulations, the improved bridging strategy and the improved robustness of this method may outweigh the cost of these full fidelity solutions.

Within multifidelity ST-MCMC, if the algorithm is at annealing level $l$ with annealing factor $\beta_{l}\in\left[0,1\right]$ and has been sampling the intermediate posterior defined by a model, $M_{m_{l}}$, in a model hierarchy $\mathcal{M}=\{M_{j}:j=1\dots K\}$, we would like to know whether $M_{m_{l}}$ still provides meaningful information about the ultimate posterior once we move to level $l+1$ with annealing factor $\beta_{l+1}$. Here we assume the ultimate posterior is $p\left(\theta\mid\mathcal{D},M_{K}\right)$, where $M_{K}$ is the highest fidelity model. Therefore, unlike the previous ESS-based method, we begin by proposing a tempering step under the assumption that the current model fidelity is valid. We find the proposed $\beta_{l+1}$ by solving Equation 17.

If $M_{m_{l}}$ no longer provides meaningful information at the next level, we use the next highest fidelity model in the algorithm, $M_{m_{l}+1}$. This criteria can be formulated using a generalization of information theory [49], where the information gained about the full posterior, $p\left(\theta\mid\mathcal{D},M_{K}\right)$, by moving from level $l$ to $l+1$ with model $M_{m_{l}}$ is:

If we choose to update the model fidelity, we consider two strategies for choosing $\beta_{l+1}$ for the next level. In the first strategy, keeping with the ESS-based tempering and bridging framework from above, is to set $\beta_{l+1}=\beta_{l}$. The second strategy is to tune $\beta_{l+1}$ to try attain an ESS target. The first strategy is often more computationally efficient, but may not be as robust if changing model fidelity introduces significant variations. To tune $\beta_{l+1}$, we first define the importance weight for transitioning from a level defined by $\beta_{i}$ and $m_{l}$ to a level defined by $\beta_{l+1}$ and $m_{l+1}=m_{l}+1$ as

Using the same approach as before, we can then tune $\beta_{l+1}$ to meet some ESS target based upon the COV of the weights. However, unlike in previous problems, this might not be achievable. If Equation (17) has a solution, we chose the largest $\beta_{l+1}$ such that the COV target is met. If Equation (17) does not have a solution, we find the $\beta_{l+1}$ that minimizes the COV and thus maximizes the ESS.

Since computing the information in Equation 19 requires marginalizing over the posterior, it can be challenging. However, this computation can be approximated using the samples from ST-MCMC. The first step is to recognize the connection between computing this criteria and estimating the model evidence:

Since we cannot yet sample $p\left(\theta\mid\mathcal{D},M_{K}\right)$ we use importance sampling to express this integral in terms of the level $l$ distribution, which we have samples for:

SMC and ST-MCMC methods not only enable robust solutions of Bayesian inference problems for parameter calibration, but also enable Bayesian model selection by providing asymptotically unbiased estimates of the model evidence. Model evidence estimates are generally highly computationally expensive since they require estimating the normalization constant,

which consists of marginalizing the likelihood over the prior distribution. If the high probability content of the prior differs significantly from the most likely parameters according to the likelihood, it is difficult to estimate this integral using Monte Carlo samples. Instead, SMC type methods break down this estimate into a series of Monte Carlo approximations over the intermediate distribution levels previously discussed. As such, a hierarchy of multifidelity models can also be used to accelerate this estimate within the Multifidelity ST-MCMC framework. Using the methods described in [28, 50] a SMC based sampler, like Multifidelity ST-MCMC, can estimate the model evidence of the highest fidelity model, $M_{K}$, by estimating the product:

Surrogate models can be derived by adding restrictive assumptions to the physics of the original model. In particular, consider a hypothetical physical biological surrogate of the original cells modeled by the full CME in which all cellular processes ‘freeze’ when molecular copy numbers reach a certain set of thresholds. As we increase these thresholds, the surrogate cells behave more freely and closer to the original cells and the master equation describing their behavior becomes closer to the original CME, illustrated in Fig. 1 .

Let $b_{1},\ldots,b_{N}$ be bounds on the copy number of species $1$ through $N$. We define an approximate SRN whose propensities are surrogates of the original SRN propensities and are given by

where $[E]$ takes value $1$ if expression $E$ is true and zero otherwise. Since there are no further transitions once the process enters a state that exceeds the bounds, the state space of the surrogate chemical master equation is effectively reduced to the hyper-rectangle $H(b)=\times_{i=1}^{N}\{0,\ldots,b_{i}\}$. Thus, the infinite-dimensional system of differential equations (3) is replaced by the finite-dimensional surrogate dynamical system

where the truncated infinitesimal generator $\widehat{A}_{H}$ is defined similar to eq. (4) but with the exact propensities replaced by the surrogate propensities given in eq. (33).

We note that eq. (34) is equivalent to eq. (6) with $\Omega=H(b)$. Thus, our surrogate propensities implicitly define a finite state projection of the original CME. We also note that the surrogate CMEs need not be constrained within a hyper-rectangle as considered here. It may be beneficial to derive a sequence of transformations to the original propensities, with the approximations chosen in such a way that alleviate the computational burden of solving the original model by, e.g., making the lower-fidelity dynamical system less stiff than the high-fidelity one. We leave this more general strategy to future work.

For the present choice hyper-rectangular state spaces, we recall the important result mentioned in section 2.2 that, as the entries of $\bm{b}$ increase monotonically, the state space $H(\bm{b})$ includes more states and the truncation error, measured as the $\ell_{1}$-distance between $\widehat{\bm{p}}_{H}$ and the true CME solution $\bm{p}(t)$ decreases monotonically. This provides us with a straightforward and natural way to form a hierarchy of surrogate models within the Multifidelity ST-MCMC framework.

With the surrogate CME models formulated, a strict hierarchy of surrogate models can be defined by a sequence of bounding vectors $\bm{b}^{(1)}\leq\bm{b}^{(2)}\leq\ldots\leq\bm{b}^{(K)}$ where the “$\leq$" sign applies element-wise. The corresponding surrogate models $M_{l}:=M(\bm{b}^{(l)})$ are then defined as in eq. (34). As mentioned earlier, the error in the surrogate CMEs decrease monotonically as we increase the bounds. Therefore, $\{M_{l}\}$ forms a hierarchy in which each level attains more fidelity than its predecessor.

At the $l$-th level, the log-likelihood function in eq. (8) is approximated by

In the surrogate log-likelihood, the data is projected onto the finite state space $H(\bm{b}_{l})$, and the probabilities of the data at different time points are computed from the surrogate Markov model $M_{l}$. Clearly, as $l$ increases, the surrogate function $L_{M_{l}}(\theta)$ becomes a more accurate approximation to the true log-likelihood $L(\mathcal{D}|\theta)$. In the ideal situation where the hierarchy is allowed to have infinite depth, these surrogates are guaranteed to form a sequence that converges pointwise to the true log-likelihood from below. This is shown formally in the following proposition.