# The $\alpha$-divergence Improves the Entropy Production Estimation via
Machine Learning

Euijoon Kwon  Department of Physics and Astronomy & Center for Theoretical
Physics, Seoul National University, Seoul 08826, Republic of Korea  Yongjoo
Baek<EMAIL_ADDRESS>Department of Physics and Astronomy & Center for
Theoretical Physics, Seoul National University, Seoul 08826, Republic of Korea

###### Abstract

Recent years have seen a surge of interest in the algorithmic estimation of
stochastic entropy production (EP) from the trajectory data via machine
learning. A crucial element of such algorithms is the identification of a loss
function whose minimization guarantees the accurate EP estimation. In this
study, we show that there exists a host of loss functions, namely those
implementing a variational representation of the $\alpha$-divergence, which
can be used for the EP estimation. Among these loss functions, the one
corresponding to $\alpha=-0.5$ exhibits the most robust performance against
strong nonequilibrium driving or slow dynamics, which adversely affects the
existing method based on the Kullback-Leibler divergence ($\alpha=0$). To
corroborate our findings, we present an exactly solvable simplification of the
EP estimation problem, whose loss function landscape and stochastic properties
demonstrate the optimality of $\alpha=-0.5$.

Figure 1:  (a) The architecture of the $\alpha$-NEEP. (b) An illustration of
the two-bead model. (c) The mean square error (MSE) of the EP estimate for the
two-bead model for various temperature differences and (d) that of the driven
Brownian particle as the potential depth $A$ is varied. The error bars
indicate standard deviations obtained from $40$ independent trainings. All the
other parameters ($T_{c}$, $\gamma$, and $k$) are fixed to unity throughout
this Letter. Also see Appendix A for more detail about the training procedure.
The dashed lines are to guide the eye.

Introduction. — How irreversible does a process look? One may pose this
question for two distinct reasons. First, whether a biological process
requires energy dissipation is often a subject of much debate Brangwynne  _et
al._ (2008); Weber  _et al._ (2012). To resolve this issue, it is useful to
note that irreversibility suggests energy dissipation. Various hallmarks of
irreversibility, such as the breaking of the fluctuation-dissipation theorem
Mizuno  _et al._ (2007) and the presence of nonequilibrium probability
currents in the phase space Battle  _et al._ (2016); Gladrow  _et al._ (2016),
have been used to determine whether energy is dissipated. Second, whether a
nonequilibrium system allows for an effective equilibrium description is an
important issue. For instance, in active matter, despite the energy
dissipation at the microscopic level, it has been argued that the large-scale
phenomena allows for an effective equilibrium description Tailleur and Cates
(2008); Speck  _et al._ (2014); Takatori  _et al._ (2014); Farage  _et al._
(2015); Solon  _et al._ (2018). If we can quantify the irreversibility of an
empirical process at various levels of coarse-graining Fodor  _et al._ (2016);
Nardini  _et al._ (2017), it will provide us with helpful clues as to whether
we should look for an effective equilibrium theory for the process.

Based on the framework of stochastic thermodynamics, modern thermodynamics
assigns entropy production (EP) to each stochastic trajectory based on its
irreversibility Seifert (2012). Thus, empirically measuring the
irreversibility of a process is closely tied to the problem of estimating EP
from sampled trajectories Roldán and Parrondo (2010, 2012); Li  _et al._
(2019); Van Vu  _et al._ (2020); Otsubo  _et al._ (2020, 2022); Lee  _et al._
; Kim  _et al._ (2020). A straightforward approach to the problem is to
evaluate the relevant transition probabilities by directly counting the number
of trajectory segments, which is called the plug-in method Roldán and Parrondo
(2010, 2012). The method, readily applicable to discrete systems, can also be
applied to continuous systems through the use of kernel functions Li  _et al._
(2019). However, while this method is simple and intuitive, it requires a huge
ensemble of lengthy trajectories for accurate estimations (curse of
dimensionality). More recent studies proposed methods based on universal lower
bounds of the average EP, such as the thermodynamic uncertainty relations Li
_et al._ (2019); Van Vu  _et al._ (2020); Otsubo  _et al._ (2020, 2022) and
the entropic bound Lee  _et al._ . While these methods do not suffer from the
curse of dimensionality and are applicable even to non-stationary processes
Otsubo  _et al._ (2022); Lee  _et al._ , their accuracy is impaired when the
underlying bounds are not tight. Moreover, these methods are applicable only
to the estimation of the average EP, not the EP of each trajectory.

Meanwhile, with the advent of machine learning techniques in physics, a novel
method for EP estimation using artificial neural networks has been developed
Kim  _et al._ (2020). This method, called the Neural Estimator for Entropy
Production (NEEP), minimizes the loss function based on a variational
representation of the Kullback-Leibler (KL) divergence. Without any
presupposed discretization of the phase space and using the rich expressivity
of neural networks, the NEEP suffers far less from the complications of the
sampling issues and is applicable to a diverse range of stochastic processes
Otsubo  _et al._ (2022).

