Neural-network solutions to stochastic reaction networks

Abstract

The stochastic reaction network in which chemical species evolve through a set of reactions is widely used to model stochastic processes in physics, chemistry and biology. To characterize the evolving joint probability distribution in the state space of species counts requires solving a system of ordinary differential equations, the chemical master equation, where the size of the counting state space increases exponentially with the type of species. This makes it challenging to investigate the stochastic reaction network. Here we propose a machine learning approach using a variational autoregressive network to solve the chemical master equation. Training the autoregressive network employs the policy gradient algorithm in the reinforcement learning framework, which does not require any data simulated previously by another method. In contrast with simulating single trajectories, this approach tracks the time evolution of the joint probability distribution, and supports direct sampling of configurations and computing their normalized joint probabilities. We apply the approach to representative examples in physics and biology, and demonstrate that it accurately generates the probability distribution over time. The variational autoregressive network exhibits plasticity in representing the multimodal distribution, cooperates with the conservation law, enables time-dependent reaction rates and is efficient for high-dimensional reaction networks, allowing a flexible upper count limit. The results suggest a general approach to study stochastic reaction networks based on modern machine learning. Stochastic reaction networks involve solving a system of ordinary differential equations, which becomes challenging as the number of reactive species grows, but a new approach based on evolving a variational autoregressive neural network provides an efficient way to track time evolution of the joint probability distribution for general reaction networks.