Language models as master equation solvers

TL;DR

This study designs a prompt-based neural network to map rate parameters, initial conditions, and time values directly to the state joint probability distribution that exactly matches the input contexts, establishing the connection between language models and master equations.

Abstract

Master equations are of fundamental importance in modeling stochastic dynamical systems.However, solving master equations is challenging due to the exponential increase in the number of possible states or trajectories with the dimension of the state space. In this study, we propose repurposing language models as a machine learning approach to solve master equations. We design a prompt-based neural network to map rate parameters, initial conditions, and time values directly to the state joint probability distribution that exactly matches the input contexts. In this way, we approximate the solution of the master equation in its most general form. We train the network using the policy gradient algorithm within the reinforcement learning framework, with feedback rewards provided by a set of variational autoregressive models. By applying this approach to representative examples, we observe high accuracy for both multi-module and high-dimensional systems. The trained network also exhibits extrapolating ability, extending its predictability to unseen data. Our findings establish the connection between language models and master equations, highlighting the possibility of using a single pretrained large model to solve any master equation.

Figures (5)

• Figure 1: Caption on next page.
• Figure 2: Solving the time-dependent state joint distribution of the genetic toggle switch model.a, Schematic of the genetic toggle switch model. b, Average species counts as a function of time for the genes and proteins as compared with RNN and Gillespie results. c, Comparison of mean and standard deviation for species counts of genes and proteins. Results from RNN and MET are labeled with red and blue colors, respectively. d, The marginal distributions obtained from MET, RNN and Gillespie at time points $t=1$, $3$ and $40$. The inset is the Hellinger distance between the distribution obtained from MET and Gillespie. e, The joint distributions of the two proteins from MET, RNN and Gillespie at time points $t=1$, $3$ and $40$. The color bar shows the joint probability values in the logarithmic scale. $10^4$ states were sampled for MET and RNN at each time points, and Gillespie simulation has also $10^4$ trajectories.
• Figure 3: Caption on next page.
• Figure 4: Exploring the parameter space of the autoregulatory feedback loop model.a, Schematic of the autoregulatory feedback loop model. b, Marginal distribution of protein counts at time $t=10$ for 3 different parameter combinations (see Supplementary Notes for more details). c, Comparison of the mean and standard deviations of 14 parameters sampled from MET and Gillespie. d, Prediction of the bimodality coefficients from samples of trained MET as compared to the groundtruth (right figure). For MET, $10^4$ pairs of $\sigma_b, \rho_b$ was sampled with each pair $10^3$ sampled states. For Gillespie, $10^3$ trajectories was generated for 100 parameter points. e, Inferencing model rate parameters with the trained MET network. Dashed vertical lines indicate the correct values of the tested case. The histograms were the results of $10^4$ Monte Carlo steps.
• Figure 5: Trajectory ensemble sampling of the birth-death model.a, Schematic of the birth-death model. b, Hellinger distance between the marginal distributions of trajectories sampled from MET and trajectories sampled by Gillespie for each time point. c, Time-dependent trajectories sampled from MET (blue) and Gillespie (red). The first 300 trajectories are used for plotting. Average counts are shown as solid lines with respective colors. d, Mean and standard deviations of time-specific species counts for trajectories sampled from MET as compared with Gillespie's results. e, The marginal count distributions at time points $t = 5$, $25$, $49$ and $98$. The inset contains the corresponding Hellinger distance between the distributions sampled from MET and Gillespie.