Still, the NEEP has its limits. Its accuracy deteriorates when the
nonequilibrium driving is strong or when the dynamics slows down so that the
phase space is poorly sampled. In this Letter, we show that the NEEP can be
significantly improved by changing the loss function. Toward this purpose, we
propose the $\alpha$-NEEP, which generalizes the NEEP. Instead of the KL
divergence, the $\alpha$-NEEP utilizes the $\alpha$-divergence, which has been
mainly used in the machine learning community Basu  _et al._ (1998); Sugiyama
_et al._ (2012); Nowozin  _et al._ (2016); Belghazi  _et al._ (2018). We
demonstrate that the $\alpha$-NEEP with $\alpha=-0.5$ shows much more robust
performance for a broader range of nonequilibrium driving and sampling
quality. This is corroborated by an analytically tractable simplification of
the $\alpha$-NEEP that shows the optimality of $\alpha=-0.5$.

Mathematical framework of the $\alpha$-NEEP. — Suppose our goal is to estimate
the EP of a Markov process in discrete time, $\mathbf{x}_{t}$, in a
$d$-dimensional space. Every pair of states, denoted by
$\mathbf{x}\equiv(\mathbf{x}_{t},\mathbf{x}_{t+1})$, is mapped by the neural
network to a real number $r(\mathbf{x})$. The goal of the $\alpha$-NEEP is to
find $r(\mathbf{x})$ that minimizes the loss function

$\displaystyle\mathcal{L}_{\alpha}[r]\equiv\left\langle-\frac{r(\mathbf{x})^{\alpha}-1}{\alpha}+\frac{q(\mathbf{x})}{p(\mathbf{x})}\frac{r(\mathbf{x})^{1+\alpha}-1}{1+\alpha}\right\rangle_{p(\mathbf{x})},$
(1)

where $p(\mathbf{x})$ and $q(\mathbf{x})$ are probability density functions,
$\langle\cdot\rangle_{p(\mathbf{x})}$ denotes the average with respect to the
distribution $p(\mathbf{x})$, and $\alpha$ is a real number other than $0$ and
$-1$. See Appendix B for discussions of these two exceptional cases.

It can be rigorously shown (see Appendix B) that $\mathcal{L}_{\alpha}[r]$
satisfies the inequality

$\displaystyle\mathcal{L}_{\alpha}[r]\geq-D_{\alpha}[p\\!:\\!q]$ (2)

for the $\alpha$-divergence, a measure of difference between the two
distributions $p(\mathbf{x})$ and $q(\mathbf{x})$, defined as

$\displaystyle
D_{\alpha}[p\\!:\\!q]\equiv\left\langle\frac{[p(\mathbf{x})/q(\mathbf{x})]^{\alpha}-1}{\alpha(1+\alpha)}\right\rangle_{p(\mathbf{x})}.$
(3)

Here, $\mathcal{L}_{\alpha}[r]=-D_{\alpha}[p\\!:\\!q]$ is achieved if and only
if $r(\mathbf{x})=p(\mathbf{x})/q(\mathbf{x})$ for all $\mathbf{x}$. In other
words, by minimizing $\mathcal{L}_{\alpha}[r]$ to find
$D_{\alpha}[p\\!:\\!q]$, we also obtain an estimate for the ratio
$p(\mathbf{x})/q(\mathbf{x})$. We note that the properties of
$\mathcal{L}_{\alpha}[r]$ used here are also valid for a much more general
class of loss functions, as discussed in Basu  _et al._ (1998); Sugiyama  _et
al._ (2012) (also in Appendix B).

Now we show how to estimate the EP using the above properties. First, we
identify $p(\mathbf{x})$ and $q(\mathbf{x})$ as the forward and reverse path
probabilities $p(\mathbf{x}_{t},\mathbf{x}_{t+1})$ and
$p(\mathbf{x}_{t+1},\mathbf{x}_{t})$, respectively. The stochastic EP of the
transition $\mathbf{x}_{t}\rightarrow\mathbf{x}_{t+1}$ is given by the log
ratio

$\displaystyle\Delta
S(\mathbf{x}_{t},\mathbf{x}_{t+1})=\log\frac{p(\mathbf{x}_{t},\mathbf{x}_{t+1})}{p(\mathbf{x}_{t+1},\mathbf{x}_{t})}\;.$
(4)

Keeping this in mind, we replace $r(\mathbf{x})$ with
$\mathrm{e}^{s_{\theta}(\mathbf{x}_{t},\mathbf{x}_{t+1})}$ in Eq. (1), where
$s_{\theta}$ plays the role of the EP estimator with $\theta$ representing the
neural network state. Then, the loss function can be rewritten as

$\displaystyle\mathcal{L}_{\alpha}(\theta)=\left\langle-\frac{e^{\alpha
s_{\theta}(\mathbf{x}_{t},\mathbf{x}_{t+1})}-1}{\alpha}+\frac{e^{-(1+\alpha)s_{\theta}(\mathbf{x}_{t},\mathbf{x}_{t+1})}-1}{1+\alpha}\right\rangle\;.$
(5)

