DeepCME: A deep learning framework for computing solution statistics of the chemical master equation
Authors: Ankit Gupta, Christoph Schwab, Mustafa Khammash
Journal: PLOS Computational Biology
Accepted: 2021-11-08T00:00:00Z


Stochastic models of biomolecular reaction networks are commonly employed in systems and synthetic biology to study the effects of stochastic fluctuations emanating from reactions involving species with low copy-numbers. For such models, the Kolmogorov’s forward equation is called the chemical master equation (CME), and it is a fundamental system of linear ordinary differential equations (ODEs) that describes the evolution of the probability distribution of the random state-vector representing the copy-numbers of all the reacting species. The size of this system is given by the number of states that are accessible by the chemical system, and for most examples of interest this number is either very large or infinite. Moreover, approximations that reduce the size of the system by retaining only a finite number of important chemical states (e.g. those with non-negligible probability) result in high-dimensional ODE systems, even when the number of reacting species is small. Consequently, accurate numerical solution of the CME is very challenging, despite the linear nature of the underlying ODEs. One often resorts to estimating the solutions via computationally intensive stochastic simulations. The goal of the present paper is to develop a novel deep-learning approach for computing solution statistics of high-dimensional CMEs by reformulating the stochastic dynamics using Kolmogorov’s backward equation. The proposed method leverages superior approximation properties of Deep Neural Networks (DNNs) to reliably estimate expectations under the CME solution for several user-defined functions of the state-vector. This method is algorithmically based on reinforcement learning and it only requires a moderate number of stochastic simulations (in comparison to typical simulation-based approaches) to train the “policy function”. This allows not just the numerical approximation of various expectations for the CME solution but also of its sensitivities with respect to all the reaction network parameters (e.g. rate constants). We provide four examples to illustrate our methodology and provide several directions for future research.
Author summary: We develop a deep learning framework for estimating solutions of the chemical master equation (CME) that is fundamental to stochastic analysis of reaction networks. The CME is a system of ordinary differential equations that describes the time-evolution of the probability density of the random state-vector, and owing to an inherent curse of dimensionality, directly solving the CME is generally impractical with existing approaches. Moreover, the commonly employed simulation-based approaches for estimating CME solutions often require an exorbitant amount of computational time, even for moderately-sized networks. To counter these issues, we develop a deep reinforcement learning based method, called DeepCME, in this paper. DeepCME not only estimates function expectations based on the CME solution, but it also solves the more challenging problem of estimating their sensitivities with respect to all the model parameters. We illustrate our approach with four carefully chosen reaction network examples with varying sizes. Our results demonstrate that DeepCME reliably estimates the expectations of interest, along with all the parametric sensitivities, at a fraction of the computational cost of simulation-based estimators. We present many directions for future research and suggest further improvements to DeepCME that can greatly enhance its accuracy and applicability.