Learning moment closure in reaction-diffusion systems with spatial dynamic Boltzmann distributions

Abstract

Many physical systems are described by probability distributions that evolve in both time and space. Modeling these systems is often challenging due to their large state space and analytically intractable or computationally expensive dynamics. To address these problems, we study a machine-learning approach to model reduction based on the Boltzmann machine. Given the form of the reduced model Boltzmann distribution, we introduce an autonomous differential equation system for the interactions appearing in the energy function. The reduced model can treat systems in continuous space (described by continuous random variables), for which we formulate a variational learning problem using the adjoint method to determine the right-hand sides of the differential equations. This approach can be used to enforce a reduced physical model by a suitable parametrization of the differential equations. The parametrization we employ uses the basis functions from finite-element methods, which can be used to model any physical system. One application domain for such physics-informed learning algorithms is to modeling reaction-diffusion systems. We study a lattice version of the Rössler chaotic oscillator, which illustrates the accuracy of the moment closure approximation made by the method and its dimensionality reduction power.

I. INTRODUCTION

Probability distributions that evolve in both space and time appear in many modeling applications, such as reaction-diffusion systems [1-4], neural population activities [5,6], and fluid dynamics [7], as well as in engineering fields such as traffic forecasting [8] and navigation of autonomous vehicles [9]. However, (1) the state space of such distributions is generally large, and (2) the dynamical systems obeyed by their observables may be unknown or intractable to solve analytically. These aspects make modeling spatiotemporal systems a computational challenge and limit the interpretability of such models.

Reaction-diffusion systems are a typical example of these problems. The distribution over system states obeys a chemical master equation (CME) [10], but the state space grows exponentially with the number of random variables that describe it [11]. Further, the time evolution of observables is not closed, i.e., the time evolution of lower-order moments depends on higher-order ones (similar to a BBGKY hierarchy [12]). Their estimation therefore requires the use of a moment closure approximation (e.g., Refs. [13,14] and others; see Ref. [15] for a review), or otherwise sampling algorithms such as the Gillespie stochastic simulation algorithm (SSA) [16] or related methods for spatial systems [17,18].

A reduced model is one which approximates both the true distribution and its dynamics and should address the challenges above by (1) having a smaller state space and (2) being more easily tractable or computationally efficient [15]. Reduced models of reaction-diffusion systems are widely studied [1,19], particularly in multiscale modeling in biology [20]. Recent work [2,4,13] has demonstrated methods based on entropic matching as a highly general approach to model reduction of reaction networks.

In this paper, we demonstrate a machine-learning (ML) approach to model reduction using Boltzmann machines (BMs) [21]. We formalize the methods of earlier work [15,22] and extend these with the introduction of latent variables. Our approach also extends work on entropic matching methods to treat spatial systems. We present examples for spatial chemical reaction systems that demonstrate the moment closure properties of the reduced model and apply the method to learn a spatial chaotic oscillator.

The area of ML most suited for model reduction of reaction-diffusion systems are generative models [23], where it is assumed that data are samples of an unknown probability distribution, with the goal of estimating this distribution by a structured approach. This structure can offer insight into the problem that has not been obtainable analytically [24] and allows new samples to be drawn using, e.g., Markov-chain Monte Carlo methods [25]. Typically, a graphical model for the distribution is introduced and learned by determining interaction parameters between random variables. Similar ML approaches have emerged as a powerful tool for studying quantum many-body problems [26,27].

Our approach introduces a differential equation (DE) model for interaction parameters in the graph. The learning problem is formulated to determine these DEs by a maximum likelihood approach. In contrast to ML methods for learning temporal data such as recurrent networks, here prior information about the system may be used to enforce a reduced physical model by parametrizing the functional forms of the DEs.

A further advantage of this strategy is that it offers a natural description of systems where neither time nor space are discretized, i.e., the system is described by random variables representing space continuously and varying continuously in time. In this case, a partial differential equation (PDE) model can be introduced. Spatially continuous descriptions are beneficial when confined geometries would introduce error into lattice-based methods, e.g., when modeling reaction-diffusion systems at synapses [17].

The algorithmic solution to this learning problem takes the form of a PDE-constrained optimization problem. The algorithm and its derivation are closely related to BM learning, but in this case data samples are trajectories in space and time rather than instantaneous snapshots or slices. A related framework, graph-constrained correlation dynamics [15], has a similar learning goal but uses spatially aggregated snapshots in time and does not consider spatial reduced models.

The outline of this paper is as follows: (1) in Sec. II we introduce spatial dynamic Boltzmann distributions as reduced models of reaction-diffusion systems in continuous space and formulate their learning problem using the adjoint method; (2) in Sec. III we demonstrate the connection to a restricted Boltzmann machine; (3) in Sec. IV we show how hidden layers implement moment closure approximations and apply the method to a spatial chaotic oscillator.

II. SPATIAL DYNAMIC BOLTZMANN DISTRIBUTIONS

In this section, we introduce the reduced model for a spatiotemporal distribution and its dynamics in continuous space from Ref. [22] and formulate the learning problem using the adjoint method. We consider the specific application of a reaction-diffusion system but note that the methods are also applicable to other spatiotemporal systems.

The state of a reaction-diffusion system at some time t is described by n particles of species labels α located at positions x in generally continuous three-dimensional (3D) space (each x_i for i = 1, . . . , n is a coordinate in 3D space). Let the true distribution over system states be denoted by p(n, α, x, t), whose time evolution can be described using the Doi-Peliti formalism [46].