Here $\langle\cdot\rangle$ denotes the average over all empirical trajectories
distributed by $p(\mathbf{x}_{t},\mathbf{x}_{t+1})$. This loss function
satisfies $\mathcal{L}_{-(1+\alpha)}(\theta)=\mathcal{L}_{\alpha}(\theta)$, so
the $\alpha$-NEEP is symmetric under the exchange
$\alpha\leftrightarrow-(1+\alpha)$. For this reason, through the rest of this
Letter, we focus on the regime $-0.5\leq\alpha\leq 0$ (the regime $\alpha>0$
leads to very poor performance and is left out). Noting that Eq. (4) implies
the antisymmetry $\Delta S(\mathbf{x}_{t},\mathbf{x}_{t+1})=-\Delta
S(\mathbf{x}_{t+1},\mathbf{x}_{t})$, we may set the estimator $s_{\theta}$ to
be related to the feedforward neural network (FNN) output $h_{\theta}$ as

$\displaystyle
s_{\theta}(\mathbf{x}_{t},\mathbf{x}_{t+1})=h_{\theta}(\mathbf{x}_{t},\mathbf{x}_{t+1})-h_{\theta}(\mathbf{x}_{t+1},\mathbf{x}_{t})\;,$
(6)

so that the neural network focuses on the estimators that satisfy the
antisymmetry of EP for more efficient training. The method described so far is
schematically illustrated in Fig. 1(a).

We note that the loss function used in the original NEEP Kim  _et al._ (2020)
is recovered in the limit $\alpha\rightarrow 0$:

$\displaystyle\mathcal{L}_{0}(\theta)$
$\displaystyle=\left\langle-1-s_{\theta}(\mathbf{x}_{t},\mathbf{x}_{t+1})+e^{-s_{\theta}(\mathbf{x}_{t},\mathbf{x}_{t+1})}\right\rangle$
$\displaystyle\geq-
D_{\text{KL}}[p(\mathbf{x}_{t},\mathbf{x}_{t+1})\\!:\\!p(\mathbf{x}_{t+1},\mathbf{x}_{t})]\;.$
(7)

Figure 2:  Ratio between the estimated value $\sigma_{\mathrm{pred}}$ and the
true value $\sigma$ of the average EP for (a) the two-bead model and (b) the
driven Brownian particle. Each data point and error bar are obtained from $40$
independent trainings.

Performance of the $\alpha$-NEEP. — To assess the performance of the
$\alpha$-NEEP, we apply the method to toy models of nonequilibrium systems,
namely the two-bead model and the driven Brownian particle.

(i) The two-bead model. This model has been used in a number of previous
studies as a benchmark for testing EP estimators Battle  _et al._ (2016); Li
_et al._ (2019); Otsubo  _et al._ (2020); Kim  _et al._ (2020). The model
consists of two 1-d overdamped beads which are connected to each other and to
the walls on both sides by identical springs, see Fig. 1(b). The beads are in
contact with heat baths at temperatures $T_{h}$ and $T_{c}$ with
$T_{h}>T_{c}$. Denoting by $x_{h}$ ($x_{c}$) the bead in contact with the hot
(cold) bath, the stochastic equations of motion are given by

$\displaystyle\dot{x}_{h}$
$\displaystyle=\frac{k}{\gamma}(-2x_{h}+x_{c})+\sqrt{\frac{2T_{h}}{\gamma}}\xi_{h}(t)\;,$
(8a) $\displaystyle\dot{x}_{c}$
$\displaystyle=\frac{k}{\gamma}(-2x_{c}+x_{h})+\sqrt{\frac{2T_{c}}{\gamma}}\xi_{c}(t)\;.$
(8b)

Here $k$ is the spring constant, $\gamma$ the friction coefficient, and
$\xi_{h,\,c}$ the Gaussian thermal noise with zero means and
$\langle\xi_{h}(t)\xi_{h}(t^{\prime})\rangle=\langle\xi_{c}(t)\xi_{c}(t^{\prime})\rangle=\delta(t-t^{\prime})$.
For infinitesimal displacements $(dx_{h},dx_{c})$, the associated stochastic
EP is given by

$\displaystyle\Delta S=\frac{k}{\gamma}\left[\frac{2x_{h}-x_{c}}{T_{h}}\circ
dx_{h}+\frac{2x_{c}-x_{h}}{T_{c}}\circ dx_{c}\right]+\Delta S_{\text{sys}}\;,$
(9)

where $\circ$ denotes the Stratonovich product and $\Delta S_{\text{sys}}$ the
change of the system’s Shannon entropy, namely

$\displaystyle\Delta
S_{\text{sys}}=-\ln\frac{p_{\mathrm{s}}(x_{h}+dx_{h},x_{c}+dx_{c})}{p_{\mathrm{s}}(x_{h},x_{c})}$
(10)

for the steady-state distribution $p_{\mathrm{s}}(x_{h},x_{c})$. Since the
system is fully linear, $p_{s}(x_{h},x_{c})$ can be calculated analytically.
Thus the EP of this model can be calculated exactly using Eq. (9) and compared
with the $\alpha$-NEEP result.

