Language models as master equation solvers

TL;DR

This paper introduces a novel approach using language models to solve master equations, enabling accurate and extrapolatable solutions for complex stochastic systems through a prompt-based neural network trained with reinforcement learning.

Contribution

It repurposes language models as a general solver for master equations, demonstrating high accuracy and extrapolation in complex, high-dimensional systems.

Findings

High accuracy in multi-module systems

Effective extrapolation to unseen data

Establishes 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.