To define the reduced model, introduce k-particle interaction functions ${\nu }_k({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t)$, where $\langle i{\rangle }_k^n$ denotes any ordered subset of k indexes with each index in {1, . . . , n}. Given a set of such interaction functions $\{\nu {\}}_{k=1}^K$ up to cutoff order K, define a spatial dynamic Boltzmann distribution as one of the form:

$$\tilde{p}(n,\alpha ,\mathit{x},t;\{\nu \})=\frac{1}{Z[\{v\}]}\exp \left[-\sum _{k=1}^K\sum _{\langle i{\rangle }_k^n}{v}_k({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t)\right],$$ (1)

where the sum over $\langle i{\rangle }_k^n$ iterates over unique kth-order interactions between n particles, and the partition function is

$$Z[\{\nu \}]=\sum _{n=0}^{\infty }\sum _{\alpha }\int d\mathit{x}\exp \left[-\sum _{k=1}^K\sum _{\langle i{\rangle }_k^n}{\nu }_k({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t)\right].$$ (2)

Boltzmann distributions are maximum entropy (MaxEnt) distributions, where each interaction function ${\nu }_k({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t)$ controls a corresponding moment ${\mu }_k({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t)$, given by:

$$\begin{gathered} {\mu }_k\left({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t\right)=\sum _{n^{\prime }=0}^{\infty }\sum _{{\alpha }^{\prime }}\int d{\mathit{x}}^{\prime }p(n^{\prime },{\alpha }^{\prime },{\mathit{x}}^{\prime },t) \\ \times \sum _{\langle j{\rangle }_k^{n^{\prime }}}\delta \left({\mathit{x}}_{\langle i{\rangle }_k^n}-{\mathit{x}}_{\langle j{\rangle }_k^{n^{\prime }}}^{\prime }\right)\delta \left({\alpha }_{\langle i{\rangle }_k^n}-{\alpha }_{\langle j{\rangle }_k^{n^{\prime }}}^{\prime }\right), \end{gathered}$$ (3)

that is, the average number of k-sized tuplets of particles of species ${\alpha }_{\langle i{\rangle }_k^n}$ at locations ${\mathit{x}}_{\langle i{\rangle }_k^n}$. Note that α′ and x′ are of size n′.

A. Moment matching

Given a set of training data drawn from p(n, α, x, t) at some instant in time, the BM learning algorithm determines parameters in the energy function such that the instantaneous distribution (1) is the MaxEnt distribution consistent with the moments in the data set. To learn a reduced model of a system that evolves in both time and space continuously, we seek the distribution that is at all times the MaxEnt solution. Define as the action the Kullback-Leibler (KL) divergence ${\mathcal{D}}_{\mathrm{KL}}$ between the true and reduced models, p and $\tilde{p}$, integrated over all times:

$$S={\int }_{t_0}^{t_f}dt{\mathcal{D}}_{\mathrm{KL}}(p\Vert \tilde{p}),$$ (4)

where the Lagrangian is $\mathcal{L}(t;\{\nu \})={\mathcal{D}}_{\mathrm{KL}}(p\Vert \tilde{p})$ for

$$\begin{gathered} {\mathcal{D}}_{\mathrm{KL}}(p\Vert \tilde{p})=\sum _{n=0}^{\infty }\sum _{\alpha }\int d\mathit{x} \\ \times p(n,\alpha ,\mathit{x},t)\ln \frac{p(n,\alpha ,\mathit{x},t)}{\tilde{p}(n,\alpha ,\mathit{x},t;\{\nu \})}. \end{gathered}$$ (5)

Minimizing S is thus equivalent to maximizing the integrated log-likelihood of the observed data given the interaction functions. Other approaches for modeling time series are discussed in Sec. III A.

The condition for extremizing the action follows from the chain rule as

$$\begin{gathered} \delta S={\int }_{t_0}^{t_f}dt\sum _{n=0}^{\infty }\sum _{\alpha }\int d\mathit{x} \\ \times \sum _{k=1}^K\sum _{\langle i{\rangle }_k^n}\Delta {\mu }_k\left({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t\right)\delta {\nu }_k\left({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t\right)=0, \end{gathered}$$ (6)

where

$$\Delta {\mu }_k\left({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t\right)={\tilde{\mu }}_k\left({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t\right)-{\mu }_k\left({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t\right),$$ (7)

where μ and $\tilde{\mu }$ are averages taken over p and $\tilde{p}$. This appearance of a difference of moments is the common result from using the KL divergence in the objective functional.

B. An adjoint method learning problem for spatial dynamic Boltzmann distributions

Introduce for each interaction function ${\nu }_k({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t)$ a functional model:

$$\frac{d}{dt}{\nu }_k\left({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t\right)={\mathcal{F}}_k[\{\nu \}](\alpha ,\mathit{x},t),$$ (8)

with initial condition ${\nu }_k({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t_0)={\eta }_k({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n})$ and where $\{\nu \}=\{{\nu }_k{\}}_{k=1}^K$ denotes possibly all interaction functions. We use $\mathcal{F}$ to denote a functional, allowing, for example, a PDE model to be introduced. Note that the arguments to the left-hand side may also appear on the right, for example, through a spatial derivative term$\nabla {\nu }_k({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t)$.

Introduce vector notation¹ ν(α, x, t) and $\mathcal{F}[\{\nu \}](\alpha ,\mathit{x},t)$ for the left- and right-hand sides of (8), which contain $N={\sum }_{k=1}^K\left(\begin{array}{c}\hfill n\hfill \\ \hfill k\hfill \end{array}\right)$ entries, one for every possible $(k,\langle i{\rangle }_k^n)$ in some order i = 1, . . . , N. To enforce the constraint (8), define the Lagrangian as the functional:

$$\begin{gathered} \mathcal{L}[\{\nu \},\{\zeta \}](t)={\mathcal{D}}_{\mathrm{KL}}(p\Vert \tilde{p})+\sum _{n=0}^{\infty }\sum _{\alpha }\int d\mathit{x}{\zeta }^{⊺}(\alpha ,\mathit{x},t) \\ \times \left[\frac{d\nu (\alpha ,\mathit{x},t)}{dt}-\mathcal{F}[\{\nu \}](\alpha ,\mathit{x},t)\right], \end{gathered}$$ (9)

where we have introduced Lagrange multiplier functions ζ(α, x, t) corresponding to ν(α, x, t) and $\{\zeta \}=\{{\zeta }_k{\}}_{k=1}^K$. Since the constraint is satisfied, then the action is as before $S={\int }_{t_0}^{t_f}dt\mathcal{L}[\{\nu \},\{\zeta \}](t)$.

Introducing perturbations δν(α, x, t) to the interaction functions gives as condition for extremizing the action:

$$\begin{gathered} \delta S={\int }_{t_0}^{t_f}dt\sum _{n=0}^{\infty }\sum _{\alpha }\int d\mathit{x}\delta {\mathit{v}}^{⊺}(\alpha ,\mathit{x},t) \\ \times \left\{\Delta \mu (\alpha ,\mathit{x},t)-\frac{d\zeta (\alpha ,\mathit{x},t)}{dt}-\frac{\delta \mathcal{J}[\{\nu \},\{\zeta \}](t)}{\delta \nu (\alpha ,\mathit{x},t)}\right\}=0, \end{gathered}$$ (10)

where the boundary terms from the integration by parts in the second term have vanished due to the boundary condition for the adjoint variables ζ(α, x, t_f) = 0, and we have defined:

$$\mathcal{J}[\{\nu \},\{\zeta \}](t)=\sum _{n^{\prime }=0}^{\infty }\sum _{{\alpha }^{\prime }}\int d{\mathit{x}}^{\prime }\times {\zeta }^{⊺}({\alpha }^{\prime },{\mathit{x}}^{\prime },t)\mathcal{F}[\{\nu \}](\alpha \prime ,{\mathit{x}}^{\prime },t).$$ (11)

From (10) we obtain the adjoint system

$$\frac{d\zeta (\alpha ,\mathit{x},t)}{dt}=\Delta \mu (\alpha ,\mathit{x},t)-\frac{\delta \mathcal{J}[\{\nu \},\{\zeta \}](t)}{\delta \nu (\alpha ,\mathit{x},t)}.$$ (12)

Depending on the form of the functional, additional boundary conditions may be enforced to evaluate the term on the right. Equations (8) and (12) can be equivalently expressed by the Hamiltonian system

$$\begin{gathered} \frac{d\nu (\alpha ,\mathit{x},t)}{dt}=\frac{\delta H[\{\nu \},\{\zeta \}](t)}{\delta \zeta (\alpha ,\mathit{x},t)}, \\ \frac{d\zeta (\alpha ,\mathit{x},t)}{dt}=-\frac{\delta H[\{\nu \},\{\zeta \}](t)}{\delta \nu (\alpha ,\mathit{x},t)}, \end{gathered}$$ (13)

where

$$H[\{\nu \},\{\zeta \}](t)=-{\mathcal{D}}_{\mathrm{KL}}(p\Vert \tilde{p})+\mathcal{J}[\{\nu \},\{\zeta \}](t).$$ (14)

Given a reduced model for the dynamics (8), Eq. (10) gives the necessary condition for extremizing the action. In a typical model reduction setting, however, the reduced model is not known beforehand. What should the form of the model (8) be to extremize the action (4)? Consider the case where the functional is specified in terms of some ordinary functions. We next set up a variational problem for these functions appearing on the right-hand side of the differential equation. Variational problems of this form have been studied previously, first in the context of optimal control theory [28,29] and later didactically in Ref. [30].

Let the functional be of the form:

$$\frac{d}{dt}{\nu }_k({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t)={\mathcal{F}}_k[\{\nu \},\{F_k\}](\alpha ,\mathit{x},t),$$ (15)

where the M_k ordinary functions appearing on the right-hand side are $F_k^{(s)}(\{\nu (\alpha ,\mathit{x},t)\})$ for s = 1, . . . , M_k, denoted by $\{F_k\}=\{F_k^{(s)}{\}}_{s=1}^{M_k}$. For arbitrary perturbations $\delta F_k^{(s)}$, extremizing the action gives

$$\begin{gathered} \delta S=-{\int }_{t_0}^{t_f}dt\sum _{n=0}^{\infty }\sum _{\alpha }\int d\mathit{x}\sum _{k=1}^K\sum _{\langle i{\rangle }_k^n}\sum _{s=1}^{M_k}\frac{\delta \mathcal{J}[\{\nu \},\{\zeta \}](t)}{\delta F_k^{(s)}(\{\nu (\alpha ,\mathit{x},t)\})} \\ \times \delta F_k^{(s)}(\{\nu (\alpha ,\mathit{x},t)\})=0. \end{gathered}$$ (16)

Equation (16) is the variational calculus form of the sensitivity equation obtained by the adjoint method when the functional model is specified in terms of some parameter vector [31]. This is particularly clear if we consider the specific form of (15) as the autonomous ordinary differential equation (ODE) system:

$$\frac{d}{dt}{\nu }_k\left({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t\right)=F_k\left(\left\{\nu ({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t)\right\}\right),$$ (17)

where $\{\nu ({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t)\}$ denotes all ν of all possible arguments appearing on the left-hand side. In this case, (16) becomes

$$\begin{gathered} \delta S=-{\int }_{t_0}^{t_f}dt\sum _{n=0}^{\infty }\sum _{\alpha }\int d\mathit{x} \\ \times [{\zeta }^{⊺}(\alpha ,\mathit{x},t)\delta \mathit{F}(\{\nu (\alpha ,\mathit{x},t)\})]=0, \end{gathered}$$ (18)

where as before we have used vectors of length N to denote possible $(k,\langle i{\rangle }_k^n)$ as before. This resembles the adjoint method sensitivity equation, where variational terms δF_k and δS replace ordinary derivatives with respect to parameters. This will be pursued further in Sec. III A. From (18) follows the common result that extremizing the action requires that the adjoint variables vanish everywhere ${\zeta }_k({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t)=0$. One case when this is satisfied is if the adjoint system is source free $\Delta {\mu }_k({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t)=0$, i.e., the moment matching condition is met.

From the Euler-Lagrange equations (12), the adjoint variables obey:

$$\frac{d\zeta (\alpha ,\mathit{x},t)}{dt}=\Delta \mu (\alpha ,\mathit{x},t)-G^{⊺}(\alpha ,\mathit{x},t)\zeta (\alpha ,\mathit{x},t),$$ (19)

where the elements of the N × N matrix G are

$$G_{i,i^{\prime }}(\alpha ,\mathit{x},t)=\frac{\partial {\mathit{F}}_k\left[\left\{\nu \left({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t\right)\right\}\right]}{\partial {\nu }_{k^{\prime }}\left({\alpha }_{\langle i{\rangle }_k^n\prime },{\mathit{x}}_{\langle i{\rangle }_k^n\prime },t\right)},$$ (20)

where $(k,\langle i{\rangle }_k^n)$ corresponds to index i and $(k^{\prime },\langle i{\rangle }_{k^{\prime }}^n)$ corresponds to index i′. Appendix A gives the formal solution to (19) and makes explicit the connection between the conditions for extrema (18) and (6).

III. DYNAMICS FOR RESTRICTED BOLTZMANN MACHINES

We next consider a specific case of the formalism of Sec. II where the system is described by discrete random variables. A Boltzmann distribution on a state v = {v₁, . . . , v_N} of N discrete random variables is of the form:

$$\tilde{p}(\mathit{v})=\frac{1}{Z}\text{exp}[-E(\mathit{v})],$$ (21)

where Z is the partition function, and the energy function E(v) is typically defined by a chosen Markov random field (MRF). For example, a BM [21] is a binary MRF, where binary units update their state based on a bias and pairwise connections to other units. A MRF where all variables v are driven by data is fully visible; otherwise, the N′ units h = {h₁, . . . , h_N′} which are not driven by data are denoted as hidden.

A restricted Boltzmann machine (RBM) [32] is a BM in which hidden and visible units are organized into layers, where a layer is defined by the property that there are no interactions among units in the same layer. For example, a typical energy function for an RBM is of the form:

$$E(\mathit{v},\mathit{h},\theta )=-\sum _{i=1}^N{b}_i{v}_i-\sum _{j=1}^{N^{\prime }}{b^{\prime }}_j{h}_j-\sum _{\{i,j\}}{W}_{i,j}{v}_i{h}_j,$$ (22)

where the summation {i, j} is determined by the graph edges and θ is the vector of length K of all interaction parameters in the graph. This defines a joint distribution over v and h:

$$\tilde{p}(\mathit{v},\mathit{h};\theta )=\frac{1}{Z(\theta )}\text{exp}[-E(\mathit{v},\mathit{h},\theta )].$$ (23)

Each parameter θ_k in this MaxEnt distribution controls a corresponding moment ${\tilde{\mu }}_k$, given by ${\tilde{\mu }}_k=\partial \ln Z(\theta )∕\partial {\theta }_k$.

Define a dynamic Boltzmann distribution as one with time-dependent interaction parameters:

$$\tilde{p}(\mathit{v},\mathit{h};\theta (t))=\frac{1}{Z[\theta (t)]}\text{exp}\{-E(\mathit{v},\mathit{h},\theta (t)]\}.$$ (24)

For example, the energy function of the RBM becomes

$$\begin{gathered} E[\mathit{v},\mathit{h},\theta (t)]=-\sum _{i=1}^N{b}_i(t)v_i-\sum _{j=1}^{N^{\prime }}{b^{\prime }}_j(t)h_j \\ -\sum _{\{i,j\}}{W}_{i,j}(t)v_i{h}_j. \end{gathered}$$ (25)

This is a specific case of a spatial dynamic Boltzmann distribution (1) in the discrete lattice limit. To see this, assign to every visible unit v_i a spatial location x_i. By taking self-interaction functions ν₁(x, t) = − ∑_i b_i(t)δ_x,xi in (1), we recover the first term in (25) with v_i ∈ {0, 1}, where δ_x,xi is unity if the coordinates are coincident and zero otherwise.

Similarly, hidden units can also be represented in continuous space. Let the species labels α_v denote visible units and β_h denote hidden units, and assign to every hidden unit h_j a spatial location y_j. The weights between layers are then obtained by taking pairwise interactions ν₂(α, β, x, y, t) = − ∑_{{i, j}} W_{i, j}(t)δ_x,xiδ_y,yjδ_α,αvδ_β,βh.

A. An adjoint method learning problem for restricted Boltzmann machines

Introduce for each interaction parameter θ_k, k = 1, . . . , K, in the interaction graph a time-evolution function F_k forming an autonomous ODE system [analogously to (17)]:

$$\frac{d}{dt}{\theta }_k(t)=F_k(\theta (t)),$$ (26)

with initial conditions θ_k(t₀) = θ_k,0. To obtain from the variational problem derived in Sec. IIB an ordinary optimization problem for parameters, further consider the paramaterization by the vectors u_k of size M_k, generally unique for every k:

$$\frac{d}{dt}{\theta }_k(t)=F_k(\theta (t);{\mathit{u}}_k).$$ (27)

Analogously to the continuous case, define as the objective function the KL divergence between the true and reduced models, p and $\tilde{p}$, over all times [analogously to (4)]:

$$\begin{gathered} S={\int }_{t_0}^{t_f}dt{\mathcal{D}}_{\mathrm{KL}}(p\Vert \tilde{p}), \\ {\mathcal{D}}_{\mathrm{KL}}(p\Vert \tilde{p})=\sum _{z}p(z)\ln \frac{p(z)}{\tilde{p}(z;\{\mathit{u}\})}. \end{gathered}$$ (28)

where $\{\mathit{u}\}=\{{\mathit{u}}_k{\}}_{k=1}^K$. Minimizing S is thus equivalent to maximizing the log-likelihood of the observed data given the parameters, i.e., $L(\{\mathit{u}\};z)=\log \tilde{p}(z;\{\mathit{u}\})$. A more common approach is to instead maximize the conditional likelihood of observations conditioned on the first observation: $L(\{\mathit{u}\};z_2,z_3,\dots \mid z_1)=\log \tilde{p}(z_2,z_3,\dots \mid z_1;\{\mathit{u}\})$ or similar causal relations. For Markov chains, this approach is highly successful (leading to, e.g., Kalman filters; see Ref. [33] for an introduction). If a prior is available, then Bayesian methods that compute the posterior $\tilde{p}(\{\mathit{u}\};z)\propto \tilde{p}(z;\{\mathit{u}\})\times \tilde{p}(\{\mathit{u}\})$ can provide further improvements. The advantage of the current approach is that a reduced physical model can be enforced through the parametrization (27). This model can be based on prior information, such as reaction networks with known solutions [22]. A second advantage is that the generalization to spatially continuous systems follows naturally using PDEs as in (8).

	Algorithm 1 Stochastic Gradient Descent for Learning Restricted Boltzmann Machine Dynamics.
1:	Initialize
2:	Parameters uk controlling the functions Fk(θ; uk) for all k = 1, . . . , K.
3:	Time interval [t0, tf], a formula for the learning rate λ.
4:	while not converged do
5:	Initialize ΔFk,i = 0 for all k = 1, . . . , K and parameters i = 1, . . . , Mk.
6:	for sample in batch do
7:	⊳ Generate trajectory in reduced space θ:
8:	Solve the PDE constraint (27) for θk(t) with a given IC θk,0 over t0 ⩽ t ⩽ tf, for all k.
9:	⊳ Wake phase:
10:	Evaluate moments μk(t) of the data for all k, t.
11:	⊳ Sleep phase:
12:	Evaluate moments ${\tilde{\mu }}_k(t)$ of the Boltzmann distribution.
13:	⊳ Solve the adjoint system:
14:	Solve the adjoint system (31) for ϕk(t) for all k, t.
15:	⊳ Evaluate the objective function:
16:	Update ΔFk,i as the cumulative moving average of the sensitivity equation (30) over the batch.
17:	⊳ Update to decrease objective function:
18:	uk,i → uk,i − λΔFk,i for all k, i.

The time integral in S can lead to undesired extrema, for example for periodic systems where the objective function may not minimize the KL divergence at each time point. One algorithmic strategy for eliminating these in practice is to shift the limits of integration during the optimization, as done in the examples of Sec. IV A.

Minimizing the objective function defines a PDE-constrained optimization problem: minimize (28) subject to the PDE constraint (27). Define the Lagrangian function [analogously to (9)]:

$$\begin{gathered} \mathcal{L}(t;\{\mathit{u}\})={\mathcal{D}}_{\mathrm{KL}}(p\Vert \tilde{p})+\sum _{k=1}^K{\varphi }_k(t) \\ \times \left\{\frac{d}{dt}{\theta }_k(t)-F_k[\theta (t);{\mathit{u}}_k]\right\}, \end{gathered}$$ (29)

where we have introduced the adjoint variables ϕ_k associated with each θ_k. Taking the derivative of the objective function $S={\int }_{t_0}^{t_f}dt\mathcal{L}(t;\{\mathit{u}\})$ with respect to a parameter gives the sensitivity equation [analogously to (18)]:

$$\frac{dS}{d{u}_{k,i}}=-{\int }_{t_0}^{t_f}dt\frac{\partial F_k[\theta (t);{\mathit{u}}_k]}{\partial u_{k,i}}{\varphi }_k(t),$$ (30)

and taking the derivative with respect to θ gives the ODE system obeyed by the adjoint variables [analogously to (19)]:

$$\frac{d}{dt}{\varphi }_k(t)={\tilde{\mu }}_k(t)-{\mu }_k(t)-\sum _{l=1}^K\frac{\partial F_l[\theta (t);{\mathit{u}}_l]}{\partial {\theta }_k(t)}{\varphi }_l(t),$$ (31)

where μ_k(t′) and ${\tilde{\mu }}_k(t^{\prime })$ are averages taken over to p and $\tilde{p}$ at time t′, and the boundary condition is ϕ_k(t_f) = 0.

Algorithm I outlines how this optimization problem can be solved in practice. The inner loop of an “wake” and “sleep” phase of sampling are identical to that of BM learning. Standard algorithmic improvements are possible, such as the use of accelerated gradient descent methods such as Adam [34], and using persistent contrastive divergence (PCD) [35] to estimate the moments of the reduced model ${\tilde{\mu }}_k(t^{\prime })$.

Adjoint methods for solving PDE-constrained optimization problems are also called “black-box” methods [36,37], since the PDE constraint (27) is eliminated in the derivation of the sensitivity equation (30). A competing class of methods (sometimes referred to as “all-at-once” methods) treat the constraint explicitly in the optimization, and may offer a computational advantage over this approach. These include sequential quadratic programming and augmented Lagrangian methods.

Additional constraints or regularization terms can be included in the optimization, such as conserved quantities identified from the left null space of the net stoichiometry matrix. For example, L₂ regularization can be incorporated into the objective function:

$$S={\int }_{t_0}^{t_f}dt{\mathcal{D}}_{\mathrm{KL}}(p\Vert \tilde{p})+{\lambda }_r{\int }_{t_0}^{t_f}dt\sum _{k=1}^K[{\theta }_k(t)-{\overline{\theta }}_k(t){]}^2,$$ (32)

where ${\overline{\theta }}_k(t)$ are some specified functions or otherwise constant and λ_r is a regularization parameter. In this case, the adjoint variables are given by:

$$\begin{gathered} \frac{d}{dt}{\varphi }_k(t)={\tilde{\mu }}_k(t)-{\mu }_k(t)+2{\lambda }_r[{\theta }_k(t)-{\overline{\theta }}_k(t)] \\ -\sum _{l=1}^K\frac{\partial F_l(\theta (t);{\mathit{u}}_l)}{\partial {\theta }_k(t)}{\varphi }_l(t). \end{gathered}$$ (33)

B. Finite-element parameterization

What choice should be made for the parametrization (27) of the right-hand sides of the differential equations? In Ref. [22], we considered simple reaction-diffusion systems from which general forms of approximate models could be inferred that maintain physical interpretations. A second approach also explored in Ref. [22] is to use a separate moment closure approximation to derive analytic solutions for simple reaction systems on 1D lattices, where the inverse Ising problem is analytically solvable. The form of (27) can then be taken as either linear or nonlinear combinations of known solutions.

Here, we take a finite-element method [38] approach to the parametrization that is more aligned with the unsupervised learning problem in a Boltzmann machine. The space of solutions to the general variational problem (16), which is some Banach space, is therefore restricted to the space of finite-element method solutions.

An important restriction is that the learning rule (30) requires C₁ finite elements. One choice for such elements is the Q₃ family of finite elements [39], which has the advantage that basis functions in dimensions higher than one are easily constructed as tensor products of 1D cubic polynomials.² For C₁ elements that control the value of the function and its derivative at the endpoints, these polynomials are just the Hermite polynomials, shown in Fig. 1(d).

FIG. 1.

Comparison of a fully visible and a latent variable model for capturing local correlations in a 1D lattice. (a) One-dimensional lattice with one hidden layer (similar to an RBM). Note that in this simplified example, W is a single translation invariant parameter rather than a matrix as common in RBMs. (b) Fully visible model for a 1D lattice including NN interactions J and NNN interactions K. (c) An example state of the hidden layer model, where blue indicates the presence of a particle in the visible layer and likewise red for the hidden layer. By learning the parameters, the hidden layer can be tuned to capture the presence of NNs. (d) The basis functions of the Q₃ family of C₁ finite elements in 1D (Hermite polynomials), used to parametrize the right-hand sides of (38) and (40). Basis functions in higher dimensions are constructed as tensor products of the 1D polynomials. (e) Moments of stochastic simulations for 10 of the 50 initial conditions used for training (each trajectory obtained from averaging over 50 lattices simulated from the same initial condition).

We introduce for each time-evolution function in (27) a domain of hypercubic cells, with 4^d degrees of freedom, where d are the number of arguments to F_k. In practice, we found it is rarely necessary to have more than d = 3 arguments (see Sec. IV). For d = 3, each cube has 64 degrees of freedom (8 degrees of freedom at each vertex, specifying the function value and derivatives). For a cubic lattice of V = L₁ × L₂ × L₃ cells, there are 8V degrees of freedom in total, with the parametrization taking the usual form in terms of the basis functions f_l associated with each degree of freedom:

$$F_k({\theta }_1,{\theta }_2,{\theta }_3;{\mathit{u}}_k)=\sum _{l=1}^{8V}{u}_l{f}_l({\theta }_1,{\theta }_2,{\theta }_3).$$ (34)

Note that here the right-hand side of the differential equation is parameterized (as opposed to the solution of the differential equation), since the objective of the learning algorithm is to determine a suitable differential equation model.

IV. LEARNING REACTION-DIFFUSION SYSTEMS ON LATTICES

Recall that the state of a reaction-diffusion system at some time is described by n particles of species α located at positions x in generally continuous 3D space. To make an explicit connection to binary random variables, we consider a simpler model of particles hopping on a discrete lattice in the single-occupancy limit. To generate stochastic simulations of such a system, we adapt the method of Takayasu and Tretyakov [40] for a lattice-based variant of the popular Gillespie SSA [16] as follows: At each time step:

(1) Perform unimolecular reactions following the standard Gillespie SSA.

(2) Iterate over all particles in random order; for each:

(a) Hop to a neighboring site, chosen at random with equal probability.

(b) If the site is unoccupied, then the move is accepted. If the site is occupied, then a bimolecular reaction occurs with some probability; else, the move is rejected and the particle is returned to the original site.

The lattice on which particles hop is designated as the visible part of the MRF. Assign a unique index i to each of the N sites in the lattice, and let the vector of possible species be s of size M in some arbitrary ordering (excluding ∅ to denote an empty site). Spins at a site i are now multinomial units, represented as a vector v_i of length M where entries v_i,α ∈ {0, 1} for α = 1, . . . , M denote the absence or presence of a particle of species s_α (an n-vector model in statistical mechanics). The single-occupancy limit corresponds to the implicit constraint that the vectors are of unit length, i.e., ${\sum }_{\alpha =0}^M{v}_{i,\alpha }=1$, where α = 0 denotes an empty site. The matrix V of size N × M describes the state of the visible part of the MRF, where each row denotes a lattice site.

Likewise, introduce hidden layer species s′ of size M′, which may be different from s. Indexing all hidden sites as j = 1, . . . , N′, hidden unit vectors are h_j of length M′. The state of the hidden units is H of size N′ × M′, with the single-occupancy constraint as before.

The dynamic Boltzmann distribution becomes $\tilde{p}[\mathit{V},\mathit{H}\mid \theta (t)]=\text{exp}\{-E[\mathit{V},\mathit{H}\mid \theta (t)]\}∕Z[\theta (t)]$, where interaction parameters θ(t) may also be species dependent. For example, the energy function for the RBM becomes

$$\begin{gathered} E[\mathit{V},\mathit{H},\theta (t)]=-\sum _{i=1}^N\sum _{\alpha =1}^M{b}_{i,\alpha }(t)v_{i,\alpha } \\ -\sum _{j=1}^{N^{\prime }}\sum _{\beta =1}^{M^{\prime }}{b^{\prime }}_{j,\beta }(t)h_{j,\beta } \\ -\sum _{\{i,j\}}\sum _{\alpha ,\beta }{W}_{i,j,\alpha ,\beta }(t)v_{i,\alpha }{h}_{j,\beta }. \end{gathered}$$ (35)

A. Learning hidden layers for moment closure

A typical problem in many-body systems is the appearance of a hierarchy of moments, where the time evolution of a given moment depends on higher-order moments. Moment closure approximations terminate this infinite hierarchy at some finite order. In this section, we develop the perspective of the learning problem (30) as a closure approximation using a simple pedagogical example. We note some similarity to previously proposed closure schemes [14,15], as well as to entropic matching [13], although the current approach differs in the objective function (28) and the formulation for spatially continuous systems in Sec. II.

Consider a bimolecular-annihilation process on a 1D lattice of length N, where particles of a single species A hop and react according to A + A → ∅. The time evolutions of the first two moments are (see Appendix B)

$$\begin{gathered} \frac{d}{dt}\langle \sum _i{v}_i\rangle =-2k_r\langle \sum _i{v}_i{v}_{i+1}\rangle , \\ \frac{d}{dt}\langle \sum _i{v}_i{v}_{i+1}\rangle =2D\langle \sum _i{v}_i{v}_{i+2}\rangle -2k_r\langle \sum _i{v}_i{v}_{i+1}{v}_{i+2}\rangle \\ +(k_r-2D)\langle \sum _i{v}_i{v}_{i+1}\rangle , \end{gathered}$$ (36)

where k_r is the reaction rate and D the diffusion rate. The simplest graph to capture such observables is a fully visible Markov random field with N units, i.e., a 1D Ising model including interactions up to some order. For example, including third-order interactions, let:

$$\begin{gathered} E[\mathit{v},b(t),J(t),K(t)]=-b(t)\sum _{i=1}^N{v}_i-J(t)\sum _{i=1}^{N-1}{v}_i{v}_{i+1} \\ -K(t)\sum _{i=1}^{N-2}{v}_i{v}_{i+1}{v}_{i+2}, \end{gathered}$$ (37)

where b is the bias, J is the nearest neighbor (NN) interaction term, and K is the next-nearest-neighbor (NNN) interaction term. Let the differential equation model be

$$\begin{gathered} \stackrel{.}{b}=F_b(b,J,K;{\mathit{u}}_b), \\ \stackrel{.}{J}=F_J(b,J,K;{\mathit{u}}_J), \\ \stackrel{.}{K}=F_K(b,J,K;{\mathit{u}}_K), \end{gathered}$$ (38)

for some parameter vectors u to be learned, where time derivatives are denoted as ẋ = d/dt. The corresponding graphical model is illustrated in Fig. 1(b). The choice of the energy function in (37) defines which moments are explicitly captured by the reduced model. The additional choice of the form of the differential equations F_γ defines the moment closure approximation made.

We next show through computational experiments that the introduction of hidden layers can improve on a fully visible closure model:

(1) In any closure scheme, moments beyond a certain order are not captured explicitly by the model, so that their approximation may be poor. The representation power of hidden layers [24] can be used to incorporate information about which higher-order moments are relevant to the data set.

(2) Two distinct states having the same lower-order moments are indistinguishable in the reduced model (the model is not sufficiently high dimensional). Hidden layers may be able to separate such states if their connectivity is suitably chosen to represent relevant higher-order correlations, even if the model remains low order.

(3) The number of higher-order terms appearing on the right of (36) grows with the order on the left. This problem is compounded if species labels are included. Hidden layers and a restriction on the number of species M′ allowed to occupy hidden units may be used to approximate such higher-order interactions with fewer parameters.

It is generally difficult to choose the optimal close approximation, i.e., to know which moments are relevant to the time evolution of a given data set. A key advantage of the present approach is that the connectivity of the hidden layers may be chosen based on the differential equations derived from the chemical master equation. For example, consider to the bimolecular annihilation system (36): If the goal is to accurately model the mean number of particles, then the right-hand side of (36) shows that the nearest-neighbor moment is relevant to the time evolution. The graphical model of the reduced system could therefore introduce a hidden unit for every pair of neighboring lattice sites (N − 1 units in the hidden layer), with corresponding energy function:

$$\begin{gathered} E[\mathit{v},\mathit{h},b(t),W(t),b^{\prime }(t)]=-b(t)\sum _{i=1}^N{v}_i-b^{\prime }(t)\sum _{j=1}^{N-1}{h}_j \\ -W(t)\sum _{j=1}^{N-1}\sum _{i\in \{j,j+1\}}{v}_i{h}_j, \end{gathered}$$ (39)

where b is bias for visible units, b′ is the bias for hidden units, and W are the weights connecting visible and hidden units. Let the differential equation model be

$$\begin{gathered} \stackrel{.}{b}=F_b(b,b^{\prime },W;{\mathit{u}}_b), \\ {\stackrel{.}{b}}^{\prime }=F_{b^{\prime }}(b,b^{\prime },W;{\mathit{u}}_{b^{\prime }}), \\ \stackrel{.}{W}=F_W(b,b^{\prime },W;{\mathit{u}}_W). \end{gathered}$$ (40)

The corresponding graphical model is shown in Figs. 1(a) and 1(c).

The time-evolution functions for (38) and (40) are learned using Algorithm 1 and compared in Fig. 2. For the visible model, cells of size 0.5 × 0.5 × 0.5 in (b, J, K) are used, and for the hidden layer model cells of size 0.5 × 0.5 × 0.05 in (b, W, b′), as shown in Fig. 2.

FIG. 2.

Top row: Learned time-evolution functions for the fully visible model (38), using the Q₃, C₁ finite-element parametrization (34) with cells of size 0.5 × 0.5 × 0.5 in (b, J, K). Left panel: Training set of initial points (b, J, K) (cyan) sampled evenly in [−1, 1]. Stochastic simulations for each initial point are used as training data (learned trajectories shown in black, endpoints in magenta). Middle three panels: The time evolution functions learned, where the heat map indicates the value of F_γ in (38). Right panel: Vertices of the finite-element cells used. Bottom row: Hidden layer model (40) and parametrization (34) with cells of size 0.5 × 0.5 × 0.05 in (b, W, b′). Initial points are generated by BM learning applied to the points of the visible model. Note that the coefficients corresponding to the other seven degrees of freedom at each vertex are also learned (not shown), i.e., the first derivatives in each parameter.

As training data, 50 points (b, J, K) are sampled evenly over (b, J, K) ∈ [−1, 1]³. Each point corresponds to an initial distribution (37), from each of which 50 lattices of length N = 1000 are sampled (top left panel of Fig. 2). The corresponding initial conditions in (b, W, b′) space are learned separately using the BM learning algorithm (bottom left panel of Fig. 2). Each lattice is simulated for 200 time steps of size Δt = 0.01 with reaction probability p_r = 0.01 on encounters for the reaction A + A → ∅, as shown in Fig. 1(e). These trajectories are pooled for Algorithm 1. Note that a single set of parameter vectors {u} in (38) and (40) is learned, i.e., the parameter vectors are shared among trajectories from all initial conditions.

For the fully visible model, sleep phase moments are estimated by running a Gibbs sampler for a single step. Similarly, for the hidden model, wake and sleep phase moments are estimated by a single step of contrastive divergence (CD), i.e., CD-1. The learning rate used in both models is λ = 1 for 200 optimization steps.

The time integral in the action (28) can lead to undesired extrema, e.g., for periodic trajectories. We use an online algorithm to shift the limits of integration in (30) as new data are available:

$$\frac{dS}{d{u}_{k,i}}={\int }_{\tau }^{\tau +\Delta \tau }dt\frac{\partial F_k(\theta (t);{\mathit{u}}_k)}{\partial u_{k,i}}{\varphi }_k(t),$$ (41)

where Δτ is fixed and τ is gradually incremented t₀ ⩽ τ ⩽ t_f − Δτ. In this case, the PDE constraint (27) is solved from t₀ to τ, decreasing the size of the trajectories early in the training. Further, the adjoint system (31) only has to be solved backward from ϕ(τ + Δτ) = 0 to ϕ(τ), which also controls the magnitude of the update steps as the length of the trajectory grows, allowing a constant learning rate to be used. For the annihilation system, we found that fixing Δτ = 5 time steps and shifting τ → τ + 1 every two optimization steps gave fast convergence.

Figure 2 shows the learned time-evolution functions and trajectories of the training data. For the visible model, these show an expected symmetric structure. As particles diffuse and NN and NNN moments decay, F_J and F_K force J, K → 0 everywhere, while the bias term tends to negative infinity. The representation learned by the hidden layer model is more compact. Figure 3(a) shows the nearest-neighbor moment ⟨∑_i v_iv_i+1⟩ overlaid onto the initial conditions, showing an almost monotonic organization from low to high values by which the model can distinguish these states (no organization is apparent in the visible model). Figure 3(b) shows the learned parameter trajectories: b monotonically decreases (not shown), W asymptotically approaches a negative value, and b′ either increases monotonically or initially decreases before increasing again. This division corresponds to the decay of spatial correlations 2⟨δ_vivi+1⟩ − 1 (such that 1 corresponds to a fully correlated lattice and −1 to a fully anticorrelated lattice), also shown in Fig. 3(b). The two types of trajectories of b′ have a clear correspondence to two types of trajectories in the correlation function, and the separation is visible in F_b′ in the negative and positive regimes. We conclude that the moment closure approximation learned by the model therefore captures relevant low-range spatial correlations to approximate the right-hand sides of the moment equations (36) identified from the CME.

FIG. 3.

(a) NN moment ⟨∑_i v_iv_i+1⟩ of the two models. The more compact representation learned by the hidden layer model (left) captures low range spatial correlations, while the fully visible model (right) shows no apparent organization. (b) The parameters W and b′ for the hidden layer model for the 50 initial conditions (b is monotonically decreasing for all trajectories). The learned parameters encode the spatial correlation 2⟨δ_vivi+1⟩ − 1 shown on the right. This shows the moment closure approximation learned by the reduced model (see text). (c) RMSE in the third-order moment ⟨∑_i v_iv_i+1v_i+2⟩ and fourth-order moment ⟨∑_i v_iv_i+1v_i+2v_i+3⟩, calculated from a set of test trajectories (not shown). Both models reproduce the observables with reasonable accuracy, however, the error in the hidden layer model is lower due to the more compact representation learned.

To assess the accuracy of the reduced models, we generate a test set of points (b, J, K) and the learn the corresponding points (b, W, b′) as before. These are evolved in time using the learned DE systems (38) and (40). Define $\epsilon (t)=\sqrt{\langle [\mu (t)-\tilde{\mu }(t){]}^2\rangle }$ as the root-mean-square error (RMSE) between some moments of the reduced model $\tilde{\mu }$ and the stochastic simulations μ, where the moments are approximated by averaging over 50 samples. Figure 3(c) shows the RMSE for the third-order moment ⟨∑_i v_iv_i+1v_i+2⟩ and fourth-order moment ⟨∑_i v_iv_i+1v_i+2v_i+3⟩. Both models have relatively low error in reproducing the observables, however, the error in the hidden layer model is lower than in the visible model. This is because the representation learned by the hidden layer model is more compact, in that states initially distributed uniformly in (b, J, K) space are mapped to an approximately 1D curve in (b, W, b′) space. Yet higher accuracies may be possible by further tailoring that parametrizations of the differential equations from the cubic finite elements used here.

B. Learning the Rössler oscillator

The Williamowski-Rössler oscillator system [41] is a chemical version of a spiral oscillator in three species. The original formulation requires additional species that are fixed at constant concentration. Recent work [42], however, has developed a volume-excluding version where these constraints are incorporated into pseudo-first-order reaction rates, eliminating the need for additional reservoir populations. We follow this approach, such that the reaction system for species A, B, and C is

$$\begin{gathered} A\underset{p_1}{\overset{k_1}{\rightleftharpoons }}2A,A+B\stackrel{p_2}{\to }2B,A+C\stackrel{p_3}{\to }\varnothing , \\ B\stackrel{k_2}{\to }\varnothing ,C\underset{p_4}{\overset{k_3}{\rightleftharpoons }}2C, \end{gathered}$$ (42)

where the unimolecular reaction rates used are k₁ = 30, k₂ = 10, k₃ = 16.5 (arbitrary units), and the probabilities for bimolecular reactions are p₁ = 0.1, p₂ = 0.4, p₃ = 0.24, p₄ = 0.36. We simulate this system on a 3D lattice of size 10 × 10 × 10 sites in the single-occupancy limit as before. Figure 4 shows snapshots of such a stochastic simulation. Figure 4(b) in particular shows the characteristic shape of the Rössler oscillator, with further structures evident in higher-order moments shown in Fig. 4(c). A snapshot of the spatial waves that occur during transitions between A-, B-, and C-dominated regimes is shown in Fig. 4(a).

FIG. 4.

Rössler oscillator on a 3D lattice. (a) Snapshots of a stochastic simulation on a 10 × 10 × 10 lattice (A, B, and C in pink, orange, and cyan). (b) Moments from a single simulation over 500 time steps, producing a stochastic version of the characteristic attractor of the well-known deterministic model. (c) Nearest-neighbor moments in the simulation of (b) show similar structure. (d) Relaxation to a stationary distribution, indicated by the convergence of the means from averaging over 300 stochastic simulations.

The time evolution of the mean number of particles in A, B, and C, denoted by μ_α, is related to the number of nearest neighbors, denoted by Δ_αβ, as follows (see Appendix B for derivation):

$$\begin{gathered} \frac{d}{dt}{\mu }_A=k_1{\mu }_A-{\kappa }_1{\Delta }_{AA}-{\kappa }_2{\Delta }_{AB}-{\kappa }_3{\Delta }_{AC}, \\ \frac{d}{dt}{\mu }_B={\kappa }_2{\Delta }_{AB}-k_2{\mu }_B, \\ \frac{d}{dt}{\mu }_C=-{\kappa }_3{\Delta }_{AC}-k_3{\mu }_C-{\kappa }_4{\Delta }_{CC}, \end{gathered}$$ (43)

where κ₁, κ₂, κ₃, and κ₄ are the reaction rates for the bimolecular reactions specified by probabilities p₁, p₂, p₃, and p₄ above. As previously, this system is not closed, such that two close initial states in Fig. 4(b) will diverge over their long-term time evolution. The challenge for the latent variables in the reduced differential equation model is to incorporate relevant higher-order correlations to separate states which are close in their lower-order moments.

As in Sec. IV A, let the visible part of the graph be the lattice of Fig. 4(a). For the hidden layer, we choose a connectivity that coarse grains the visible lattice by one unit in each spatial dimension as shown in Fig. 5. Note that the hidden layer is also of size 10 × 10 × 10 units that implement periodic boundary conditions. The visible layer of the graph is multinomial in one of {A, B, C, ∅}, and similarly the hidden layer in {X, Y, Z, ∅}. The corresponding energy model is

$$\begin{gathered} E(V,H,\theta (t)) \\ =-\sum _i\sum _{\alpha \in \{A,B,C\}}{b}_{\alpha }{v}_{i,\alpha }-\sum _j\sum _{\alpha \in \{X,Y,Z\}}{b}_{\alpha }{h}_{j,\alpha } \\ -\sum _{\{i,j\}}(W_{AX}{v}_{i,A}{h}_{j,X}+W_{BY}{v}_{i,B}{h}_{j,Y}+W_{CZ}{v}_{i,C}{h}_{j,Z}), \end{gathered}$$ (44)

where H refers to the hidden layer and the sum over {i, j} implements the connectivity shown in Fig. 5 and

$$\stackrel{.}{\gamma }=F_{\gamma }(b_A,b_B,b_C;{\mathit{u}}_{\gamma })$$ (45)

for γ ∈ {b_A, b_B, b_C, W_AX, W_BY, W_CZ, b_X, b_Y, b_Z}. The right-hand side of the differential equation is parameterized (34) by cubic C₁ finite elements as before. To reduce the complexity of the model, we have purposefully omitted interactions W_AY, W_AZ, W_BX, W_BZ, W_CX, W_CY. With this choice, the latent species X coarse grains the visible species A, and similarly for Y, B and C, Z. Note that all differential equation models share the same domain in (b_A, b_B, b_C) space. While the biases h_A, h_B, h_C are the Lagrange multipliers corresponding to the constraints for the number of particles of each species, through the energy function (44) both biases and weights together control all spatial correlations of the model.

FIG. 5.

(a) Graph to learn for the Rössler oscillator. The lattice on the left corresponds to the visible layer, equivalent to the 10 × 10 × 10 cube in Fig. 4; the right corresponds to the hidden layer. Gray units in the hidden layer denote those units which implement periodic boundary conditions to the visible layer. (b) Connectivity of hidden layer. Each cube of eight neighboring units in the visible layer (green circles) is connected to a single unit (blue triangle) in the hidden layer (connections shown in red), resembling a body-centered cubic structure. Biases for the units are not shown.

Stochastic simulations are generated from an initial state with b_A = b_B = b_C = − ln(2), W_AX = W_BY = W_CZ = W_XY = W_YZ = 0, and b_X = b_Y = b_Z = − ln(1/7). By setting the initial weights to zero, this is the MaxEnt state given that the number of particles is μ_A = μ_B = μ_C = 200, since with zero weight:

$${\mu }_{\alpha }=1000\frac{e^{b_{\alpha }}}{1+\sum _{\beta =A,B,C}{e}^{b_{\beta }}}$$ (46)

for α ∈ {A, B, C}, and where the factor 1000 results from summing over all visible sites. With zero weight, the choice for the initial hidden layer bias is free—by choosing to set it to − ln(1/7), we are setting the target sparsity to approximately half of that of the visible layer with approximately 100 particles of each species as given by (46). Simulations are run for 500 time steps of size Δt = 0.01. Figure 4(d) shows the relaxation of the distribution to equilibrium [43].

For training, we use Algorithm 1 with learning rate λ = 0.05 for the weights and λ = 0.8 for the biases for 10 000 optimization steps. To estimate the wake phase moments, we sample $\tilde{p}(\mathit{H}=1\mid \mathit{V})$ for each sample in a batch size of η = 5, where V is a data vector. To estimate the sleep phase moments, we alternate between sampling $\tilde{p}[{\mathit{H}}^{(r)}=1\mid {\mathit{V}}^{(r)}]$ and $\tilde{p}[{\mathit{V}}^{(r)}=1\mid {\mathit{H}}^{(r-1)}]$ for r = 1, . . . , 10 steps, starting from a random configuration V⁽⁰⁾. Alternatively, we also found fast convergence using k = 10 steps of CD, as well as using PCD. To reduce the noise in the estimates, we use as is common raw probabilities instead of multinomial states for the hidden units when estimating both the wake and sleep phase moments.

As before, we use the online variant (41) of Algorithm 1 where the limits of integration are shifted during training, with window size Δτ = 10, and τ is gradually incremented τ → τ + 1 every 100 optimization steps. To learn smooth trajectories and avoid jumps in the learned differential equation model, each time step is divided into 10 substeps when solving the differential equations (44) and (45).

We compare the learned trajectories to a simplified MaxEnt model in Figs. 6(a)-6(c). The side length of the cubic finite elements used was 0.1 on all sides, centered at the initial condition, as shown in Fig. 6(d). Figure 6(a) shows the mean number of particles over the first 100 time steps, as in Fig. 4(d). Figure 6(b) transforms these points to the parameters (b_A, b_B, b_C) of a simple MaxEnt model constrained on these lowest-order moments as given by (46). Figure 6(c) shows the learned model (45), where the biases now control both the means and spatial correlations together with the weights. The trajectory no longer resembles a periodic trajectory, having learned to separate close states in Fig. 6(b). Figure 7 shows the learned time evolution functions for the Rössler oscillator over the first 100 time steps.

FIG. 6.

(a) The first 100 time steps of the mean number of A, B and C in the Rössler oscillator system. (b) Interaction parameters for a MaxEnt model constrained on the moments in (a) given by Eq. (46). (c) The learned trajectory of (44) in (b_A, b_B, b_C) space, with initial condition [− ln(2), − ln(2), − ln(2)]. The bias parameters have been tuned to control both the means and spatial correlations, together with the weights (not shown). Grayscale value indicates b_C component for clarity, scaled from dark [min(b_C)] to light [max(b_C)]. Initial point is shown in cyan, and endpoint in magenta. (d) Vertices of the finite-element cells of side length 0.1 used to parametrize the differential equations (45).

FIG. 7.

Learned time-evolution functions (45) in (b_A, b_B, b_C) space [see Fig. 6(d) for the vertices used], and the resulting trajectory in black [see Fig. 6(c)].

The agreement between the stochastic simulations and reconstructed observables is shown in Fig. 8(a). At each time point, 100 samples are drawn from the reduced model by running 25 steps of CD sampling, starting from a random configuration. Nearest neighbors, which determine the time evolution of the means in (43), are reasonably approximated, primarily due to the connectivity chosen in Fig. 5.

FIG. 8.

(a) Example of correlations learned by the reduced model compared to stochastic simulations, obtained by sampling over 100 samples. Top row: Mean number of A, B, C particles. Bottom: Neighboring pairs of (B, B), (C, C), and (A, B). Short range spatial correlations relevant to the moment equations (43) are reasonably approximated due to the chosen connectivity. (a) Sampled state V from the learned model (top left), and the activated hidden layer probabilities $\tilde{p}(\mathit{H}\mid \mathit{V})$ at time point 20. After training, the hidden layers coarse grain nearest neighbors in the visible layer.

Figure 8(b) shows a sampled state V from the learned model, and the activated hidden layer probabilities $\tilde{p}(\mathit{H}\mid \mathit{V})$ at time point 20. With the learned parameters, the hidden units coarse grain nearest neighbors in the lattice, as needed to approximate the right-hand side of (43). A deeper network such as a deep Boltzmann machine (DBM) may approximate yet-higher spatial correlations and can therefore be used to close differential equation systems depending on higher-order moments.

V. DISCUSSION

We have presented a learning problem for spatiotemporal distributions that estimates differential equation systems controlling a time-varying Boltzmann distribution. The ability to estimate a reduced physical model makes the method interesting for many modeling applications, including chemical kinetics as presented here. Mapping to a differential equation model can likewise be useful for engineering applications, allowing constraints to be efficiently introduced into BM learning as discussed in Sec. IIIA.

The moment closure approximation presented in Sec. II is broadly applicable due to the use of latent variables that can be trained to capture relevant higher-order correlations rather than deciding a priori what correlations to include as in typical closure schemes. Minimizing the KL divergence between the reduced and true models at all times is closely related to entropic matching but differs by the introduction of a differential equation system. We also make the connection to spatially continuous reaction systems explicit.

The finite-element parametrization is similar to the unsupervised learning setting of RBMs in the sense that it is independent of the system under consideration. For deeper architectures such as DBMs as discussed in Sec. IV B, recycling the same time-evolution functions across multiple layers may be effective, similarly to convolution layers in convolutional neural networks. Factoring weights has also been used effectively in deep learning [44] and may similarly reduce the computational burden here. The main advantage of the current DE formalism, however, is to use a parametrization (26) that enforces a physically relevant model.

We have illustrated the advantage of using latent variables in the learning problem, as opposed to a fully visible model. In the fully visible model of Sec. IV A, two and three particle correlations are explicitly captured. In the competing hidden layer model, we use a locally connected RBM (as opposed to fully connected layers) to control the range of correlations captured through the connectivity of the hidden layer. This has the advantage that the representation learned by the hidden layers is easily interpretable as it coarse grains the visible layer. Further, the local connectivity used can be inferred from the moment equations derived from the CME. Deeper networks with multiple hidden layers can be constructed in this fashion to learn hierarchical statistics, with the ability to infer long-range spatial correlations that may become relevant over long timescales.

A popular alternative class of generative models to RBMs are variational autoencoders (VAEs). An adaptation of the proposed method may be possible for these models; however, the main advantages of the current RBM framework is that the form of the energy function can be used interpret the reduced model [22] and that the distribution over the latent variables is not chosen as in VAEs (typically a standard normal distribution) but rather learned from data.

A closely related problem to model reduction is the problem of data assimilation, where noisy measurements and an incomplete model for the dynamics are combined to estimate the true state of the system and unknown parameters in the model [45]. Model reduction methods complement the data assimilation problem by replacing the physical model with a reduced one which can increase the efficiency of data assimilation methods.

We view the present work as progress toward linking models across scales in biology [20]. Reaction-diffusion systems illustrate many of the common problems in this field. While much machinery (CME or field-theoretic methods) exist to formulate problems for observables, their solution is nontrivial in most applications. Even without analytic challenges such as moment closure, the numerical solution of PDE systems is difficult for systems with high spatial organization or where interactions with other scales (e.g., molecular dynamics) or physics (e.g., electrodiffusion) become relevant. Learning reduced models in the form of spatial dynamic Boltzmann distributions may abstract many of these nontrivial interactions.

APPENDIX A: FORMAL SOLUTION FOR THE ADJOINT SYSTEM

The connection between (6) and (18) can be made more explicitly. A differential equation system for the perturbations $\delta {\nu }_k({\alpha }_{\langle i{\rangle }_k^n},{\mathit{x}}_{\langle i{\rangle }_k^n},t)$ in (6) can be derived by linearizing the differential equation around a particular solution [22,30]. For the autonomous system (17), this leads to the linear ODE system:

$$\frac{d}{dt}\delta \nu (\alpha ,\mathit{x},t)=\delta \mathit{F}(\alpha ,\mathit{x},t)+G(\alpha ,\mathit{x},t)\delta \nu (\alpha ,\mathit{x},t),$$ (A1)

with some given initial condition δν(α, x, t₀) = δη(α, x). Here we have used the vector notation introduced in Sec. IIB.

Let the homogenous part of this system

$$\frac{d}{dt}\delta \nu (\alpha ,\mathit{x},t)=G(\alpha ,\mathit{x},t)\delta \nu (\alpha ,\mathit{x},t)$$ (A2)

have solution given by the nonsingular fundamental matrix A(α, x, t). Then (A1) has as formal solution

$$\delta \nu (\alpha ,\mathit{x},t)=A(\alpha ,\mathit{x},t)\left[\delta \eta (\alpha ,\mathit{x})+{\int }_{t_0}^{t}d{t}^{\prime }{A}^{-1}(\alpha ,\mathit{x},t^{\prime })\delta \mathit{F}(\alpha ,\mathit{x},t^{\prime })\right],$$ (A3)

which substituted into (6) gives:

$$\begin{gathered} \delta S={\int }_{t_0}^{t_f}dt\sum _{n=0}^{\infty }\sum _{\alpha }\int d\mathit{x}\Delta {\mu }^{⊺}(\alpha ,\mathit{x},t)A(\alpha ,\mathit{x},t) \\ \times \left[\delta \eta (\alpha ,\mathit{x})+{\int }_{t_0}^{t}d{t}^{\prime }{A}^{-1}(\alpha ,\mathit{x},t^{\prime })\delta \mathit{F}(\alpha ,\mathit{x},t^{\prime })\right]=0, \end{gathered}$$ (A4)

where Δμ^⊺(t) is the vector with components (7). Applying integration by parts on the term in parentheses to move the integral over time gives

$$\begin{gathered} \left[\delta \eta (\alpha ,\mathit{x})+{\int }_{t_0}^{t_f}d{t}^{\prime }{A}^{-1}(\alpha ,\mathit{x},t^{\prime })\delta \mathit{F}(\alpha ,\mathit{x},t^{\prime })\right] \\ \times {\phantom{\mid }\left[{\int }_{t_0}^{t}d{t}^{\prime }\Delta {\mu }^{⊺}(\alpha ,\mathit{x},t^{\prime })A(\alpha ,\mathit{x},t^{\prime })\right]\mid }_{t=t_0}^{t_f} \\ -{\int }_{t_0}^{t_f}dt{\int }_{t_0}^{t}d{t}^{\prime }\Delta {\mu }^{⊺}(\alpha ,\mathit{x},t^{\prime })A(\alpha ,\mathit{x},t^{\prime }) \\ \times A^{-1}(\alpha ,\mathit{x},t)\delta \mathit{F}(\alpha ,\mathit{x},t), \end{gathered}$$ (A5)

where the adjoint functions ζ(t) can be identified as:

$${\zeta }^{⊺}(\alpha ,\mathit{x},t)={\int }_{t_0}^{t}d{t}^{\prime }\Delta {\mu }^{⊺}(\alpha ,\mathit{x},t^{\prime })A(\alpha ,\mathit{x},t^{\prime })A^{-1}(\alpha ,\mathit{x},t).$$ (A6)

By choosing the adjoint functions to satisfy the boundary condition ζ(α, x, t_f) = 0, the boundary term in (A5) vanishes and we obtain the previous result (16).

APPENDIX B: DERIVATION OF MOMENT EQUATIONS FROM THE CHEMICAL MASTER EQUATION

The moment equations (36) and (43) can be derived from the chemical master equation using the Doi-Peliti [46] formalism and its equivalent generating function representation. We demonstrate this for the Rössler system (43).

For notational convenience, we do not consider the single-occupancy limit here. The state of the system is described by the N × M matrix V′ with entries v_i,α ∈ {0, 1, 2, . . . }, where N = 10 × 10 × 10 rows denote lattice sites, and M = 3 columns denote occupancies of species {A, B, C}.

Define the N × M single-entry matrix e_ij with entries zero everywhere except at index (i, j) where it is one. The creation and annihilation operators ${\widehat{a}}_{i,\alpha }$ and a_i,α create and destroy particles of species α at unit i:

$$\begin{gathered} {\widehat{a}}_{i,\alpha }\mid {\mathit{V}}^{\prime }\rangle =\mid {\mathit{V}}^{\prime }+{\mathit{e}}_{i,\alpha }\rangle , \\ {a}_{i,\alpha }\mid {\mathit{V}}^{\prime }\rangle =v_{i,\alpha }\mid {\mathit{V}}^{\prime }-{\mathit{e}}_{i,\alpha }\rangle . \end{gathered}$$ (B1)

The operators corresponding to reactions in the Rössler system (excluding diffusion) are then:

$$\begin{gathered} A\to 2A:k_1\sum _{i=1}^N({\widehat{a}}_{i,A}-1){\widehat{a}}_{i,A}{a}_{i,A}, \\ 2A\to A:{\kappa }_1\sum _{\langle ij\rangle }(1-{\widehat{a}}_{j,A}){\widehat{a}}_{i,A}{a}_{i,A}{a}_{j,A}, \\ A+B\to 2B:{\kappa }_2\sum _{\langle \langle ij\rangle \rangle }({\widehat{a}}_{i,B}-{\widehat{a}}_{i,A}){\widehat{a}}_{j,B}{a}_{i,A}{a}_{j,B}, \\ A+C\to \varnothing :{\kappa }_3\sum _{\langle \langle ij\rangle \rangle }(1-{\widehat{a}}_{i,A}{\widehat{a}}_{j,C})a_{i,A}{a}_{j,C}, \\ B\to \varnothing :k_2\sum _{i=1}^N(1-{\widehat{a}}_{i,B})a_{i,B}, \\ C\to 2C:k_3\sum _{i=1}^N({\widehat{a}}_{i,C}-1){\widehat{a}}_{i,C}{a}_{i,C}, \\ 2C\to C:{\kappa }_4\sum _{\langle ij\rangle }(1-{\widehat{a}}_{j,C}){\widehat{a}}_{i,C}{a}_{i,C}{a}_{j,C}, \end{gathered}$$ (B2)

where ∑_{⟨i j⟩} sums over all neighboring sites without double counting, ∑_{⟪i j⟫} sums over all neighboring sites with double counting, and we specify the species {A, B, C} instead of an index α = 1, . . . , M for clarity in the subscripts. Here we place new particles resulting from fission reactions with rates k₁ and k₃ at the same site - in the single-occupancy limit, they must be placed at a neighboring site. For bimolecular reactions with rates κ₁ and κ₄, we make the in this case ambiguous choice to place new species at site i versus j. The time evolution operator W for the Rössler system is the sum of all terms in (B2).

The system state and the ladder operators admit an equivalent generating function representation:

$$\mid {\mathit{V}}^{\prime }\rangle \to \prod _{i=1}^N\prod _{\alpha =1}^M{z}_{i,\alpha }^{v_{i,\alpha }},{\widehat{a}}_{i,\alpha }\to z_{i,\alpha },a_{i,\alpha }\to \frac{\partial }{\partial z_{i,\alpha }}.$$ (B3)

An observable ⟨X⟩ with generating function representation X_z according to (B3) evolves as:

$$\frac{d\langle X\rangle }{dt}={\phantom{\mid }\left(X_{z}W\prod _{i=1}^N\prod _{\alpha =1}^M{z}_{i,\alpha }^{v_{i,\alpha }}\right)\mid }_{z=1},$$ (B4)

where W is now the sum of terms (B2) in the generating function representation (B3). From the number operator ${\widehat{a}}_{k,\beta }{a}_{k,\beta }$ which counts the number of particles of species β at position k, the time evolution of the mean number of particles of species β is then

$$\frac{d{\mu }_{\beta }}{dt}={\phantom{\mid }\left(\sum _{k=1}^N{z}_{k,\beta }\frac{\partial }{\partial z_{k,\beta }}W\prod _{i=1}^N\prod _{\alpha =1}^M{z}_{i,\alpha }^{v_{i,\alpha }}\right)\mid }_{z=1},$$ (B5)

which can be directly evaluated to give the moment equations (43). For a review on field theoretic methods for reaction-diffusion systems, we refer to Mattis and Glasser [47]. The formalism can also describe systems in continuous space [46] where it has a similar generation function representation [22].