(ii) The driven Brownian particle. As an example of a nonlinear system, we
also consider a 1-d overdamped Brownian particle in a periodic potential
$U(x)=A\sin x$ driven by a constant force $f$. The motion of the particle is
described by the Langevin equation

$\displaystyle\dot{x}=f-U^{\prime}(x)+\sqrt{2T}\xi(t)\;,$ (11)

where $\xi(t)$ is a Gaussian white noise with unit variance. For this model,
the stochastic EP associated with the infinitesimal displacement $dx$ is given
by

$\displaystyle\Delta S=\frac{-f\,dx+U(x+dx)-U(x)}{T}+\Delta S_{\text{sys}}\;,$
(12)

where $\Delta S_{\text{sys}}=-\ln p_{\mathrm{s}}(x+dx)/p_{\mathrm{s}}(x)$
again denotes the Shannon entropy change for the steady-state distribution
$p_{\mathrm{s}}(x)$. Since the system is 1-d, it is straightforward to obtain
$p_{\mathrm{s}}(x)$ by numerical integration. Thus, the stochastic EP of this
model can also be calculated exactly and compared to the $\alpha$-NEEP result.

To see how the predicted EP differs from the true EP, we observe the behavior
of the mean square error (MSE) $\langle(s_{\theta}-\Delta S)^{2}\rangle$. In
Fig. 1(c), we observe that strengthening the nonequilibrium driving (by
increasing $T_{h}$ while keeping $T_{c}=1$) tends to impair the EP estimation.
This is because a stronger driving makes the reverse trajectories of typical
trajectories rarer, lowering the sample quality. Remarkably, the $\alpha$-NEEP
with $\alpha=-0.5$ exhibits the most robust performance. Similarly, in Fig.
1(d), we observe that increasing the amplitude of the periodic potential
$U(x)$ results in the increase of the MSE. This may be the consequence of
rarer movements across the system as the potential well gets deeper, which
means rare trajectories are poorly sampled. Again, the MSE tends to be the
smallest for $\alpha=-0.5$.

As an alternative measure of the estimator’s performance, we also observe the
ratio between the predicted average EP $\sigma_{\text{pred}}$ and the exact
average EP $\sigma$. The results for the two-bead model are shown in Fig.
2(a), which exhibit two different regimes. As $T_{h}$ increases, the estimator
first overestimates and then underestimates the average EP, with
$\sigma_{\text{pred}}/\sigma$ staying closest to $1$ for $\alpha=-0.5$.
Similarly, in Fig. 2(b), we observe that increasing the potential depth $A$ of
the driven Brownian particle leads to an overestimation of the average EP for
$\alpha=0$, while the extent of overestimation stays largely the same for the
other cases; in this case, we could not clearly observe the switching of the
overestimation regime to the underestimation regime as $A$ is increased
further. Nonetheless, again the $\alpha$-NEEP with $\alpha=-0.5$ gives the
best estimation overall.

Figure 3:  Performance of the exactly solvable one-parameter model. (a) The
shift $\Delta\theta$ of the loss function minimum as a function of the
truncation parameter $k$. The circles are results obtained by numerical
minimization, while the solid lines are from the small $1/k$ expansion.
(Inset) The loss function landscapes, with the circles indicating the minima.
We fixed $\mu=3$, $\sigma=1$, and $\alpha=0$. (b) The ratio of the estimated
minimum $\theta^{*}$ to the true minimum $\theta_{0}$ as the bias $\mu$ is
varied. The optimal points are calculated using the criterion that the loss
function gradient satisfies
$|\partial_{\theta}\mathcal{L}_{\alpha}(\theta)|<10^{-3}$ for the first time
as $\theta$ increases from $0$. We fixed $k=4$ and $\sigma=1$. (Inset) The
loss function landscape. While the open diamonds indicate the true minima
$\theta_{0}$, the filled diamonds represent the estimated minima $\theta^{*}$.
The parameters $\alpha=-0.5$ and $k=4$ are fixed. (c) The MSE of $\theta$. The
vertical dashed line shows that the error is minimized at $\alpha=-0.5$. (d)
The distribution of the loss function gradient
$\partial_{\theta}\mathcal{L}_{\alpha}$ at the minimum $\theta_{0}=2$ for
$\mu=\sigma=1$.

An exactly solvable one-parameter model. — To understand the observed
optimality of $\alpha=-0.5$, we simplify the problem to the density-ratio
estimation problem for a 1-d random variable. To be specific, we estimate the
log ratio $s(x)\equiv\log p(x)/p(-x)$ given samples drawn from the
distribution $p(x)$. It is intuitively clear that this problem is structurally
equivalent to EP estimation.

For further simplification, we set

$\displaystyle
p(x)=\begin{cases}\mathcal{N}\exp\left[-\frac{(x-\mu)^{2}}{2\sigma^{2}}\right]&\text{for
$|x-\mu|\leq k\sigma$,}\\\ 0&\text{otherwise.}\end{cases}$ (13)

