# Physics-informed graph neural networks enhance scalability of variational
nonequilibrium optimal control

Jiawei Yan (闫嘉伟<EMAIL_ADDRESS>Department of Chemistry, Stanford
University, Stanford, CA 94305, USA  Grant M. Rotskoff<EMAIL_ADDRESS>Department of Chemistry, Stanford University, Stanford, CA 94305, USA

###### Abstract

When a physical system is driven away from equilibrium, the statistical
distribution of its dynamical trajectories informs many of its physical
properties. Characterizing the nature of the distribution of dynamical
observables, such as a current or entropy production rate, has become a
central problem in nonequilibrium statistical mechanics. Asymptotically, for a
broad class of observables, the distribution of a given observable satisfies a
large deviation principle when the dynamics is Markovian, meaning that
fluctuations can be characterized in the long-time limit by computing a scaled
cumulant generating function. Calculating this function is not tractable
analytically (nor often numerically) for complex, interacting systems, so the
development of robust numerical techniques to carry out this computation is
needed to probe the properties of nonequilibrium materials. Here, we describe
an algorithm that recasts this task as an optimal control problem that can be
solved variationally. We solve for optimal control forces using neural network
ansätze that are tailored to the physical systems to which the forces are
applied. We demonstrate that this approach leads to transferable and accurate
solutions in two systems featuring large numbers of interacting particles.

gbsn

##  I Introduction

At equilibrium, a physical system relaxes to a time-invariant distribution
with a statistical weight that can, in many cases, be computed easily.
Straightforward characterization of the relative probabilities of distinct
states facilitates the use of Markov Chain Monte Carlo algorithms, and, as a
result, expectations of observables are tractable to compute in many cases.
When we drive a system away from equilibrium, either through the application
of an external perturbation or with an energy-consuming dynamics that
inherently breaks detailed balance, we lose the luxury of this facile
description and knowledge of the universal distribution of states. While it is
possible in some cases to characterize steady-state distributions of complex,
interacting, nonequilibrium systems Toner and Tu (1995, 1998); Lazarescu
(2015); Martin  _et al._ (2021); Fodor, Jack, and Cates (2022); Cates and
Tailleur (2015), as an alternative, we adopt a manifestly _dynamical_ point of
view.

The large deviation perspective on nonequilibrium statistical mechanics
focuses less on static states of a system and more on time-dependent
trajectories Garrahan  _et al._ (2007); Touchette (2009). Ruelle developed a
thermodynamic formalism Ruelle (2009) based on Donsker-Varadhan large
deviation theory that has become a widely-used framework to understand the
statistical fluctuations of dynamical behavior, from glassiness Garrahan  _et
al._ (2007, 2009); Hedges  _et al._ (2009), to dynamical phase transitions
Gingrich, Vaikuntanathan, and Geissler (2014); GrandPre  _et al._ (2021), to
thermodynamic uncertainty relations Barato and Seifert (2015); Gingrich  _et
al._ (2016). The central object of study within this theory is the scaled
cumulant generating function (SCGF) for a dynamical observable $A_{T}.$ This
object is defined as the limit of an expectation over trajectories,

$\psi(\lambda)=\lim_{T\to\infty}\frac{1}{T}\log\mathbb{E}_{\boldsymbol{X}}e^{\lambda
TA_{T}},$ (1)

plays a role akin to that of a free energy for an equilibrium system. The
Legendre-Fenchel transform of $\psi(\lambda)$ defines the _large deviation
rate function_ $I(a)$, which describes the rate of decay of the probability
that $A_{T}\in[a,a+da]$

$P(A_{T}\in[a,a+da])\asymp e^{-TI(a)}da.$ (2)

These two objects contain identical statistical information in the long-time
limit, but both (1) and (2) are notoriously difficult to compute because it
involves exponential averages that are dominated by rare events.

Here, we outline a robust and transferable strategy for estimating (1) that
applies to two broad classes of interacting models: interacting particles
systems with diffusive dynamics and Markov jump processes where local
interactions dictate jump rates. In both cases, we construct a neural network
ansatz for _virtual_ control forces that realize the rare trajectories
necessary to compute (1). Our results demonstrate that carefully incorporating
physical symmetries into the architecture of the neural network ansatz leads
to an accurate, straightforward algorithm that can be deployed on a wide
variety of systems.

The present work builds on recent developments using machine learning Yan,
Touchette, and Rotskoff (2022); Das  _et al._ (2021); Casert  _et al._ (2021),
that provide a generic framework for learning optimal control forces to
minimize a variational objective for (1). These control forces guide the
dynamics to realize the statistically rare fluctuations that lead to the
corresponding rare values of the observable $A_{T}$. Here, we demonstrate that
this algorithm can be specifically tailored for interacting systems with
appropriate inductive biases. Physical constraints, such as translational
invariance, are built into the neural network to achieve fast and transferable
training. The resulting control forces efficiently define a many-body control
force to achieve rare fluctuations. One key advantage to the approach taken
here is that, unlike existing approaches, the new neural network architecture
we employ allows _transfer training_ , in which a small, easy-to-simulate
system is used to generate training data, but the networks are subsequently
deployed on larger systems.

###  I.1 Existing numerical approaches

Many numerical algorithms have been proposed to access the rare dynamical
fluctuations that determine a large deviation rate function. These algorithms
vary in their approach; those based on importance sampling Bucklew (2004);
Ray, Chan, and Limmer (2017); Ferré and Touchette (2018) sample a dynamical
process distinct from the original process, but re-weight those samples
accordingly. On the other hand, cloning algorithms Grassberger (2002);
Giardinà, Kurchan, and Peliti (2006); Lecomte and Tailleur (2007); Ray and
Chan (2020) implement a population dynamics at the level of
trajectories—trajectories are killed and cloned with to re-weight expectation
values with the exponential statistical bias in the scaled cumulant generating
function. State-of-the-art methods employ a combination of cloning and
adaptive sampling Nemoto  _et al._ (2016) using parameter dependent potential
functions to guide the birth-death dynamics. Along these lines, a series of
papers has employed genetic algorithms Jacobson and Whitelam (2019); Whitelam,
Jacobson, and Tamblyn (2020); Whitelam and Tamblyn (2020); Casert  _et al._
(2021); Sprague  _et al._ (2021) and has shown that this approach can
powerfully exploit the variational nature of the problem.

Numerically exact approaches are tractable for sufficiently small state
spaces; the limit in (1) can be computed by obtaining the maximum eigenvalue
of a $\lambda$-dependent operator Lebowitz and Spohn (1999). For some lattice
models, this eigenvalue problem is therefore amenable to treatment using
matrix product states (MPS) and numerical techniques based on density matrix
renormalization group (DMRG) yield very accurate results when used with a
tensor network ansatz Helms, Ray, and Chan (2019); Helms and Chan (2020);
Bañuls and Garrahan (2019). Of course, the systems for which DMRG can be used
are inherently limited and it remains challenging to extend these methods to
models with continuous state spaces like interacting particle systems.

Machine learning approaches to this problem have relied for the most part on
Malliavin weight sampling Das and Limmer (2019) or reinforcement learning
Rose, Mair, and Garrahan (2021); Das  _et al._ (2021) to parameterize
polynomial force functions that drive sampling. Most closely related to the
present work Casert  _et al._ (2021); Yan, Touchette, and Rotskoff (2022) use
a neural network to represent a biasing force.

Our approach to estimating large deviation rate functions falls broadly under
the umbrella of importance sampling for dynamical trajectories, a topic on
which there has been substantial progress in recent years. Improvements in
transition path sampling Das, Kuznets-Speck, and Limmer (2022) and
nonequilibrium sampling methods Rotskoff and Vanden-Eijnden (2019); Thin  _et
al._ (2021). Also related to the present work, machine learning algorithms
dedicated to learning a force field or many-body potential in interacting
particle systems have become a topic of intense study, including Zhang  _et
al._ (2018); Schütt  _et al._ (2018); Wang  _et al._ (2018); Frishman and
Ronceray (2020). Many of these works incorporate physical inductive biases
that facilitate the effort to construct a force that preserves physical
invariances and equivariances Zhang  _et al._ (2018); Satorras, Hoogeboom, and
Welling (2021); Schütt  _et al._ (2018).

##  II Sampling as an optimal control problem

In this paper we consider interacting particle systems described by a
Markovian dynamics, which could either be a Markov jump process (cf. Appendix
A) or a stochastic differential equation (SDE):

$d\boldsymbol{X}_{t}=b(\boldsymbol{X}_{t})dt+\sigma d\boldsymbol{W}_{t},$ (3)

where $\boldsymbol{X}_{t}\in\mathbb{R}^{d}$ is the state of the system,
$b:\mathbb{R}^{d}\to\mathbb{R}^{d}$ is the drift function, and
$\boldsymbol{W}_{t}$ is a Wiener process acting as a noise source, which is
multiplied by a matrix $\sigma$. For systems in equilibrium $b$ is the
gradient of a potential energy function describing the interactions among the
particles, but for systems out of equilibrium it may have non-conservative
contributions. For simplicity, we assume that $\sigma$ is independent of space
and that the corresponding diffusion tensor $D=\sigma\sigma^{T}$ is
invertible. Moreover, we assume that $\boldsymbol{X}_{t}$ is ergodic, which
means that it has a unique stationary probability density, reached from any
initial distribution in the long-time limit.

For nonequilibrium systems, fluctuations in time-extensive observables, such
as currents, entropy production, and work yield information about dynamical
response Maes (2020) and constraints on the dissipation rate Barato and
Seifert (2015); Gingrich  _et al._ (2016). Following the now standard
formulation Touchette (2009), we consider a class of dynamical observables
$A_{T}$ which has the following generic form:

$A_{T}=\frac{1}{T}\int_{0}^{T}f(\boldsymbol{X}_{t})dt+\frac{1}{T}\int_{0}^{T}g(\boldsymbol{X}_{t})\circ
d\boldsymbol{X}_{t},$ (4)

where $f:\nobreak\mathbb{R}^{d}\to\mathbb{R}$ we call “density-like” and
$g:\mathbb{R}^{d}\to\mathbb{R}^{d}$ we call “current-like”. Depending on the
choices of these two functions, (4) can represent a large class of physical
quantities, including the current, dynamical activity, work, and entropy
production Touchette (2018). Computing the statistical distribution of $A_{T}$
is, of course, not analytically tractable in all but the simplest systems,
however it often satisfies a large deviation principle:

$p(A_{T}=a)\asymp e^{-TI(a)}$ (5)

where $I(a)$ is called the “rate function” and the symbol $\asymp$ denotes
equivalence up to logarithmic corrections. The asymptotic expression (5) can
be regarded as a generalization of the central limit theorem—Gaussian
fluctuations around the mean value lead to a quadratic expression for $I(a)$.
To compute $I(a)$ we can equivalently evaluate the Legendre-Fenchel transform
of the scaled cumulant generating function (SCGF) $\psi(\lambda)$:

$\psi(\lambda)=\lim_{T\to\infty}\frac{1}{T}\log\mathbb{E}_{\boldsymbol{X}}e^{\lambda
TA_{T}},$ (6)

where $\mathbb{E}_{\boldsymbol{X}}$ denotes an ensemble average over
trajectories generated from (3), and $\lambda$ is a real number. When
$\lambda\neq 0$, the average in (6) is dominated by trajectories that realize
rare values of the observable $A_{T}$, and, as a result, it is difficult to
evaluate $\psi(\lambda)$ by directly simulating trajectories.

###  II.1 Rigorous formulation of Active Importance Sampling

Importance sampling can ameliorate the sampling problem required to compute
(6) to some extent, but it is often difficult to use because one must
determine an appropriate biased sampling distribution. In the context of
sampling dynamical trajectories, we must simulate the dynamics of an
alternative SDE

$d\boldsymbol{X}^{u}_{t}=u_{t}(\boldsymbol{X}^{u}_{t})dt+\sigma
d\boldsymbol{W}_{t},$ (7)

and then the SCGF can be estimated simply by reweighting the average

$\psi(\lambda)=\lim_{T\to\infty}\frac{1}{T}\log\mathbb{E}_{\boldsymbol{X}^{u}}\left(e^{\lambda
TA_{T}}\frac{\mathrm{d}\mathbb{P}[\boldsymbol{X}^{u}]}{\mathrm{d}\mathbb{P}_{u}[\boldsymbol{X}^{u}]}\right).$
(8)

Here $\mathbb{E}_{\boldsymbol{X}^{u}}$ denotes the fact that the expectation
is taken over trajectories from (7) rather than the original dynamics (3).

The main task here is to identify an optimal or near optimal control force
$u_{t}(\boldsymbol{X}^{u}_{t})$ which can turn the rare event we aim to sample
into a typical one. To achieve this goal, we adapt the weak convergence
approach developed by Dupuis and Wang (2004) and extend it to continuous-time
Markov processes. Let $\mathrm{d}\mathbb{P}_{\nu}$ be a path measure on
$[0,T]$ such that its Radon-Nikodym derivative with respect to the path
measure of the original dynamics is

$\frac{\mathrm{d}\mathbb{P}_{\nu}}{\mathrm{d}\mathbb{P}}(\omega)=\frac{e^{\lambda
TA_{T}}}{\mathbb{E}[e^{\lambda TA_{T}}]}.$ (9)

Then for any fixed $T$, the Kullback-Leibler (KL) divergence between two path
measure $\mathrm{d}\mathbb{P}_{u}$ and $\mathrm{d}\mathbb{P}$ is

$\displaystyle\frac{1}{T}\mathcal{D}_{KL}[\mathrm{d}\mathbb{P}_{u}\|\mathrm{d}\mathbb{P}]=$
$\displaystyle\frac{1}{T}\int\log\frac{\mathrm{d}\mathbb{P}_{u}}{\mathrm{d}\mathbb{P}}(\omega)\mathrm{d}\mathbb{P}_{u}(\omega),$
$\displaystyle=$
$\displaystyle\frac{1}{T}\int\log\frac{\mathrm{d}\mathbb{P}_{u}}{\mathrm{d}\mathbb{P}_{\nu}}(\omega)\mathrm{d}\mathbb{P}_{u}(\omega)+\int\log\frac{\mathrm{d}\mathbb{P}_{\nu}}{\mathrm{d}\mathbb{P}}(\omega)\mathrm{d}\mathbb{P}_{u}(\omega),$
$\displaystyle=$
$\displaystyle\frac{1}{T}\mathcal{D}_{KL}[\mathrm{d}\mathbb{P}_{u}\|\mathrm{d}\mathbb{P}_{\nu}]+\lambda\mathbb{E}_{u}[A_{T}]-\frac{1}{T}\log\mathbb{E}[e^{\lambda
TA_{T}}].$ (10)

Simply rearranging the expression (10), we see that

$\frac{1}{T}\log\mathbb{E}[e^{\lambda
TA_{T}}]-\frac{1}{T}\mathcal{D}_{KL}[\mathrm{d}\mathbb{P}_{u}\|\mathrm{d}\mathbb{P}_{\nu}]=\lambda\mathbb{E}_{u}[A_{T}]-\frac{1}{T}\mathcal{D}_{KL}[\mathrm{d}\mathbb{P}_{u}\|\mathrm{d}\mathbb{P}].$
(11)

Assuming all the limits exist, we take the limit $T\rightarrow\infty$ and then
maximize over $u$ on both sides:

$\displaystyle\sup_{u}\left\\{\lim_{T\rightarrow\infty}\frac{1}{T}\log\mathbb{E}[e^{\lambda
TA_{T}}]-\lim_{T\rightarrow\infty}\frac{1}{T}\mathcal{D}_{KL}[\mathrm{d}\mathbb{P}_{u}\|\mathrm{d}\mathbb{P}_{\nu}]\right\\}$
(12) $\displaystyle=$
$\displaystyle\sup_{u}\left\\{\lim_{T\rightarrow\infty}\mathbb{E}_{u}[A_{T}]-\lim_{T\rightarrow\infty}\frac{1}{T}\mathcal{D}_{KL}[\mathrm{d}\mathbb{P}_{u}\|\mathrm{d}\mathbb{P}]\right\\}$
(13)

Since the KL-divergence is always non-negative and the first term on the left
is independent with $u$, the left side reaches its supremal value only when
the KL-divergence term vanishes, therefore,

$\displaystyle\psi(\lambda)=\lim_{T\rightarrow\infty}\frac{1}{T}\log\mathbb{E}[e^{\lambda
TA_{T}}]=\sup_{u}\lim_{T\rightarrow\infty}\left\\{\mathbb{E}_{u}[A_{T}]-\frac{1}{T}\mathcal{D}_{KL}[\mathrm{d}\mathbb{P}_{u}\|\mathrm{d}\mathbb{P}]\right\\}.$
(14)

The derivations above indicate that the optimal control force $u^{*}$ has the
property

$\frac{1}{T}\mathcal{D}_{KL}[\mathrm{d}\mathbb{P}_{u^{*}}\|\mathrm{d}\mathbb{P}_{\nu}]=0,$
(15)

so $\mathrm{d}\mathbb{P}_{u^{*}}$ is equivalent with the path measure
$\mathrm{d}\mathbb{P}_{\nu}$ we construct in (9). However, the Radon-Nikodym
derivative does not necessarily imply the explicit form of the controlled
process (7), nevertheless to mention that knowing the explicit form of (9)
requires knowing $\mathbb{E}[e^{\lambda TA_{T}}]$, which is exactly we aim to
compute. Therefore, we cannot direct construct a control force to match (9)
but we can optimize (14) to determine the control force.

The KL-divergence in (14) between path measure of the controlled dynamics (7)
and the original dynamics (3) can be evaluated explicitly since (7) and (3)
have the same noise tensor $\sigma$. According to Girsanov theorem, the Radon-
Nikodym derivative is determined by the Martingale process Oksendal (1992)

$\frac{\mathrm{d}\mathbb{P}_{u}}{\mathrm{d}\mathbb{P}}=\exp\left\\{\int_{0}^{T}(u-b)\sigma^{-1}dW_{t}+\frac{1}{2}\int_{0}^{T}(u-b)^{\top}D^{-1}(u-b)dt\right\\},$
(16)

and thus the stochastic integral term vanishes when taking the expectation
over the path space. Therefore,

$\displaystyle\frac{1}{T}\mathcal{D}_{KL}[\mathrm{d}\mathbb{P}_{u}\|\mathrm{d}\mathbb{P}]=\frac{1}{2T}\mathbb{E}_{u}\left[\int_{0}^{T}(u-b)^{\top}D^{-1}(u-b)dt\right].$
(17)

##  III Building physically motivated representations of optimal control
forces

Figure 1: ($a$) The information gathered by the graph neural network is
limited to a particle’s local environment, defined by a cutoff distance
$d_{\rm c}$. ($b$) The graph neural network architecture, where $f_{g}$,
$f_{e}$, and $f_{\theta}$ are three independent multilayer perceptrons (see
(22)).

Computing SCGF through optimizing (14) can be carried by representing the
control force $u$ by a deep neural network. We have previously shown that this
approach can obtain accurate results with significantly lower computational
overhead than adaptive cloning methods Yan, Touchette, and Rotskoff (2022).
However, interacting particle systems pose a particular challenge because
probing collective effects such as phase transitions require simulations with
hundreds of particles. Though it is still possible to use one large deep
neural network to represent the control force of the entire system, i.e., all
particles, training becomes very difficult as the number of particle
increases, and also the memory usage becomes extensive. One solution to this
problem is to leverage the fact that the control force for any given particle
should depend most strongly on its local environment, i.e., the relative
positions of other particles nearby, as illustrated in Fig. 1 ($a$). In
addition, the control force should also be restricted by physical laws and
constraints, including translational invariance, permutation invariance, and
equivariance. With these constraints, we here propose a representation of the
control forces that is tailored for representing the control force for
interacting particle systems.

As mentioned above, we assume that the control force $u$ depends only on one
particle’s local environment, and should be constrained by physical laws and
invariance including:

  * •

Permutation invariance: the control force for a given particle remains the
same when exchanging the index of its two neighboring particles

  * •

Translational invariance: the control force remains the same when shifting the
coordinates by a constant vector

  * •

Equivariance: the direction of the control force rotates by the same angle as
the coordinates rotate.

A general neural network, e.g., a multilayer perceptron (MLP), does not
automatically satisfy the constraints outlined above, though it is possible to
approximately realize these properties by augmenting the data to learn this
physical information. A more natural approach is to incorporate a physical
inductive bias directly into the architecture of the neural network, sometimes
called physics-informed machine learning Karniadakis  _et al._ (2021). We use
a representation that satisfies all the three properties above, without
sacrificing the expressiveness of deep neural networks. The local environment
of a given particle, defined by a cutoff distance $d_{\rm c}$, limits the
spatial extent of interactions that determine the control force. For particle
$i$, the control force is a function of all the neighboring particles $j$ such
that $|\boldsymbol{r}_{j}-\boldsymbol{r}_{i}|\leqslant d_{\rm c}$ (see Fig. 1
for illustrations).

As described above, a generic neural network does not maintain equivariance of
the forces, therefore we compute the direction and magnitude of the control
force separatelySatorras, Hoogeboom, and Welling (2021). Furthermore, the
output of the neural network should be insensitive to any re-ordering of
neighboring particles to ensure permutation invariance. We use a message-
passing graph neural network architecture, illustrated in Fig. 1(b), in which
the messages

$m_{ij}=f_{\rm
e}(\|\boldsymbol{x}_{j}-\boldsymbol{x}_{i}\|)\cdot\phi(\|\boldsymbol{x}_{j}-\boldsymbol{x}_{i}\|)$
(18)

are passed between the particle labelled $i$ and neighboring particles indexed
by $j$. The function $\phi(x)$ is a smooth decay function in the range of
$[0,d_{\rm c}]$ satisfying $\phi(0)=1$ and $\phi(d_{\rm c})=0$ since the
control force should smoothly decay to $0$ as it reaches the cutoff distance.
In the present work, we use $\phi(x)=\sqrt{[\cos(\pi x/d_{\rm c})+1]/2}$,
though others functions may also function well. The function $f_{\rm e}$
represents some neural network, here we also use an MLP. The magnitude of the
control force is parameterized by an additional neural network which takes as
an input the sum over messages,

$|\boldsymbol{u}_{i}|=f_{g}(\sum_{j}m_{ij})=f_{g}\left[\sum_{j}f_{e}(\|\boldsymbol{x}_{j}-\boldsymbol{x}_{i}\|)\cdot\phi(\|\boldsymbol{x}_{j}-\boldsymbol{x}_{i}\|)\right].$
(19)

We include the sum over messages to maintain permutation invariance.

A rotationally equivariant transformation can be implemented by parameterizing
an orthogonal matrix that depends only on the inter-particle distance; we use

$\frac{\boldsymbol{u}_{i}}{|\boldsymbol{u}_{i}|}=\sum_{j}R\left[f_{\theta}(\|\boldsymbol{x}_{j}-\boldsymbol{x}_{i}\|)\right]\frac{\boldsymbol{x}_{j}-\boldsymbol{x}_{i}}{\|\boldsymbol{x}_{j}-\boldsymbol{x}_{i}\|},$
(20)

where $R(\cdot)$ is the rotation matrix

$R(\theta)=\begin{bmatrix}\cos(\theta)&-\sin(\theta)\\\
\sin(\theta)&\cos(\theta)\end{bmatrix}.$ (21)

Rotation matrices in two-dimensions are normal matrices, i.e.,
$R(\theta_{1})R(\theta_{2})=R(\theta_{2})R(\theta_{1})$ for any two
$\theta_{1}$ and $\theta_{2}$, if we rotate the coordinate by an angle
$\theta$, the direction of the control force will also rotate by an angle
$\theta$. The message-passing framework we use is similar to Satorras,
Hoogeboom, and Welling (2021) but we include the rotation matrix (21) to allow
more expressive representations of the direction of the control force. Because
we use the commutativity of rotation matrices in two dimensions, the
representation we use will require further development to be extended to
three-dimensional inputs. For those systems, how to construct a flexible
representation of the direction of the control force while remaining
equivariant is a challenge for future studies.

The total control force on particle $i$ can be represented by a combination of
three neural networks $f_{\theta}$, $f_{g}$ and $f_{e}$,

$u_{i}=\sum_{j}R[\phi_{\theta}(\|\boldsymbol{x}_{j}-\boldsymbol{x}_{i}\|)]\frac{\boldsymbol{x}_{j}-\boldsymbol{x}_{i}}{\|\boldsymbol{x}_{j}-\boldsymbol{x}_{i}\|}\cdot
f_{g}\left[\sum_{j}f_{e}(\|\boldsymbol{x}_{j}-\boldsymbol{x}_{i}\|)\right]$
(22)

where $f_{\theta}:\mathbb{R}^{1}\rightarrow\mathbb{R}^{1}$,
$f_{e}:\mathbb{R}^{1}\rightarrow\mathbb{R}^{h}$ and
$f_{g}:\mathbb{R}^{h}\rightarrow\mathbb{R}^{2}$ are three different neural
networks composed independently by a stack of MLPs with linear residue
followed by a linear layer to match the dimension. That is,

$\displaystyle f_{g,e,\theta}$ $\displaystyle=P\circ L_{1}\circ
L_{2}\circ\cdots\circ L_{n}\circ W$ (23) $\displaystyle L_{i}(x)$
$\displaystyle=\phi[W_{i,2}\cdot\phi(W_{i,1}\cdot x+b_{i,1})+b_{i,2}]+x,$ (24)

where $\phi$ is a nonlinear activation function. We use $\phi(x)=\max(0,x)$,
the ReLU nonlinearity, through this work. The map $P$ is a padding layer that
matches the input dimension with the hidden layer dimension. The input of
$f_{g}$ is $0$ when there are no particles within the cutoff distance.
Accordingly, we set the bias $b_{i,j}$ parameters to be $0$ for $f_{g}$ to
ensure the control force is $0$ in this case. We document further
implementation details in the our code, which is available under an open-
source license All code for this project is available on github . The
resulting representation of the control force has no information on the total
number of particles, thus can be trained in a small system to quickly obtain
the optimal control forces and scale up to large systems thereafter.

##  IV Sampling dynamical fluctuations near phase transitions

###  IV.1 Current fluctuations in the asymmetric simple exclusion process

Figure 2: The 1D ASEP with lattice size $L=[10,50,100]$. The variational
results (dots) in different lattice size $L$ and DMRG results (dashed lines)
of (a) the SCGF scaled by $L$, (b) the average current per site. The
variational results of the average current shown in (b) are obtained by
numerically computing the derivative of the convex envelope of the
corresponding SCGF from (a). The convex envelope is obtained by performing the
Legendre-Fenchel twice. The insert of (c) shows the absolute error between the
variational results and the DMRG results. The parameters used throughout this
example are $p=0.1$, $q=0.9$, $\alpha=0.5$, $\beta=0.5$, $\gamma=0.5$,
$\delta=0.5$. For the parameters used in the neural network, the dimension of
the hidden layer is $20$, and the number of $L_{i}$ in (23) is $3$. The DMRG
results are adapted from Helms, Ray, and Chan (2019). Figure 3: (a) A typical
training process for $\lambda=-0.1$. (b)-(d) Typical trajectories of the
current in the biased ensemble. The yellow line depicts the average current
from the mean field theory. It is clearly that the control force changes the
current for different $\lambda$.

While the variational formulation of our approach gives us a natural metric to
compare the quality of the results, without an analytical solution (and in the
absence of competitive numerical techniques) we cannot guarantee that the
algorithm finds a true maximizer of (14). To test the accuracy of our
algorithm and whether it can be applied to broad classes of systems, we study
the current fluctuations in the asymmetric simple exclusion process (ASEP).
Exclusion processes are a paradigmatic model for nonequilibrium transport
phenomena and their complex dynamical fluctuations Derrida (1998); Chou,
Mallick, and Zia (2011). Although this model has been studied extensively and
many analytical results have been derived Lazarescu (2015), it is still
computationally challenging due to the high dimensionality of its state space
and its stochastic dynamics. Together, these properties make this model ideal
for testing new computational methods.

Recently, tensor network algorithms including the matrix product state (MPS)
and density-matrix renormalization group (DMRG) have been adapted to study
rare fluctuations in lattice models such as the ASEP and the kinetically
constrained models Bañuls and Garrahan (2019); Helms, Ray, and Chan (2019);
Helms and Chan (2020). Here we compute the current SCGF in the ASEP with open
boundaries and compare with previous numerical results obtained through DMRG
by Helms, Ray, and Chan (2019). Consider a 1-dimensional lattice with $L$
sites. Each site can be occupied by a hard-core particle, which randomly jumps
to its neighboring right or left site with continuous-time transition rates
$p$ and $q$, respectively. Hops to occupied sites are excluded from the
dynamics. A mass reservoir is coupled to the system at each boundary; the
effective chemical potential difference across the lattice is dictated by the
rates $\alpha,\beta,\gamma$, and $\delta$ (see the insert of Fig. 2 (a) for
the illustration). The dimension of the state space for a lattice with $L$
sites will be $2^{L}$, and the probability distribution evolves according to a
Kolmogorov forward equation $\partial_{t}P=WP$, where $W$ is the continuous-
time transition matrix.

The total current $J_{T}$ of the ASEP for time $T$ is $\int_{0}^{T}j_{t}dt$
where $j_{t}$ is the total rightward hops minus the total leftward hops per
unit time. The empirical total current rate, $j_{T}$, satisfies a large
deviation principle, which is the function we aim to compute. The SCGF,
similarly, can be computed through a variational formula:

$\psi(\lambda)=\sup_{u}\lim_{T\rightarrow\infty}\left\\{\lambda
j_{T}[u]-\frac{1}{T}\int_{0}^{T}W_{u}(x^{u}_{t^{-}},x^{u}_{t^{+}})\left[\frac{W(x^{u}_{t^{-}},x^{u}_{t^{+}})}{W_{u}(x^{u}_{t^{-}},x^{u}_{t^{+}})}-1-\log\frac{W(x^{u}_{t^{-}},x^{u}_{t^{+}})}{W_{u}(x^{u}_{t^{-}},x^{u}_{t^{+}})}\right]dt\right\\}$
(25)

where a different transition matrix $W_{u}$ is optimized. Eq. (25) is
different from (14) because the Radon-Nikodym derivative for Markov jump
processes has a different form than the one for SDEs Segall and Kailath
(1975). We review the derivation in Appendix A.

Denote $S_{i}=1$ when the $i$th site of the lattice is occupied by a particle
and $S_{i}=0$ otherwise. Practically, for the $i$th site, the modified
transition rate $[q_{u}(i),p_{u}(i)]=[q+\delta q_{u}(i),p+\delta p_{u}(i)]$
when $S_{i}=1$ and $[q_{u}(i),p_{u}(i)]=0$ when $S_{i}=0$. The rate
modification $[\delta q_{u}(i),\delta p_{u}(i)]$ is determined by a neural
network which input is $[S_{i-m},S_{i-m+1},\cdots,S_{i},\cdots,S_{i+m}]$. The
parameter $m$ plays the role of a cutoff distance, determining the range of
neighboring information the neural network can gather. In this example, force
equivariance does not arise because the system is one-dimensional, and the
input of $f_{e}(\cdot)$ becomes $x_{j}-x_{i}$ where $x_{j}$ is the location of
a particle within the cutoff distance. Throughout, we used $m=10$, meaning the
input can be both negative or positive depending on the location of the
neighboring particle.

To compute the SCGF we use automatic differentiation to backpropagate through
the dynamics. The results are compared with the asymptotic analytical results
in the limit of $L\rightarrow\infty$ derived from the Bethe ansatz Lazarescu
(2015), as well as the previous numerical results by the DMRG algorithm from
Helms, Ray, and Chan (2019), shown in Fig. 2-3. The DMRG algorithm with
sufficient bond-dimension is nearly exact for this model, which we verified by
comparing with the numerically exact spectral solution for $L=10$ (data not
shown). Fig. 2 shows that our method is very accurate for both the SCGF and
the cumulants obtained from this cumulant generating functions, and the
absolute errors are all less than the order of $10^{-2}$.

###  IV.2 Dynamical phase transitions in Active Brownian particle systems

Figure 4: (a) The SCGF of entropy production in the ABP system, with particle
number $N=[200,400,800]$. The data suggest that as $N\rightarrow\infty$ there
will be a singularity at $\lambda=0$, indicating a first-order dynamical phase
transition. The neural network used here is (22) with hidden layer dimension
equal to $50$ and number of $L_{i}$ in (23) is $3$. (b)-(d) Typical snapshots
of ABPs under different biasing field $\lambda$ ($N=400$). The red and grey
arrows depict the magnitude and direction of the control force and active
force $\boldsymbol{b}^{(i)}_{t}$ for each particle, respectively.

Recent advances on active matter, such as active Brownian particles (ABPs),
active nematics, and driven polymer networks, have articulated a relation
between macroscopic pattern formation and energy dissipation Lamtyugina  _et
al._ (2021); Bowick  _et al._ (2022); GrandPre  _et al._ (2021); Fodor,
Nemoto, and Vaikuntanathan (2020); Fodor  _et al._ (2019); Giomi  _et al._
(2013); Narayan, Ramaswamy, and Menon (2007); Nguyen, Qiu, and Vaikuntanathan
(2021). One challenge of characterizing macroscopic pattern formation that
results from rare collective fluctuations, is that it requires large system
sizes and high densities, making numerical sampling extremely difficult. Here
we examine our algorithm in the context of the active Brownian particle model,
in which the $i$th particle evolves as

$\displaystyle d\boldsymbol{X}_{t}^{(i)}$ $\displaystyle=[-\mu\frac{\partial
U(\boldsymbol{X}_{t})}{\partial\boldsymbol{x}^{(i)}}+v\boldsymbol{b}^{(i)}_{t}]dt+\sqrt{2D_{t}}d\boldsymbol{W}_{t}^{(i)},$
(26) $\displaystyle\boldsymbol{b}^{(i)}$
$\displaystyle=[\cos\phi_{t}^{(i)},\sin\phi_{t}^{(i)}]^{\top},\quad
d\phi_{t}^{(i)}=\sqrt{6D_{t}}d\boldsymbol{W}_{t}^{\phi^{(i)}}.$

where $U$ is a purely repulsive WCA potential, $\boldsymbol{x}^{(i)}$ denotes
the position of the $i$th particle, and $\boldsymbol{b}^{(i)}$ denotes the
direction of its active velocity. Throughout we consider systems with particle
density $\rho=N/L^{2}=0.6$, and $v=100$ which is within clustered phase as a
result of motility-induced phase separation (MIPS) Hagan, Baskaran, and Redner
(2013). The non-conservative self-propulsion term $v\boldsymbol{b}^{(i)}$
represents the dissipative “active” force. In the ABP model,
$\boldsymbol{b}_{t}^{(i)}$ are unit vectors which rotate diffusively and $v$
is the magnitude of the active force. Here, $\boldsymbol{W}_{t}^{(i)}$ and
$\boldsymbol{W}_{t}^{\phi^{(i)}}$ are independent standard Wiener processes.
The phenomenology of motility induced phase separation is relatively well-
understood compared with other nonequilibrium phase transitions: for
sufficiently large Péclet number and density, a macroscopic aggregate of
particles formsHagan, Baskaran, and Redner (2013). GrandPre  _et al._ (2021)
demonstrated that this transition correlates with the average entropy
production

$s=\frac{1}{NT}\sum_{i=1}^{N}\int_{0}^{T}v\boldsymbol{b}^{(i)}_{t}D^{-1}_{t}\circ
d\boldsymbol{X}_{t}^{(i)}.$ (27)

Mechanistically, when the system enters the phase separated state, much of the
directional motion ceases, leading to a drop in the average entropy production
compared to an unclustered trajectory.

Previously, we have shown that by using a neural network ansatz, it is
possible to compute the many-body control force and the SCGF for systems with
up to 200 particles. Here, by using the formalism proposed in Section III, we
can employ _transfer learning_. That is, we train the neural network in a
small system (e.g., $N=200$), because the GNN representation is naturally
independent of the total size of the system, we can use the trained as an
initial condition for larger systems. In Appendix B, we show that the GNN
formulation of the present work outperforms an “uninformed” neural network.

We show the results for the entropy production rate function in Fig. 4. In
Fig. 4(a) we see clear convergence as $N$ increases, and a singularity at
around $\lambda=-1$ indicates a first-order dynamical phase transition,
consistent with previous studiesGrandPre  _et al._ (2021); Yan, Touchette, and
Rotskoff (2022). For positive $\lambda$, the control force induces a
collective flow to increase the average entropy production rate as shown in
Fig. 4 (d). When $\lambda$ is negative, below the transition point, particles
cluster, reducing mobility and subsequently the entropy production.

One conceptual advantage of the optimal control formalism is that we can
directly access the control forces that realize the target rare event. To
better understand the physical forces that drive low entropy production
clustered dynamics, we plot the learned control force of one single particle
as a function of the distance between it and the center of a cluster (Fig. 5).
The results show that for particles deeply inside the cluster,r the control
force is nearly zero, as it is stabilized by the active forces. The force
gradually increases near the boundary of the cluster surface and reaches a
maximum when a particle is just outside the cluster. The force vanishes at the
cutoff distance $d_{\rm c}=4\sigma$ by construction, but its relatively long
range emphasizes the effective attraction that drives clustering.

Figure 5: The learned control force of one single particle as it gradually
approaching an established cluster ($\lambda=-10$, $\rho=0.1$, $v=100$). We
plot the magnitude of the control force in two orthogonal directions,
perpendicular to and along the vector $\vec{w}$, which connects the particle
and the cluster’s center of mass. The force is in the unit of the velocity
$v$, and the distance is in the unit of the particle size $\sigma$. The yellow
dashed line represents the position of the interface of the cluster. As
expected, the control force is small when a particle is inside the cluster,
and reaches its maximum when closing to the cluster.

##  V Conclusion and future work

Many of neural networks used in practice, such as convolutional neural
networks and transformers Vaswani  _et al._ (2017) have become popular because
of their success on “big data” tasks, computer vision and natural language
processing. Problems in scientific computing demand a distinct class of
function representations. The physics-informed graph neural networks we have
developed in this work not only provide a representation that is specially
tailored to the symmetries and properties of the systems we consider, they
also are naturally invariant to the extent of the system. The ansatz that we
use builds in translational invariance and rotational equivariance. Such a
constraint limits the function class, likely accelerating learning, as well.
Together, these properties manifest in scalability: we can optimize the
control forces on small systems where data collection is inexpensive
computationally, and then deploy and refine them on systems large enough to
explore collective phenomena.

Recently, tensor network algorithms including MPS and DMRG have shown promise
in dynamical large deviation calculations, yet these approaches are much more
tailored to lattice models with MJP dynamics. The ansatz used for DMRG uses
much a longer effective range of correlations when the bond-dimension is high.
The results we obtain for the ASEP model show our approach achieves good
performance compared with DMRG, providing an alternative when DMRG is ill-
suited to the system in question.

In this work, we consider only homogeneous systems in which there is only one
type of particle. Extending the current framework to heterogeneous systems,
for example, multi-component active materials or molecular glass formers,
remains a significant future challenge. Furthermore, since our physics-
informed GNNs provide an expressive representation of potentially any physical
force, we could apply these networks to many other problems, including
nonequilibrium response calculations.

###### Acknowledgements.

The authors thank Todd Gingrich for several helpful discussions. We also thank
Phillip Helms for generously sharing the DMRG data used in Fig. 2. This work
was partially supported by the Terman faculty fellowship.

## References

  * Toner and Tu (1995) J. Toner and Y. Tu, Phys. Rev. Lett. 75, 4326 (1995).
  * Toner and Tu (1998) J. Toner and Y. Tu, Phys. Rev. E 58, 4828 (1998).
  * Lazarescu (2015) A. Lazarescu,  _Exact Large Deviations of the Current in the Asymmetric Simple Exclusion Process with Open Boundaries_ , PhD Thesis, Institut de Physique Théorique, CEA-Saclay (2015).
  * Martin  _et al._ (2021) D. Martin, J. O’Byrne, M. E. Cates, E. Fodor, C. Nardini, J. Tailleur, and F. van Wijland, Phys. Rev. E 103, 032607 (2021).
  * Fodor, Jack, and Cates (2022) É. Fodor, R. L. Jack, and M. E. Cates, Ann. Rev. Cond. Matter Phys. 13 (2022).
  * Cates and Tailleur (2015) M. E. Cates and J. Tailleur, Ann. Rev. Cond. Matter Phys. 6, 219 (2015).
  * Garrahan  _et al._ (2007) J. P. Garrahan, R. L. Jack, V. Lecomte, E. Pitard, K. van Duijvendijk, and F. van Wijland, Phys. Rev. Lett. 98, 195702 (2007).
  * Touchette (2009) H. Touchette, Phys. Rep. 478, 1 (2009).
  * Ruelle (2009) D. Ruelle,  _Thermodynamic Formalism_, 2nd ed., The Mathematical Structures of Equilibrium Statistical Mechanics (Cambridge University Press, Cambridge, 2009).
  * Garrahan  _et al._ (2009) J. P. Garrahan, R. L. Jack, V. Lecomte, E. Pitard, K. van Duijvendijk, and F. van Wijland, J. Phys. A 42, 075007 (2009).
  * Hedges  _et al._ (2009) L. O. Hedges, R. L. Jack, J. P. Garrahan, and D. Chandler, Science 323, 1309 (2009).
  * Gingrich, Vaikuntanathan, and Geissler (2014) T. R. Gingrich, S. Vaikuntanathan, and P. L. Geissler, Phys. Rev. E 90, 042123 (2014).
  * GrandPre  _et al._ (2021) T. GrandPre, K. Klymko, K. K. Mandadapu, and D. T. Limmer, Phys. Rev. E 103, 012613 (2021).
  * Barato and Seifert (2015) A. C. Barato and U. Seifert, Phys. Rev. Lett. 114, 158101 (2015).
  * Gingrich  _et al._ (2016) T. R. Gingrich, J. M. Horowitz, N. Perunov, and J. L. England, Phys. Rev. Lett. 116, 120601 (2016).
  * Yan, Touchette, and Rotskoff (2022) J. Yan, H. Touchette, and G. M. Rotskoff, Phys. Rev. E 105, 024115 (2022).
  * Das  _et al._ (2021) A. Das, D. C. Rose, J. P. Garrahan, and D. T. Limmer, arXiv:2105.04321 (2021), arXiv:2105.04321 .
  * Casert  _et al._ (2021) C. Casert, T. Vieijra, S. Whitelam, and I. Tamblyn, Phys. Rev. Lett. 127, 120602 (2021).
  * Bucklew (2004) J. A. Bucklew,  _Introduction to Rare Event Simulation_ (Springer, New York, 2004).
  * Ray, Chan, and Limmer (2017) U. Ray, G. K.-L. Chan, and D. T. Limmer, Phys. Rev. Lett. 120, 210602 (2017).
  * Ferré and Touchette (2018) G. Ferré and H. Touchette, J Stat Phys 172, 1525 (2018).
  * Grassberger (2002) P. Grassberger, Comp. Phys. Comm. 147, 64 (2002).
  * Giardinà, Kurchan, and Peliti (2006) C. Giardinà, J. Kurchan, and L. Peliti, Phys. Rev. Lett. 96, 120603 (2006).
  * Lecomte and Tailleur (2007) V. Lecomte and J. Tailleur, J. Stat. Mech. 2007, P03004 (2007).
  * Ray and Chan (2020) U. Ray and G. K.-L. Chan, J. Chem. Phys. 152, 104107 (2020).
  * Nemoto  _et al._ (2016) T. Nemoto, F. Bouchet, R. L. Jack, and V. Lecomte, Phys. Rev. E 93, 062123 (2016).
  * Jacobson and Whitelam (2019) D. Jacobson and S. Whitelam, Phys. Rev. E 100, 052139 (2019).
  * Whitelam, Jacobson, and Tamblyn (2020) S. Whitelam, D. Jacobson, and I. Tamblyn, J. Chem. Phys. 153, 044113 (2020).
  * Whitelam and Tamblyn (2020) S. Whitelam and I. Tamblyn, Phys. Rev. E 101, 052604 (2020).
  * Sprague  _et al._ (2021) K. Sprague, J. Carrasquilla, S. Whitelam, and I. Tamblyn, Mach. Learn.: Sci. Technol. 2, 02LT02 (2021).
  * Lebowitz and Spohn (1999) J. L. Lebowitz and H. Spohn, J. Stat. Phys. 95, 333 (1999).
  * Helms, Ray, and Chan (2019) P. Helms, U. Ray, and G. K.-L. Chan, Phys. Rev. E 100, 022101 (2019).
  * Helms and Chan (2020) P. Helms and G. K.-L. Chan, Phys. Rev. Lett. 125, 140601 (2020).
  * Bañuls and Garrahan (2019) M. C. Bañuls and J. P. Garrahan, Phys. Rev. Lett. 123, 200601 (2019).
  * Das and Limmer (2019) A. Das and D. T. Limmer, J. Chem. Phys. 151, 244123 (2019).
  * Rose, Mair, and Garrahan (2021) D. C. Rose, J. F. Mair, and J. P. Garrahan, New J. Phys. 23, 013013 (2021).
  * Das, Kuznets-Speck, and Limmer (2022) A. Das, B. Kuznets-Speck, and D. T. Limmer, Phys. Rev. Lett. 128, 028005 (2022).
  * Rotskoff and Vanden-Eijnden (2019) G. M. Rotskoff and E. Vanden-Eijnden, Physical Review Letters 122, 150602 (2019).
  * Thin  _et al._ (2021) A. Thin, Y. Janati, S. L. Corff, C. Ollion, A. Doucet, A. Durmus, E. Moulines, and C. Robert, arXiv:2103.10943 [stat] (2021), arXiv:2103.10943 [stat] .
  * Zhang  _et al._ (2018) L. Zhang, J. Han, H. Wang, R. Car, and W. E, Phys. Rev. Lett. 120, 143001 (2018).
  * Schütt  _et al._ (2018) K. T. Schütt, H. E. Sauceda, P.-J. Kindermans, A. Tkatchenko, and K.-R. Müller, J. Chem. Phys. 148, 241722 (2018).
  * Wang  _et al._ (2018) H. Wang, L. Zhang, J. Han, and W. E, Comput. Phys. Commun. 228, 178 (2018).
  * Frishman and Ronceray (2020) A. Frishman and P. Ronceray, Phys. Rev. X 10, 021009 (2020).
  * Satorras, Hoogeboom, and Welling (2021) V. G. Satorras, E. Hoogeboom, and M. Welling, in  _Proceedings of the 38th International Conference on Machine Learning_, Proceedings of Machine Learning Research, Vol. 139, edited by M. Meila and T. Zhang (PMLR, 2021) pp. 9323–9332.
  * Maes (2020) C. Maes, Frontiers in Physics 8, 229 (2020).
  * Touchette (2018) H. Touchette, Physica A 504, 5 (2018).
  * Dupuis and Wang (2004) P. Dupuis and H. Wang, Stoch. Stoch. Rep. 76, 481 (2004).
  * Oksendal (1992) B. Oksendal,  _Stochastic Differential Equations: An Introduction with Applications_ , Universitext (Springer, Berlin, 1992).
  * Karniadakis  _et al._ (2021) G. E. Karniadakis, I. G. Kevrekidis, L. Lu, P. Perdikaris, S. Wang, and L. Yang, Nat. Rev. Phys. 3, 422 (2021).
  * (50) All code for this project is available on github, https://github.com/quark-strange/GNNLDP .
  * Derrida (1998) B. Derrida, Phys. Rep. 301, 65 (1998).
  * Chou, Mallick, and Zia (2011) T. Chou, K. Mallick, and R. K. Zia, Rep. Prog. Phys 74, 116601 (2011).
  * Segall and Kailath (1975) A. Segall and T. Kailath, Ann. Probab. , 449 (1975).
  * Lamtyugina  _et al._ (2021) A. Lamtyugina, Y. Qiu, É. Fodor, A. R. Dinner, and S. Vaikuntanathan, “Thermodynamic control of activity patterns in cytoskeletal networks,” (2021), arXiv:2111.08677 .
  * Bowick  _et al._ (2022) M. J. Bowick, N. Fakhri, M. C. Marchetti, and S. Ramaswamy, Phys. Rev. X 12, 010501 (2022).
  * Fodor, Nemoto, and Vaikuntanathan (2020) É. Fodor, T. Nemoto, and S. Vaikuntanathan, New J. Phys. 22, 013052 (2020).
  * Fodor  _et al._ (2019) É. Fodor, T. Nemoto, S. Vaikuntanathan, and L. Tociu, Phys. Rev. X 9, 041026 (2019).
  * Giomi  _et al._ (2013) L. Giomi, M. J. Bowick, X. Ma, and M. C. Marchetti, Phys. Rev. Lett. 110, 228101 (2013).
  * Narayan, Ramaswamy, and Menon (2007) V. Narayan, S. Ramaswamy, and N. Menon, Science 317, 105 (2007).
  * Nguyen, Qiu, and Vaikuntanathan (2021) M. Nguyen, Y. Qiu, and S. Vaikuntanathan, Annu. Rev. Condens. Matter Phys. 12, 273 (2021).
  * Hagan, Baskaran, and Redner (2013) M. F. Hagan, A. Baskaran, and G. S. Redner, Phys. Rev. Lett. 110, 055701 (2013).
  * Vaswani  _et al._ (2017) A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, in  _Advances in Neural Information Processing Systems_ , Vol. 30, edited by I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (Curran Associates, Inc., 2017).
  * Chetrite and Touchette (2015) R. Chetrite and H. Touchette, J. Stat. Mech. 2015, P12001 (2015).

##  Appendix A Variational formula of SCGF for Markov jump processes

The Radon-Nikodym derivative (16) is a result of the Girsanov theorem Oksendal
(1992). For Markov jump processes, the change of path measure is obtained via
a similar procedure, cf. Segall and Kailath (1975). For simplicity, here we
also review an alternative proof by using the contraction principle of the
large deviation theory Chetrite and Touchette (2015).

Consider a general Markov jump process on the state space $\mathcal{X}$, and
$x,y\in\mathcal{X}$. Let $W(x,y)$ be the continuous-time transition rate from
$x$ to $y$. Level 2.5 large deviation theory Chetrite and Touchette (2015)
(cf. Appendix B, therein) states that the joint rate function of the empirical
density $\rho(x)$, i.e., the proportion of time that a trajectory stays at
$x$, and the empirical flow $q(x,y)$, i.e., time average of the number of
jumps from $x$ to $y$, can be explicitly written as

$\displaystyle K(\rho,q)$
$\displaystyle=\sum_{x,y}\left[q(x,y)\log\frac{q(x,y)}{\rho(x)W(x,y)}-q(x,y)+\rho(x)W(x,y)\right].$
(28)

From the contraction principle, (28) can be obtained by contracting over all
control forces $u$, implemented through the transition rate matrix $W_{u}$,

$K(\rho,q)=\min_{u}K(\rho_{u},q_{u})=\sum_{x,y}q_{u}(x,y)\left[\frac{\rho_{u}(x)W(x,y)}{q_{u}(x,y)}-1-\log\frac{\rho_{u}(x)W(x,y)}{q_{u}(x,y)}\right],$
(29)

where $\rho_{u}$ and $q_{u}(x,y)=\rho_{u}(x)W_{u}(x,y)$ and stationary density
and flow associated with the new transition rate matrix $W_{u}$. Simplifying
the expression above, we have

$K(\rho_{u},q_{u})=\sum_{x,y}\rho_{u}(x)W_{u}(x,y)\left[\frac{W(x,y)}{W_{u}(x,y)}-1-\log\frac{W(x,y)}{W_{u}(x,y)}\right].$
(30)

Because the dynamics is ergodic, the ensemble average (30) is equivalent to
averaging over trajectories $\\{x_{t}^{u}\\}_{0\leqslant t\leqslant T}$
determined by the rate matrix $W_{u}$; as a result, we can write

$K(\rho_{u},q_{u})=\lim_{T\rightarrow\infty}\frac{1}{T}\int_{0}^{T}W_{u}(x^{u}_{t^{-}},x^{u}_{t^{+}})\left[\frac{W(x^{u}_{t^{-}},x^{u}_{t^{+}})}{W_{u}(x^{u}_{t^{-}},x^{u}_{t^{+}})}-1-\log\frac{W(x^{u}_{t^{-}},x^{u}_{t^{+}})}{W_{u}(x^{u}_{t^{-}},x^{u}_{t^{+}})}\right]dt.$
(31)

The SCGF for the current can be estimated by optimizing the transition rate
matrix $W_{u}$ to perform the contraction in (29),

$\displaystyle\psi(\lambda)=\sup_{u}\lim_{T\rightarrow\infty}\left\\{\lambda
j_{T}[u]-\frac{1}{T}\int_{0}^{T}W_{u}(x^{u}_{t^{-}},x^{u}_{t^{+}})\left[\frac{W(x^{u}_{t^{-}},x^{u}_{t^{+}})}{W_{u}(x^{u}_{t^{-}},x^{u}_{t^{+}})}-1-\log\frac{W(x^{u}_{t^{-}},x^{u}_{t^{+}})}{W_{u}(x^{u}_{t^{-}},x^{u}_{t^{+}})}\right]dt\right\\}.$
(32)

##  Appendix B Comparison with a general purpose Neural Networks

Figure 6: Comparison of the performances of the GNN approach and the NN
approach using the ABPs model. We test two set of parameters: (a)-(b)
$\rho=0.1$, $v=10$ and (c)-(d) $\rho=0.6$, $v=100$ and plot the SCGF in (a),
(c); and the average entropy production per particle in (b), (d). For the GNN
approach, we use two $L_{i}$ (see (23)) with hidden layer width equal to
$100$. For the NN approach, we use two $L_{i}$ with hidden layer width equal
to $2000$.

In previous work Yan, Touchette, and Rotskoff (2022) we used a single neural
network to generate a many-body control force for all particles. To compare
the performance of this general purpose neural network (NN) approach with the
graph neural network (GNN) approach presented in this work, we here test two
methods in the ABPs model. We tested two set of parameters (number density
$\rho$ and active force velocity $v$) which lays in different phases in the
ABPs model Hagan, Baskaran, and Redner (2013). The results are shown in Fig.
6. Due to the variational principle we derive, the larger value of
$\psi(\lambda)$ must be closer, if not equal, to the true results. In the ABP
model, it seems that both method perform equally well for small biasing
$\lambda$, but the GNN approach proposed in this work outperforms the previous
approach when $\lambda$ becomes more negative. Furthermore, when $v$ is large,
the previous approach even fails to capture the nature of the dynamical phase
transition, characterized by a singularity of $\psi(\lambda)$ (Fig. 6(c)-(d)),
while the GNN approach still remains robust.

##  Appendix C Numerical details of the ABPs model

In the active Brownian particle model (26), the conservative inter-particle
force is modeled as a purely repulsive WCA pair potential $U=\sum_{i\neq
j}u(l_{ij})$ which depends on the relative distance between any two particle
$i$ and $j$, denoted $l_{ij}$. Explicitly,

$u(l_{ij})=\left\\{\begin{array}[]{lcl}4\epsilon\left[(\sigma/l_{ij})^{12}-(\sigma/l_{ij})^{6}\right]+\epsilon,&l_{ij}\leqslant
2^{1/6}\sigma\\\ 0,&l_{ij}>2^{1/6}\sigma,\end{array}\right.$ (33)

In the simulation, the unit of length is normalized by $\sigma$ and we set
$\epsilon=1$. We used periodic boundary condition with box length $L$. The
particle number density $\rho$ in Fig. 4 is $0.6$ and $v=100$. In Fig. 5 we
use $\rho=0.1$ and $v=100$. When $\rho$ is high enough, particles aggregate
into a band rather than a circular cluster.