Here $\mathcal{N}$ is a suitable normalization factor, $\mu$ the positive mean
of the distribution, $\sigma$ the width of the distribution, and $k$ a
positive number truncating the tails of the distribution. While $k=\infty$
corresponds to the perfect sampling of a Gaussian distribution, a finite $k$
corresponds to the case where the tails of the distribution are poorly
sampled.

For $k=\infty$, the correct answer to the problem is a linear function
$s(x)=\theta_{0}x$, where $\theta_{0}\equiv 2\mu/\sigma^{2}$. Thus, for
further simplicity, we focus on the one-parameter model $s_{\theta}(x)=\theta
x$, which estimates $s(x)$ using only a single parameter $\theta$. For this
problem, the suitable loss function is obtained as an analog of Eq. (5):

$\displaystyle\mathcal{L}_{\alpha}(\theta)=\left\langle-\frac{(1+\alpha)e^{\alpha\theta
x}-1}{\alpha}+e^{-(1+\alpha)\theta x}\right\rangle_{p(x)}\;.$ (14)

If $k$ is large but finite, the minimum of this loss function shifts to
$\theta_{0}+\Delta\theta$, where $\Delta\theta$ can be expanded to the leading
orders in $1/k$:

$\displaystyle\Delta\theta\sim\exp\left[-\frac{k^{2}}{2}+\frac{2\mu}{\sigma}\left(\left|\alpha+\frac{1}{2}\right|+\frac{1}{2}\right)k\right]\;.$
(15)

This clearly shows that $\alpha=-0.5$ gives the least shift $\Delta\theta$, as
also illustrated by various results shown in Fig. 3.

In Fig. 3(a), we show that the shift of the minimum $\Delta\theta$ tends to
increase as the tail sampling becomes poorer (i.e., $k$ decreases). The
landscapes of the loss function $\mathcal{L}_{\alpha}(\theta)$, shown in the
inset of Fig. 3(a), also confirm this observation. The increase of the error
with the potential depth $A$ in Figs. 1(d) and 2(b) may primarily be due to
the same effect.

In Fig. 3(b), we plot the ratio between the estimated minimum $\theta^{*}$ and
the true minimum $\theta_{0}$ as a function of the mean $\mu$, which is an
analog of the nonequilibrium driving. We note that here $\theta^{*}$ is the
lowest value of $\theta$ at which the slope of the loss function becomes less
then $10^{-3}$. We observe that an overestimation regime
($\theta^{*}/\theta_{0}>1$) crosses over to an underestimation regime
($\theta^{*}/\theta_{0}<1$) as $\mu$ grows. This is in striking agreement with
the trends shown in Fig. 2(a). The reason why $\theta^{*}$ underestimates
$\theta_{0}$ for large $\mu$ can be understood by the flattened loss function
landscapes shown in the inset of Fig. 3(b). In this regime, the dynamics of
$\theta$ (starting from $\theta$ = 0) slows down, ending up at a value (filled
diamonds) even lower than $\theta_{0}$ (empty diamonds). This effect is due to
the samples with $x<0$ vanishing when $\mu$ is too large.

The one-parameter model also allows us to examine the effects of the finite
mini-batch size $M$. While the ideal loss function is given in Eq. (14), the
loss function used in the actual training looks like

$\displaystyle\mathcal{L}_{\alpha}(\theta;M)=\frac{1}{M}\sum_{i}{\left[-\frac{(1+\alpha)e^{\alpha\theta
X_{i}}-1}{\alpha}+e^{-(1+\alpha)\theta X_{i}}\right]}\;,$ (16)

where $X_{1},\,\ldots,\,X_{M}$ are i.i.d. Gaussian random variables of mean
$\mu$ and variance $\sigma^{2}$. When $M$ is large and finite, using the
central limit theorem (CLT), the gradient of this loss function can be
approximated as Li  _et al._ (2017); Chaudhari and Soatto (2018)

$\displaystyle\frac{\partial\mathcal{L}_{\alpha}}{\partial\theta}\bigg{|}_{\theta=\theta_{t}}=\bar{K}(\theta_{t}-\theta_{0})+\sqrt{\frac{\Lambda}{M}}N_{t}+o(M^{-1/2})\;,$
(17)

where $\theta_{0}=\text{argmin}(\langle\mathcal{L}_{\alpha}(\theta)\rangle)$,
$\bar{K}=\partial^{2}_{\theta}\langle\mathcal{L}_{\alpha}\rangle\big{|}_{\theta=\theta_{0}}$,
and
$\Lambda=\text{Var}\big{[}\partial_{\theta}\mathcal{L}_{\alpha}\big{|}_{\theta=\theta_{0},\,M=1}\big{]}$.
When the stochastic gradient descent
$\theta_{t+1}=\theta_{t}-\lambda(\partial\mathcal{L}_{\alpha}/\partial\theta)\big{|}_{\theta=\theta_{t}}$
reaches the steady state, the MSE of $\theta$ is given by

$\displaystyle\langle\left(\theta-\theta_{0}\right)^{2}\rangle_{s}=\frac{\lambda\Lambda}{M\bar{K}(2-\lambda\bar{K})}+o(M^{-1})\;.$
(18)

This leading-order behavior is shown in Fig. 3(c) for various values of $\mu$.
For all cases, the MSE of $\theta$ is minimized at $\alpha=-0.5$, which is
consistent with the smallest error bars observed at $\alpha=-0.5$ in Figs. 1
and 2. Hence, $\alpha=-0.5$ yields the most consistent EP estimator.

Direct measurements of the loss function gradient at the minimum also confirm
the above result. As shown in Fig. 3(d), the gradient
$\partial_{\theta}\mathcal{L}_{\alpha}$ is far more broadly distributed for
$\alpha=0$ than for $\alpha=-0.5$. Moreover, due to the subleading effects
(beyond the CLT) of finite $M$, the gradient for $\alpha=0$ features a large
skewness. These show that the training dynamics for the original NEEP
($\alpha=0$) tends to be far more volatile and unstable than for the
$\alpha$-NEEP with $\alpha=-0.5$.

Summary and outlook. — We proposed the $\alpha$-NEEP, a generalization of the
NEEP for estimating the steady-state EP at the trajectory level. Comparing the
performance for various values of $\alpha$, we found that $\alpha=-0.5$ is the
best in the following regards: (i) it is the most robust against poor sampling
caused by strong nonequilibrium driving or deep potential wells; (ii) it gives
the most consistent EP estimate with smallest error bars. We also demonstrated
the optimality of $\alpha=-0.5$ using a simplification of the original EP
estimation problem, whose loss function landscape and relaxation properties
are analytically tractable. The $\alpha$-NEEP thus provides a powerful method
for estimating the EP for much broader range of the nonequilibrium driving
force and the time scale of dynamics. Identification of even better loss
functions and optimization of other hyperparameters (network size, number of
iterations, etc.) are left as future works. It would also be interesting to
apply the $\alpha$-NEEP to estimations of the EP of the Brownian movies Bae
_et al._ (2022) and stochastic systems with odd-parity variables Kim  _et al._
(2022), which have been studied using the original NEEP method.

Acknowledgments. — This work was supported by the National Research Foundation
of Korea Grant funded by the Korean Government (NRF-2020R1C1C101443613). EK
and YB also thank Junghyo Jo and Sangyun Lee for helpful comments.

##  Appendix A Training details

We always use the fully connected network (FCN) with three hidden layers, with
each layer composed of $512$ nodes. Each training dataset consists of $10^{6}$
trajectories. The neural network parameters are updated using the ReLU
activation function and the Adam optimizer. The learning rate is fixed to
$10^{-5}$ and the weight decay is fixed to $5\times 10^{-5}$. We halt the
training after $10,000$ iterations, except for the results shown in Figs. S3
and S4. All trainings are done on PyTorch with NVIDIA GeForce RTX 3090.

In Fig. 1(c, d), each mini-batch consists of $10^{4}$ trajectories. In Fig.
2(a, b), each mini-batch consists of $10^{5}$ trajectories.

##  Appendix B Density ratio estimation using $f$-divergence

Consider a convex, twice-differentiable real-valued function $f(u)$. Then, the
inequality

$\displaystyle-pf^{\prime}(u)+q\left(uf^{\prime}(u)-f(u)\right)\geq-qf(p/q)\;$
(19)

holds. We can verify this by differentiating the lhs with respect to $u$,
which yields $f^{\prime\prime}(u)(-p+qu)$. Thus the lhs has a local minimum at
$u=p/q$, and this is the only local minimum since $f$ is convex. In addition,
the second derivative of the lhs at $u=p/q$ equals $qf^{\prime\prime}(p/q)$,
which is positive by the convexity. This proves the inequality (19). Using
this, we can design a loss function whose minimum is equal to the negative
$f$-divergence between two probability distributions $p(\mathbf{x})$ and
$q(\mathbf{x})$. To be specific, for any function $r(\mathbf{x})$, we define

$\displaystyle\mathcal{L}_{f}[r]$
$\displaystyle=\left\langle-f^{\prime}(r(\mathbf{x}))+\frac{q(\mathbf{x})}{p(\mathbf{x})}\left\\{r(\mathbf{x})f^{\prime}(r(\mathbf{x}))-f(r(\mathbf{x}))\right\\}\right\rangle_{p(\mathbf{x})}$
$\displaystyle=\int
d^{2d}\mathbf{x}\,[-p(\mathbf{x})f^{\prime}\\!\left(r(\mathbf{x})\right)+q(\mathbf{x})\left\\{r(\mathbf{x})f^{\prime}\\!\left(r(\mathbf{x})\right)-f\\!\left(r(\mathbf{x})\right)\right\\}]\;.$
(20)

Using Eq. (B), we conclude that

$\displaystyle\mathcal{L}_{f}[r]\geq-\int
d^{2d}\mathbf{x}\,q(\mathbf{x})\,f\\!\left(\frac{p(\mathbf{x})}{q(\mathbf{x})}\right)=-D_{f}[p\\!:\\!q]\;,$
(21)

where $D_{f}[p{{\,:\,}}q]$ is the $f$-divergence between the distributions
$p(\mathbf{x})$ and $q(\mathbf{x})$, and the equality holds if and only if
$r(\mathbf{x})=p(\mathbf{x})/q(\mathbf{x})$ for all $\mathbf{x}$. By
minimizing $\mathcal{L}_{f}[r]$, we can estimate $p(\mathbf{x})/q(\mathbf{x})$
as well as $D_{f}[p\\!:\\!q]$.

The loss function $\mathcal{L}_{\alpha}$ and the associated
$\alpha$-divergence discussed in the main text are obtained by choosing the
function $f$ to be

$\displaystyle
f_{\alpha}(u)=\begin{cases}\frac{u^{1+\alpha}-(1+\alpha)u+\alpha}{\alpha(1+\alpha)}&\text{for
$\alpha\neq 0,\,-1$,}\\\ u\log u&\text{for $\alpha=0$,}\\\ \log
u+1-u&\text{for $\alpha=-1$.}\end{cases}\;.$ (22)

Note that $f_{0}(u)=\lim_{\alpha\to 0}f_{\alpha}(u)$ and
$f_{-1}(u)=\lim_{\alpha\to-1}f_{\alpha}(u)$. It is straightforward to obtain
Eq. 5 of the main text and its extensions to the cases $\alpha=0$ and
$\alpha=-1$ from this choice.

Figure 4:  Coefficient of determination $R^{2}$. $10^{4}$ trajectories are
used for each mini-batch, and error bars indicate the standard deviations
obtained from $40$ independent trainings. (a) For the two-bead model, $R^{2}$
exhibits a non-monotonic behavior as a function of $T_{h}$ (while $T_{c}=1$).
The decrease of $R^{2}$ near equilibrium ($T_{h}=1$) is due to the overfitting
phenomenon, which disrupts the linear relationship between $s_{\theta}$ and
$\Delta S$. (b) For the driven Brownian particle, $R^{2}$ always increases
with $A$. This may seem contradictory to how the MSE increases with $A$ in
Fig. 1(d) of the main text. Indeed, higher $R^{2}$ only means that there is a
good linear relationship between the EP estimate $s_{\theta}$ and the true EP
$\Delta S$, not that $s_{\theta}$ and $\Delta S$ are close to each other. This
clearly shows that $R^{2}$ is not an adequate measure of the performance of EP
estimators. Figure 5:  Effects of the mini-batch size on the performance of
the $\alpha$-NEEP for the two-bead model with $T_{h}=1000$. We use three
different measures of the performance. Error bars indicate the standard
deviations obtained from $40$ independent trainings. For small mini-batches,
the highly skewed distribution of the stochastic gradient causes
underestimation of the EP. For large mini-batches, the distortion of the loss
function landscape stemming from poor sampling leads to overestimation of the
EP. All three measures confirm the superiority of $\alpha=-0.5$ to $\alpha=0$.
Figure 6:  Training dynamics of the $\alpha$-NEEP for the two-bead model at
$T_{h}=10$. Each mini-batch consists of $10^{5}$ trajectories. The first
(last) two columns show the dynamics for $\alpha=-0.5$ ($\alpha=0$). The first
and the third columns show the three different measures of performance for the
training dataset and the test dataset. The second and the fourth columns show
the difference between the corresponding measures obtained for two datasets.
We note that the test dataset is independently generated. The overfitting
phenomena are manifest from the decrease of $R^{2}$ and the increase of the
MSE towards the end of the training. Interestingly, overfitting leads to an
overestimation of the average EP only for the training dataset. We also note
that the value of $\alpha$ is largely irrelevant to the extent of overfitting.
This phenomenon can be explained as follows. Near equilibrium, the neural
network swiftly reaches the loss function minimum. However, as the training
continues, the neural network starts to see the detailed fluctuations of the
training dataset. This makes the functional form of $s_{\theta}$ very rough,
leading to the increase of the MSE and $R^{2}$ for both datasets. But while
the neural network now believes all trajectories in the training dataset to be
highly irreversible and assigns high EP to them, the EP assigned to the
trajectories in the test dataset stay unbiased.

Figure 7:  Training dynamics of the $\alpha$-NEEP for the two-bead model at
$T_{h}=3000$. Each mini-batch consists of $10^{5}$ trajectories. The
subfigures are organized in exactly the same way as in Fig. S3. In this case,
the overfitting effects are not pronounced, and the differences between the
training and the test datasets are very small. Note that $\alpha=0$ exhibits
strong noise effects, which is in agreement with the gradient distribution
shown in Fig. 3(d) of the main text.

## References

  * Brangwynne  _et al._ (2008) C. P. Brangwynne, G. H. Koenderink, F. C. MacKintosh, and D. A. Weitz, J. Cell Biol. 183, 583 (2008).
  * Weber  _et al._ (2012) S. C. Weber, A. J. Spakowitz, and J. A. Theriot, Proc. Natl. Acad. Sci. 109, 7338 (2012).
  * Mizuno  _et al._ (2007) D. Mizuno, C. Tardin, C. F. Schmidt, and F. C. MacKintosh, Science 315, 370 (2007).
  * Battle  _et al._ (2016) C. Battle, C. P. Broedersz, N. Fakhri, V. F. Geyer, J. Howard, C. F. Schmidt, and F. C. MacKintosh, Science 352, 604 (2016).
  * Gladrow  _et al._ (2016) J. Gladrow, N. Fakhri, F. C. MacKintosh, C. F. Schmidt, and C. P. Broedersz, Phys. Rev. Lett. 116, 248301 (2016).
  * Tailleur and Cates (2008) J. Tailleur and M. E. Cates, Phys. Rev. Lett. 100, 218103 (2008).
  * Speck  _et al._ (2014) T. Speck, J. Bialké, A. M. Menzel, and H. Löwen, Phys. Rev. Lett. 112, 218304 (2014).
  * Takatori  _et al._ (2014) S. C. Takatori, W. Yan, and J. F. Brady, Phys. Rev. Lett. 113, 028103 (2014).
  * Farage  _et al._ (2015) T. F. F. Farage, P. Krinninger, and J. M. Brader, Phys. Rev. E 91, 042310 (2015).
  * Solon  _et al._ (2018) A. P. Solon, J. Stenhammar, M. E. Cates, Y. Kafri, and J. Tailleur, Phys. Rev. E 97, 020602(R) (2018).
  * Fodor  _et al._ (2016) E. Fodor, C. Nardini, M. E. Cates, J. Tailleur, P. Visco, and F. van Wijland, Phys. Rev. Lett. 117, 038103 (2016).
  * Nardini  _et al._ (2017) C. Nardini, E. Fodor, E. Tjhung, F. van Wijland, J. Tailleur, and M. E. Cates, Phys. Rev. X 7, 021007 (2017).
  * Seifert (2012) U. Seifert, Rep. Prog. Phys. 75, 126001 (2012).
  * Roldán and Parrondo (2010) E. Roldán and J. M. R. Parrondo, Phys. Rev. Lett. 105, 150607 (2010).
  * Roldán and Parrondo (2012) E. Roldán and J. M. R. Parrondo, Phys. Rev. E 85, 031129 (2012).
  * Li  _et al._ (2019) J. Li, J. M. Horowitz, T. R. Gingrich, and N. Fakhri, Nat. Commun. 10, 1666 (2019).
  * Van Vu  _et al._ (2020) T. Van Vu, V. T. Vo, and Y. Hasegawa, Phys. Rev. E 101, 042138 (2020).
  * Otsubo  _et al._ (2020) S. Otsubo, S. Ito, A. Dechant, and T. Sagawa, Phys. Rev. E 101, 062106 (2020).
  * Otsubo  _et al._ (2022) S. Otsubo, S. K. Manikandan, T. Sagawa, and S. Krishnamurthy, Commun. Phys. 5, 11 (2022).
  * (20) S. Lee, D.-K. Kim, J.-M. Park, W. K. Kim, H. Park, and J. S. Lee, “Multidimensional entropic bound: Estimator of entropy production for general langevin dynamics with an arbitrary time-dependent protocol,” arXiv:2207.05961 [cond-mat.stat-mech] .
  * Kim  _et al._ (2020) D.-K. Kim, Y. Bae, S. Lee, and H. Jeong, Phys. Rev. Lett. 125, 140604 (2020).
  * Basu  _et al._ (1998) A. Basu, I. R. Harris, N. L. Hjort, and M. C. Jones, Biometrika 85, 549 (1998).
  * Sugiyama  _et al._ (2012) M. Sugiyama, T. Suzuki, and T. Kanamori,  _Density ratio estimation in machine learning_ (Cambridge University Press, 2012).
  * Nowozin  _et al._ (2016) S. Nowozin, B. Cseke, and R. Tomioka, in  _Proceedings of the 30th International Conference on Neural Information Processing Systems_ (2016) pp. 271–279.
  * Belghazi  _et al._ (2018) M. I. Belghazi, A. Baratin, S. Rajeshwar, S. Ozair, Y. Bengio, A. Courville, and D. Hjelm, in  _International Conference on Machine Learning_ (PMLR, 2018) pp. 531–540.
  * Li  _et al._ (2017) Q. Li, C. Tai, and W. E, in  _Proceedings of the 34th International Conference on Machine Learning_, Proceedings of Machine Learning Research, Vol. 70, edited by D. Precup and Y. W. Teh (PMLR, 2017) pp. 2101–2110.
  * Chaudhari and Soatto (2018) P. Chaudhari and S. Soatto, in  _2018 Information Theory and Applications Workshop (ITA)_ (IEEE, 2018) pp. 1–10.
  * Bae  _et al._ (2022) Y. Bae, D.-K. Kim, and H. Jeong, Phys. Rev. Res. 4, 033094 (2022).
  * Kim  _et al._ (2022) D.-K. Kim, S. Lee, and H. Jeong, Phys. Rev. Res. 4, 023051 (2022).