# Artificial-intelligence-based surrogate solution of dissipative quantum
dynamics: physics-informed reconstruction of the universal propagator

Jiaji Zhang  0000-0003-2978-274X<EMAIL_ADDRESS>Zhejiang
Laboratory, Hangzhou 311100, China  Carlos L. Benavides-Riveros
0000-0001-6924-727X<EMAIL_ADDRESS>Pitaevskii BEC Center, CNR-
INO and Dipartimento di Fisica, Università di Trento, I-38123 Trento, Italy
Lipeng Chen  0000-0003-1415-4767<EMAIL_ADDRESS>Zhejiang Laboratory,
Hangzhou 311100, China

###### Abstract

The accurate (or even approximate) solution of the equations that govern the
dynamics of dissipative quantum systems remains a challenging task for quantum
science. While several algorithms have been designed to solve those equations
with different degrees of flexibility, they rely mainly on highly expensive
iterative schemes. Most recently, deep neural networks have been used for
quantum dynamics but current architectures are highly dependent on the physics
of the particular system and usually limited to population dynamics. Here we
introduce an artificial-intelligence-based surrogate model that solves
dissipative quantum dynamics by parameterizing quantum propagators as Fourier
neural operators, which we train using both dataset and physics-informed loss
functions. Compared with conventional algorithms, our quantum neural
propagator avoids time-consuming iterations and provides a universal super-
operator that can be used to evolve any initial quantum state for arbitrarily
long times. To illustrate the wide applicability of the approach, we employ
our quantum neural propagator to compute population dynamics and time-
correlation functions of the Fenna-Matthews-Olson complex.

##  I Introduction

Since atomic and molecular systems are never completely isolated from the
surrounding environment, a proper description of their dynamics requires the
theory of open quantum systems [1, 2]. The environment is usually treated as a
heat bath, and thus a commonly used strategy in the theory of open quantum
systems is to trace out bath degrees of freedom and derive the corresponding
equations of motion (EOM) for the system, also referred to as quantum master
equations (QMEs) [3, 4]. Under perturbative and Markovian approximations, the
best-known QMEs are the Lindblad [5, 6] and Redfield [7] equations.
Numerically exact QME approaches, such as the hierarchy equations of motion
[8, 9] and the quasiadiabatic propagator path integral [10, 11], are developed
to capture the intrinsic non-perturbative and non-Markovian system dynamics.
The modern technique of ultrafast spectroscopy provides essential insights
into the intricate system-bath interactions and bath-induced electronic
dephasing and energy relaxation processes [12, 13, 14]. The key quantity for
simulating optical spectroscopy is polarization, which is linked to various
response functions made calculable by QMEs [15]. QMEs approaches have also
been utilized to simulate both linear and nonlinear spectra of molecular
systems [16, 17, 18].

Apart from their different concrete forms or range of validity, a common
ingredient of the EOM for open quantum systems is the presence of a partial
differential equation with respect to time which is usually well-suited for
iterative numerical solvers [19]. The most widely used numerical solvers are
the family of Runge-Kutta methods, which iteratively integrate the EOM with a
small enough time step $(\delta t)$ from a given initial condition. As a
universal formalism, no restrictions are a priori made on the explicit form of
the EOM or the initial states. Runge-Kutta methods thus provide a versatile
tool to solve EOM for a wide variety of scenarios such as rate equations in
chemical reactions [20, 21], population dynamics in charge and energy transfer
processes [17, 22], and spectroscopic simulations [23, 24]. Despite its
straightforward implementation, the drawbacks of Runge-Kutta methods are the
large computational cost and the presence of numerical instabilities
associated with the iterations, especially for long-time dynamics.
Improvements have been proposed to overcome these limitations by deriving
modified schemes of EOM [25, 26, 27, 28, 29] or implementing other numerical
solvers (e.g., proper orthogonal decomposition [30]). However, since they are
usually limited to a few specific scenarios, an efficient universal solver is
still to be proposed.

Over recent years, artificial intelligence (AI) has emerged as a powerful tool
that, integrated into research, is augmenting and accelerating scientific
discovery [31]. It is also increasingly being seen as a useful tool to solve
various physical and chemical problems that have remained so far very
challenging for conventional methods [32, 33]. Deep neural networks are the
best-known example of AI tools applied to quantum systems. The universal
approximation theorem of (Borel-measurable) functions [34] provides the basis
for the neural-network parameterization of wavefunctions for quantum spin
systems [35, 36]. When further extended to many-electron systems, they
outperformed traditional variational ansätze [37, 38, 39, 32]. Yet, while they
are somehow routinely applied to study static (ground-state) properties, there
are relatively few attempts to simulate the dynamics of quantum systems [40,
41, 42, 43], and existing dynamical deep neural-network models are usually
highly dependent on the physics of specific systems, limited to population
dynamics [42], or tested on toy models only [44, 45, 46].

While the aforementioned works are concerned with the approximation of quantum
states, AI can also be used as a surrogate model to approximate highly
nonlinear operators. Based on the universal approximation theorem of
functionals and mathematical operators [47], the solution operator of a given
partial differential equation can be parameterized as a deep neural network
that describes the mapping between initial states and the evolved state at
some subsequent time [48, 49, 50]. Recently developed architectures, such as
DeepONet [51, 49] and Fourier Neural Operator (FNO) [52, 53], have
demonstrated their superiority on classical differential equations in weather
forecasting [54, 55] and latent diffusion models in text-to-image
transformation [56]. Despite this progress, surrogate models for quantum
operators have only been applied, to the best of our knowledge, to scattering
processes [57].

In this work, our objective is to develop an AI-based surrogate model for the
universal solution of EOM for open quantum systems. To this end, we train a
neural operator as the (universal) quantum propagator of dissipative quantum
dynamics using both dataset and physics-informed loss functions. We employ the
FNO architecture to parameterize the propagator for the Lindblad QME, which is
the quantum analog of the solution operator in a typical surrogate model. Our
approach is quite versatile and can be easily extended to solve other EOM for
quantum dynamics. The trained propagator allows for the direct computation of
dynamics up to a chosen time limit $t_{\mathrm{max}}$ for any initial state
through a single-step operation without invoking the tedious, expensive
iterations.

Moreover, since the neural propagator obeys the usual composition property of
quantum operators, the method can easily be extended to arbitrarily long
times. In addition to the simulation of population dynamics, our neural
propagator can be used to compute the more challenging time-correlation
functions of system operators. We test our approach by training a neural
operator as the quantum propagator of the well-known Fenna-Matthews-Olson
complex and compute population dynamics as well as multi-time correlation
functions of system operators.

##  II Result

###  II.1 The QME propagator and time-correlation functions

We consider the typical Fenna-Matthews-Olson pigment-protein complex found in
green sulfur bacteria as our model system [58, 59]. The interplay between
molecular excitations, environment interaction, and quantum coherence effects
makes the complex one of the most important workhorses for developing and
improving simulation methods [60, 61, 62]. The dissipative dynamics of the
reduced density operator $\hat{\rho}$ for the electronic sub-system is
described by the Lindblad QME,

$\frac{\partial}{\partial
t}\hat{\rho}(t)=-\frac{i}{\hbar}\left[\hat{H}_{el},\leavevmode\nobreak\
\hat{\rho}(t)\right]-\sum_{j=1}^{N}\frac{\lambda_{j}}{2}\left(\hat{V}_{j}^{\dagger}\hat{V}_{j}\hat{\rho}(t)+\hat{\rho}(t)\hat{V}_{j}^{\dagger}\hat{V}_{j}-2\hat{V}_{j}\hat{\rho}(t)\hat{V}_{j}^{\dagger}\right),$
(1)

where $\hat{H}_{el}$ is the Hamiltonian of the electronic states [22]

$\hat{H}_{el}=\sum_{j=1}^{N}\varepsilon_{j}|{j}\rangle\langle{j}|+\sum_{j\neq
j^{\prime}}\Delta_{j,j^{\prime}}|{j}\rangle\langle{j^{\prime}}|,$ (2)

with $\varepsilon_{j}$ being the energy of the $j$-th state $|j\rangle$ (there
are $N$ of them) and $\Delta_{j,j^{\prime}}$ the inter-state couplings. The
second part on the right-hand side of Eq. 1 represents a pure-dephasing
Lindblad operator with $\hat{V}_{j}=|j\rangle\langle j|$ and a pure-dephasing
rate $\lambda_{j}$.

To alleviate the notation, we introduce the abbreviated index
$x=(j,j^{\prime})$ and align the matrix entries of the reduced density
operator as column vector,
$\vec{\rho}_{t}=\left\\{\rho(x_{0},t),\rho(x_{1},t),...\right\\}$, with
$\rho(x,t)=\langle{j}|\hat{\rho}(t)|{j^{\prime}}\rangle$. Eq. (1) can be
recast to a matrix-vector form as

$\frac{\partial}{\partial t}{\vec{\rho}}_{t}={\bm{L}}{\vec{\rho}}_{t},$ (3)

where the matrix entries of ${\bm{L}}$ is inferred from the right-hand side of
Eq. (1).

The QME propagator ${\bm{G}}_{t}$ is defined through the integration of Eq.
(3) as

${\vec{\rho}}_{t}={\bm{G}}_{t}{\vec{\rho}}_{0}=e^{t{\bm{L}}}{\vec{\rho}}_{0},$
(4)

where ${\bm{G}}_{t}$ can be regarded as a time-dependent matrix that always
acts on the vector to its right. The evaluation of dynamics up to arbitrarily
long times can be inferred from the well-known composition property of quantum
operators as

${\vec{\rho}}_{t_{1}+t_{2}}={\bm{G}}_{t_{2}}{\vec{\rho}}_{t_{1}}={\bm{G}}_{t_{2}}{\bm{G}}_{t_{1}}{\vec{\rho}}_{0}.$
(5)

The first-order and second-order time-correlation functions (TCFs) are defined
through the propagator as (see “Methods” for more details):

$R^{(1)}(t_{1})={\bm{X}}_{\rm
tr}{\bm{G}}_{t_{1}}{\bm{X}}_{\times}{\vec{\rho}}_{0},$ (6)
$R^{(2)}(t_{1},t_{2})={\bm{X}}_{\rm
tr}{\bm{G}}_{t_{2}}{\bm{X}}_{\times}{\bm{G}}_{t_{1}}{\bm{X}}_{\times}{\vec{\rho}}_{0},$
(7)

where ${\bm{X}}_{\times}$ and ${\bm{X}}_{\rm tr}$ recover the operations for
any operator $\hat{X}$ as

${\bm{X}}_{\times}\vec{\rho}=\frac{i}{\hbar}\left[\hat{X},\,\hat{\rho}\right],$
(8) ${\bm{X}}_{\rm tr}\vec{\rho}=\sum_{j}\langle
j|\hat{X}\hat{\rho}|j\rangle.$ (9)

Conventional methods such as Runge-Kutta evaluate Eqs. (6) and (7) by starting
from an initial state ${\vec{\rho}}_{0}$ and iteratively propagate Eq. (3) for
all combinations of $t_{1}$ and $t_{2}$. The overall computational cost thus
scales exponentially with the order of TCFs, limiting, as a consequence,
iteration-based methods to low-order cases.

###  II.2 Neural propagator within the FNO architecture

Within the AI surrogate framework, we use the FNO architecture to parametrize
the propagator ${\bm{G}}_{t}(\theta)$ ($\theta$ denote the set of trainable
parameters) up to a chosen time limit $t_{\mathrm{max}}$. The model takes the
initial density matrix ${\vec{\rho}}_{0}$ and time $t\in[0,t_{\mathrm{max}}]$
as the input, and outputs the density matrix
${\vec{\rho}}_{t}(\theta)={\bm{G}}_{t}(\theta)\vec{\rho}_{0}$ that satisfies
Eq. (3). As this parametrization retains the composition property of Eq. (5),
the model can be extended far beyond $t_{\mathrm{max}}$ by iteratively
applying the neural operator. In addition, the constructed FNO model works for
any form of ${\vec{\rho}}_{0}$, including for instance the
${\bm{X}}_{\times}{\vec{\rho}}_{0}$. Thus, to make use of the full power of
the AI-based architecture we will substitute the algebraic propagator in Eqs.
(6) and (7) by the neural propagator ${\bm{G}}_{t}(\theta)$.

To train the FNO model, we prepare a dataset by first randomly sampling a set
of initial states and then evaluating their time evolution using fourth-order
Runge-Kutta (RK4) with a chosen time step $\delta t$. The dataset loss
function is defined as follows:

$\mathcal{L}_{\mathrm{data}}=\int_{0}^{t_{\mathrm{max}}}dt\,\frac{\left|\left|{\bm{G}}_{t}(\theta)\vec{\rho}_{0}-\vec{\rho}_{t}\right|\right|_{F}}{\left|\left|{\bm{G}}_{t}(\theta)\vec{\rho}_{0}\right|\right|_{F}},$
(10)

where $\vec{\rho}_{t}$ is the corresponding data, and $||\cdot||_{F}$ denotes
the Frobenius-norm. In addition, we add to this loss function a physics-
informed part, defined as

$\mathcal{L}_{\mathrm{phys}}=\int_{0}^{t_{\mathrm{max}}}dt\frac{\left|\left|\frac{\partial}{\partial
t}{\bm{G}}_{t}(\theta)\vec{\rho}_{0}-\bm{L}({\bm{G}}_{t}(\theta)\vec{\rho}_{0})\right|\right|_{F}}{\left|\left|\bm{L}({\bm{G}}_{t}(\theta)\vec{\rho}_{0})\right|\right|_{F}}+\frac{\left|\left|{\bm{G}}_{t=0}(\theta)\vec{\rho}_{0}-\vec{\rho}_{0}\right|\right|_{F}}{\left|\left|\vec{\rho}_{0}\right|\right|_{F}}.$
(11)

The first term is the residual of Eq. (3), and the second term is introduced
to ensure the model reproduces the identity when $t=0$. Importantly, notice
that only the inputs $\vec{\rho}_{0}$ are required for the evaluation of the
losses in $\mathcal{L}_{\mathrm{phys}}$. We can thus employ the on-the-fly
sampling to generate those additional initial states. More details on the
model and the training can be found in Section “Methods”.

###  II.3 Applications to dissipative quantum dynamics

We apply the trained FNO propagator to the excitation energy transfer of the
seven-site Fenna-Matthews-Olson complex. To test our model, we compare results
for population dynamics and TCFs with those from the RK4. Fig. 1 shows the
population dynamics obtained from both FNO propagator and RK4 up to
$50t_{\mathrm{max}}$ with $t_{\mathrm{max}}=30$fs. Here we consider two
typical initial states: an excitation located on (a) site 1, i.e.,
$\hat{\rho}_{0}=|{1}\rangle\langle{1}|$ and (b) site 6, i.e.,
$\hat{\rho}_{0}=|{6}\rangle\langle{6}|$. Following the experimentally
demonstrated energy transfer pathway [22, 59], we focus on the time evolution
of populations $p_{n}(t)=\langle n|\hat{\rho}_{t}|n\rangle$ for sites (a) $n=$
1, 2, and 3, and (b) $n=$ 4, 5, and 6. Within the range
$t\in[0,50t_{\mathrm{max}}]$, the FNO propagator yields population dynamics in
excellent agreement with those obtained from the RK4. In both cases, the
results are essentially exact up to $20t_{\mathrm{max}}$. It is thus quite
remarkable that our FNO propagator can infer truly long-time dynamics from the
relatively short-time training data we used (recall that the propagator was
trained with data in the interval $[0,t_{\mathrm{max}}]$).

Figure 1: Population dynamics of the FMO complex computed using RK4 (solid
line) and FNO neural propagator (dashed line) for two initial states. The time
is presented in units of the maximum training time $t_{\mathrm{max}}=30$fs.

Next, we apply the FNO propagator to calculate first-order and second-order
TCFs using the initial condition of $\hat{\rho}_{0}=|1\rangle\langle{1}|$ and
the system operator

$\hat{X}=\sum_{j=1}^{N-1}\left(|{j}\rangle\langle{j+1}|+|{j+1}\rangle\langle{j}|\right),$
(12)

which is arbitrarily chosen for the demonstration of our algorithm. Although
we can choose any other type of operator (e.g., to measure the correlation of
populations of different sites or even coherences), the rich structure of
$\hat{X}$, which corresponds to the hopping term of the lattice Hamiltonian,
makes it an excellent test-bed for evaluating the performance of our neural
propagator when tasked with computing TCFs. To facilitate the comparison of
numerical results, we focus on the Fourier transform of TCFs over the time
variables.

Figure 2: The imaginary part of Fourier transform of the first-order TCF (6).
The blue and red curves are results obtained from the FNO propagator and the
RK4, respectively. In each case, the peak intensities are normalized with
respect to their maximum value.

In Fig. 2, we show the imaginary part of Fourier transform of the first-order
TCF ($t_{1}\in[0,50t_{\mathrm{max}}]$) evaluated using both the FNO propagator
and the RK4. For each case, we normalize peak intensities with respect to
their maximum value. In Fourier domain, the spectrum obtained from the FNO
propagator matches perfectly with that from the RK4. The small deviation
around 5 $\rm{cm^{-1}}$ may be attributed to the specific architecture of the
FNO. It is important to note that the quantum system in this work is described
by matrix entries over a discrete Hilbert space, which may differ in nature
from the classical system in continuous coordinate space. Indeed, it has been
pointed out that in certain circumstances, applying the FNO architecture
(which relies on Fourier transformations) to discontinuous systems may yield
non-optimal solutions [49]. Nevertheless, overall, our FNO propagator produces
results in excellent agreement with those obtained using the RK4 method.

In Fig. 3, we show the imaginary part of Fourier transform of the second-order
TCF ($t_{1},t_{2}\in[0,40t_{\mathrm{max}}]$). As defined in Eq. (7), the
presence of the second operator $\bm{X}_{\times}$ between propagators
$\bm{G}_{t_{1}}(\theta)$ and $\bm{G}_{t_{2}}(\theta)$ generates additional
transitions during propagation. Accurate simulation thus requires the
propagator to be universally applicable to any density matrix. As shown in
Fig. 3(a) and 3(b), the results obtained from the FNO propagator are again in
good agreement with those from the RK4, featuring very precise peak
localization. Despite the complexity of the signal, the mean absolute error
(i.e., $|R^{(2)}_{\rm{FNO}}-R_{\rm{RK4}}^{(2)}|$) is just $4.8\times 10^{-3}$.
Still, we identify a slight tendency of the neural propagator to overestimate
the envelopes of the peaks. It’s worth noting that computing the second-order
TCF takes only a few minutes when using the FNO propagator, whereas, with the
RK4 method, it requires around one hour due to the iterations involved. The
extension to the calculations of TCFs involving more time variables is
straightforward, and one can reasonably expect an even more prominent
computational speed-up. In our opinion, this demonstrates the broad
applicability of our FNO propagator.

Figure 3: The imaginary part of Fourier transform of the second-order TCF (7).
The panels (a) and (b) are results obtained from the FNO propagator and the
RK4.

##  III Discussion and outlook

In this work, we develop an AI-based surrogate model to simulate dissipative
quantum dynamics. Unlike typical deep neural networks that parametrize quantum
states, our surrogate model parameterizes quantum propagators as universal
neural operators, which are capable of mapping initial quantum states to their
evolved counterparts at any subsequent time. The neural operator is trained
using both dataset and physics-informed loss functions that ensure the
universality of the trained propagator. The accuracy of the surrogate model is
gauged by calculating both population dynamics and multi-time correlation
functions of the Fenna-Matthews-Olson complex, yielding results in excellent
agreement with those from conventional RK4. Quite remarkably, our FNO
propagator, besides avoiding the conventional time-consuming iterations,
captures quantum dynamics far beyond the small time window chosen in our
training sets.

This work lays the foundation for further exploration of neural propagators.
Currently, our focus is on the time-independent case, where the dynamics
depends solely on the Hamiltonian. Future investigations may extend to
scenarios involving driven quantum systems, where the time evolution is
influenced by the interaction with the external fields or by non-Markovian
dynamics [4, 63]. In addition, in many non-dissipative scenarios, using wave
functions alone may suffice. For such physical settings, we expect an
excellent performance of our neural propagators (which were applied here to
non-pure quantum states). Furthermore, our AI-based approach could be extended
to compute photoemission spectra or transport dynamics in ab-initio quantum
chemistry, where many existing numerical methods such as the $GW$
approximation, coupled-cluster techniques, or nonequilibrium Green’s
functions, still face several shortcomings at both the fundamental and
practical levels [64, 65, 66, 67, 68].

##  IV Methods

###  IV.1 The model Hamiltonian and propagator

The Fenna-Matthews-Olson protein complex is modeled as an electronic system
interacting with heat baths which consist of a set of harmonic oscillators.
The total Hamiltonian is defined as

$\hat{H}_{tot}=\hat{H}_{el}+\hat{H}_{ph}+\hat{H}_{el-ph}+\hat{H}_{reorg}.$
(13)

The first term $\hat{H}_{el}$ is the Hamiltonian of the electronic system for
Fenna-Matthews-Olson complex (in ${\rm{cm^{-1}}}$) [58]:

$\hat{H}_{el}=\begin{bmatrix}12410&-87.7&5.5&-5.9&6.7&-13.7&-9.9\\\
-87.7&12530&30.8&8.2&0.7&11.8&4.3\\\ 5.5&30.8&12210&-53.5&-2.2&-9.6&6.0\\\
-5.9&8.2&-53.5&12320&-70.7&-17.0&-63.6\\\
6.7&0.7&-2.2&-70.7&12480&81.1&-1.3\\\ -13.7&11.8&-9.6&-17.0&81.1&12630&39.7\\\
-9.9&4.3&6.0&-63.3&-1.3&39.7&12440\end{bmatrix}.$ (14)

The second term is the Hamiltonian of phonon heat baths,

$\hat{H}_{ph}=\sum_{j=1}^{N}\sum_{\alpha}\left(\frac{\hat{p}_{j,\alpha}^{2}}{2}+\frac{\omega_{j,\alpha}^{2}\hat{x}_{j,\alpha}^{2}}{2}\right),$
(15)

where $\hat{p}_{j,\alpha}$, $\hat{x}_{j,\alpha}$, and $\omega_{j,\alpha}$ are
the dimensionless momentum, coordinate, and frequency of the $\alpha$-th
phonon mode of the $j$-th bath. The third term is the exciton-phonon coupling
Hamiltonian,

$\hat{H}_{el-
ph}=-\sum_{j=1}^{N}\hat{V}_{j}\sum_{\alpha}c_{j,\alpha}\hat{x}_{j,\alpha},$
(16)

where $c_{j,\alpha}$ represents the coupling constant between the $j$-th state
and the $\alpha$-th phonon mode. Finally, the last term of Eq. (13) serves as
the counter term,

$\hat{H}_{reorg}=\sum_{j=1}^{N}\lambda_{j}|{j}\rangle\langle{j}|,$ (17)

with $\lambda_{j}$ being the reorganization energy of the $j$-th electronic
state. Under the perturbative and Markovian approximation for the heat baths,
the dissipative dynamics is described by the Lindblad QME, as defined in Eq.
(1). Throughout this work, we choose $\lambda_{j}=35{\rm{cm^{-1}}}$, following
the experimentally determined value at high temperature [22].

The first-order and second-order TCFs for any operator $\hat{X}$ are defined
as

$R^{(1)}(t_{1})=\frac{i}{\hbar}\left\langle\left[\bar{X}(t_{1}),\,\bar{X}(0)\right]\right\rangle,$
(18)
$R^{(2)}(t_{1},t_{2})={\left(\frac{i}{\hbar}\right)}^{2}\left\langle\left[\left[\bar{X}(t_{2}),\,\bar{X}(t_{1})\right],\,\bar{X}(0)\right]\right\rangle,$
(19)

where $[\hat{A},\,\hat{B}]=\hat{A}\hat{B}-\hat{B}\hat{A}$,
$\bar{X}(t)=e^{i\hat{H}_{tot}t/\hbar}\hat{X}e^{-i\hat{H}_{tot}t/\hbar}$ is the
Heisenberg representation of operator $\hat{X}$, and $\langle\cdot\rangle$
denotes the ensemble average over the density matrix of total system. Using
the cyclic invariance, the above equations can be recast as [12]

$R^{(1)}(t_{1})=\frac{i}{\hbar}{\rm{Tr}}\left\\{\hat{X}\hat{G}_{tot}(t_{1})\hat{X}^{\times}\hat{\rho}_{tot}(0)\right\\},$
(20)
$R^{(2)}(t_{1},t_{2})={\left(\frac{i}{\hbar}\right)}^{2}{\rm{Tr}}\left\\{\hat{X}\hat{G}_{tot}(t_{2})\hat{X}^{\times}\hat{G}_{tot}(t_{1})\hat{X}^{\times}\hat{\rho}_{tot}(0)\right\\},$
(21)

where $\hat{\rho}_{tot}(0)$ and
$\hat{G}_{tot}(t)=\exp(-i\hat{H}_{tot}^{\times}t/\hbar)$ are the initial
density operator and the quantum propagator of the total system, and we have
introduced the notation
$\hat{A}^{\times}\hat{B}=\hat{A}\hat{B}-\hat{B}\hat{A}$. To obtain a reduced
description of Eqs. 20 and 21, one traces the bath degrees of freedom, leading
to the formalism with $\hat{\rho}_{tot}(0)$ and $\hat{G}_{tot}(t)$ replaced by
the reduced density operator $\hat{\rho}_{0}$ and the reduced propagator
$\hat{G}(t)$, respectively (see Eq. (1)). The first-order TCF can be evaluated
as follows. Starting from $\hat{\rho}_{0}$, we update the matrix entries by
$\hat{X}^{\times}\hat{\rho}_{0}\to\hat{\rho}^{\prime}_{0}$, and propagate
$\hat{\rho}^{\prime}_{0}$ up to $t_{1}$ using iterative method (Eq. (1)).
$R^{(1)}(t_{1})$ is then obtained by computing the trace of
$\hat{X}\hat{\rho}^{\prime}_{t_{1}}$ for a series of $t_{1}$. To obtain the
second-order TCF, we update $\hat{\rho}^{\prime}_{t_{1}}$ by
$\hat{X}^{\times}\hat{\rho}^{\prime}_{t_{1}}\to\hat{\rho}^{{}^{\prime\prime}}_{t_{1}}$,
propagate up to $t_{2}$ and compute the trace of
$\hat{X}\hat{\rho}^{{}^{\prime\prime}}_{t_{1},t_{2}}$. The overall
computational cost thus increases dramatically with the number of time
variables, which restricts the iteration-based propagation methods (such as
the RK4) to lower-order cases.

###  IV.2 Quantum propagator in the FNO architecture

Figure 4: The architecture of the FNO propagator taken from Ref. [52].
$P_{in}$ and $P_{out}$ are linear projection layers. $W_{1}$ and $W_{2}$ are
two different convolution layers. $\mathrm{FT}$ and $\mathrm{FT}^{-1}$ denote
the Fourier transform and its inverse. $+$ and $\sigma$ represent the element-
wise sum and the activation function. The learnable parameters are those in
$P_{in}$, $P_{out}$, $W_{1}$, and $W_{2}$.

The vanilla FNO architecture is used in this work to parametrize the
propagator ${\bm{G}}_{t}(\theta)$ with a set of trainable parameters $\theta$
for time $t$ up to a chosen maximum time limit $t_{\mathrm{max}}$. The model
architecture is shown in Fig. 4. In Fig. 4 (a), $P_{in}$ and $P_{out}$ denote
the linear projections between the physical and latent spaces, which are
parametrized as linear multilayer perceptrons with one hidden layer and a
rectified linear unit (ReLU) activation function. The structure of the Fourier
layer is shown in Fig. 4 (b), where $\mathrm{FT}$, $\mathrm{FT}^{-1}$, $+$,
and $\sigma$ denote the Fourier transform and its inverse, element-wise sum,
and ReLU, respectively. On the upper route, the linear convolution $W_{1}$ is
applied to the layer-wise input in Fourier space, while on the lower route,
another convolution $W_{2}$ is directly applied without FT. For $W_{1}$, only
the lowest $k_{\mathrm{max}}$ Fourier modes are explicitly included in the
convolution, while all the modes with higher frequencies are truncated for
numerical stability. The results from two routes are summed and activated by
ReLU before passing to the next layer. The learnable parameters $\theta$ are
those in $P_{in}$, $P_{out}$, $W_{1}$, and $W_{2}$. Note that all the
parameters $\theta$ are complex-valued, and ReLU is separately applied to the
real and imaginary parts as $\sigma(a+ib)=\sigma(a)+i\sigma(b)$.

The FNO model takes the initial density matrix $\vec{\rho}_{t_{0}}$ and time
$t_{1}$ as input and outputs $\vec{\rho}_{t_{0}+t_{1}}$ corresponding to the
solution of the QME. The dynamics can be easily extended to arbitrarily long
times by using the composition property of quantum operators. Moreover, no
restrictions are a priori made on the explicit form of the initial states
$\vec{\rho}_{t_{0}}$. We can thus directly employ the FNO propagator to
calculate TCFs by taking ${\bm{X}}_{\times}{\vec{\rho}}_{t_{0}}$ as input and
inferring the dynamics up to $t_{1}$ and $t_{2}$. Each inference requires
running a forward pass of the model, which needs far less computational cost
than conventional iterative methods such as the RK4.

###  IV.3 Training and validation tests

We use 8 layers to construct the neural propagator and parameterize the
projection parts $P_{in}$ and $P_{out}$ as 2-layer feedforward neural networks
with a hidden channel of size 512. The others are all Fourier layers, each of
which has 32 Fourier modes and a hidden channel of size 256. The dataset was
prepared by randomly sampling initial density matrices, which are drawn from
the Gaussian Unitary Ensemble. We prepare in total 400 random samples, 200 of
which are used as the training set, and the rest are held back as the
validation set. The EOM of Eq. 1 is integrated by the RK4 for time up to
$t_{\mathrm{max}}=30{\rm{fs}}$ with a time step of $\delta t=0.6{\rm{fs}}$. We
train the neural propagator with up to $2000$ epochs using Adam optimizer and
loss functions defined in Eqs. (10) and (11). All the computation is carried
out on a single Nvidia 4090 GPU with 24GB memory.

Figure 5: The relative error of the data loss $\mathcal{L}_{\mathrm{data}}$
for the validation set. The horizontal axis represents the index of the
samples contained in the validation set.

We now briefly discuss the on-the-fly sampling algorithm for
$\mathcal{L}_{\mathrm{phys}}$. We first prepare an initial dataset by randomly
sampling an additional 400 density matrices, which are not included in either
training or validation datasets. This dataset can be used directly as
additional inputs for $\mathcal{L}_{\mathrm{phys}}$ in Eq. (11) during the
training process and is re-generated at each epoch. Noticeably, since only the
information at $t=0$ is needed, this approach expands the effective training
set without solving any additional time evolution. Minimization of
$\mathcal{L}_{\mathrm{phys}}$ for this on-the-fly sampled dataset enforces
crucial algebraic properties of the neural propagator (i.e., its derivative
should match ${\bm{L}}$ in Eq. (3) and its value at $t=0$ should give the
identity), enabling it applicable to any initial state.

Finally, we present the validation test by focusing on the data loss
$\mathcal{L}_{\mathrm{data}}$ only. In Fig. 5, we show the relative error of
$\mathcal{L}_{\mathrm{data}}$ for each sample in the validation set. For all
samples, the data loss is always around $\sim 1.0\%$. To further reduce the
validation error, a straightforward solution is to increase the size of the
FNO propagator by using more layers and more hidden channels. We thus conclude
that our constructed neural operator can be regarded as a universal propagator
for any initial state and time $t$ within $[0,t_{\mathrm{max}}]$.

## Data availability

All the datasets in this work were generated directly from the code. They are
also available from the corresponding authors upon reasonable request.

## Code availability

The code is available at https://github.com/MrLightless/Neural-Quantum-
Propagator.

## acknowledgement

J.Z. and L.P.C. acknowledge support from the Key Research Project of Zhejiang
Lab (No. 2021PE0AC02). C.L.B.-R. gratefully acknowledges the European Union’s
Horizon Europe Research and Innovation program under the Marie Skłodowska-
Curie Grant Agreement n°101065295–RDMFTforbosons. Views and opinions expressed
are however those of the authors only and do not necessarily reflect those of
the European Union or the European Research Executive Agency.

## References

  * Breuer and Petruccione [2007] H.-P. Breuer and F. Petruccione,  _The Theory of Open Quantum Systems_ (Oxford University Press, 2007).
  * Weiss [2012] U. Weiss,  _Quantum Dissipative Systems_, 4th ed. (World Scientific, 2012).
  * de Vega and Alonso [2017] I. de Vega and D. Alonso, “Dynamics of non-markovian open quantum systems,” Rev. Mod. Phys. 89, 015001 (2017).
  * Alicki [2003] R. Alicki, “Controlled quantum open systems,” in  _Irreversible Quantum Dynamics_ (Springer, Berlin, Heidelberg, 2003) pp. 121–139.
  * Lindblad [1976] G. Lindblad, “On the generators of quantum dynamical semigroups,” Commun. Math. Phys. 48, 119–130 (1976).
  * Gorini, Kossakowski, and Sudarshan [1976] V. Gorini, A. Kossakowski, and E. C. G. Sudarshan, “Completely positive dynamical semigroups of N‐level systems,” J. Math. Phys. 17, 821–825 (1976).
  * Redfield [1957] A. G. Redfield, “On the theory of relaxation processes,” IBM J. Res. Dev. 1, 19–31 (1957).
  * Ishizaki and Tanimura [2005] A. Ishizaki and Y. Tanimura, “Quantum dynamics of system strongly coupled to low-temperature colored noise bath: Reduced hierarchy equations approach,” J. Phys. Soc. Jpn. 74, 3131–3134 (2005).
  * Tanimura [2020] Y. Tanimura, “Numerically “exact” approach to open quantum dynamics: The hierarchical equations of motion (HEOM),” J. Chem. Phys. 153, 020901 (2020).
  * Makarov and Makri [1994] D. E. Makarov and N. Makri, “Path integrals for dissipative systems by tensor multiplication. condensed phase quantum dynamics for arbitrarily long time,” Chem. Phys. Lett. 221, 482–491 (1994).
  * Makri [2021] N. Makri, “Small matrix decomposition of feynman path amplitudes,” J. Chem. Theory Comput. 17, 3825–3829 (2021).
  * Mukamel [1995] S. Mukamel,  _Principles of Nonlinear Optical Spectroscopy_, Oxford series in optical and imaging sciences (Oxford University Press, 1995).
  * Abramavicius  _et al._ [2009] D. Abramavicius, B. Palmieri, D. V. Voronine, F. Šanda, and S. Mukamel, “Coherent multidimensional optical spectroscopy of excitons in molecular aggregates; quasiparticle versus supermolecule perspectives,” Chem. Rev. 109, 2350–2408 (2009).
  * Gelin, Chen, and Domcke [2022] M. F. Gelin, L. Chen, and W. Domcke, “Equation-of-motion methods for the calculation of femtosecond time-resolved 4-wave-mixing and N-wave-mixing signals,” Chem. Rev. 122, 17339–17396 (2022).
  * Fetherolf and Berkelbach [2017] J. Fetherolf and T. Berkelbach, “Linear and nonlinear spectroscopy from quantum master equations,” J. Chem. Phys. 147, 244109 (2017).
  * Tanimura and Ishizaki [2009] Y. Tanimura and A. Ishizaki, “Modeling, calculating, and analyzing multidimensional vibrational spectroscopies,” Acc. Chem. Res. 42, 1270–1279 (2009).
  * May and Kühn [2011] V. May and O. Kühn,  _Charge and Energy Transfer Dynamics in Molecular Systems_ (Wiley, 2011).
  * Ishizaki and Fleming [2012] A. Ishizaki and G. R. Fleming, “Quantum coherence in photosynthetic light harvesting,” Annu. Rev. Condens. Matter Phys. 3, 333–361 (2012).
  * Butcher [2016] J. C. Butcher,  _Numerical Methods for Ordinary Differential Equations_ (Wiley, 2016).
  * Ishizaki and Fleming [2009a] A. Ishizaki and G. R. Fleming, “On the adequacy of the redfield equation and related approaches to the study of quantum dynamics in electronic energy transfer,” J. Chem. Phys. 130, 234110 (2009a).
  * Zhang, Borrelli, and Tanimura [2020] J. Zhang, R. Borrelli, and Y. Tanimura, “Proton tunneling in a two-dimensional potential energy surface with a non-linear system–bath interaction: Thermal suppression of reaction rate,” J. Chem. Phys. 152, 214114 (2020).
  * Ishizaki and Fleming [2009b] A. Ishizaki and G. R. Fleming, “Theoretical examination of quantum coherence in a photosynthetic system at physiological temperature,” Proc. Natl. Acad. Sci. U.S.A. 106, 17255–17260 (2009b).
  * Zhang, Borrelli, and Tanimura [2021] J. Zhang, R. Borrelli, and Y. Tanimura, “Probing photoinduced proton coupled electron transfer process by means of two-dimensional resonant electronic–vibrational spectroscopy,” J. Chem. Phys. 154, 144104 (2021).
  * Ikeda and Tanimura [2017] T. Ikeda and Y. Tanimura, “Probing photoisomerization processes by means of multi-dimensional electronic spectroscopy: The multi-state quantum hierarchical Fokker-Planck equation approach,” J. Chem. Phys. 147, 014102 (2017).
  * Nazir [2009] A. Nazir, “Correlation-dependent coherent to incoherent transitions in resonant energy transfer dynamics,” Phys. Rev. Lett. 103, 146404 (2009).
  * Kimura and Fujihashi [2014] A. Kimura and Y. Fujihashi, “Quantitative correction of the rate constant in the improved variational master equation for excitation energy transfer,” J. Chem. Phys. 141, 194110 (2014).
  * Mulvihill and Geva [2022] E. Mulvihill and E. Geva, “Simulating the dynamics of electronic observables via reduced-dimensionality generalized quantum master equations,” J. Chem. Phys. 156, 044119 (2022).
  * Schlimgen  _et al._ [2021] A. W. Schlimgen, K. Head-Marsden, L. M. Sager, P. Narang, and D. A. Mazziotti, “Quantum simulation of open quantum systems using a unitary decomposition of operators,” Phys. Rev. Lett. 127, 270503 (2021).
  * Wang  _et al._ [2023] Y. Wang, E. Mulvihill, Z. Hu, N. Lyu, S. Shivpuje, Y. Liu, M. Soley, E. Geva, V. Batista, and S. Kais, “Simulating Open Quantum System Dynamics on NISQ Computers with Generalized Quantum Master Equations,” J. Chem. Theory Comput. 19, 4851 (2023).
  * Liu  _et al._ [2023] W. Liu, Z.-H. Chen, Y. Su, Y. Wang, and W. Dou, “Predicting rate kernels via dynamic mode decomposition,” J. Chem. Phys. 159, 144110 (2023).
  * LeCun, Bengio, and Hinton [2015] Y. LeCun, Y. Bengio, and G. Hinton, “Deep learning,” Nature 521, 436–444 (2015).
  * Hermann  _et al._ [2023] J. Hermann, J. Spencer, K. Choo, A. Mezzacapo, W. M. C. Foulkes, D. Pfau, G. Carleo, and F. Noé, “Ab initio quantum chemistry with neural-network wavefunctions,” Nat. Rev. Chem. 7, 692–709 (2023).
  * Hermann, Schätzle, and Noé [2020] J. Hermann, Z. Schätzle, and F. Noé, “Deep-neural-network solution of the electronic schrödinger equation,” Nat. Chem. 12, 891–897 (2020).
  * Hornik, Stinchcombe, and White [1989] K. Hornik, M. Stinchcombe, and H. White, “Multilayer feedforward networks are universal approximators,” Neural Netw. 2, 359–366 (1989).
  * Sharir  _et al._ [2020] O. Sharir, Y. Levine, N. Wies, G. Carleo, and A. Shashua, “Deep autoregressive models for the efficient variational simulation of many-body quantum systems,” Phys. Rev. Lett. 124, 020503 (2020).
  * Carleo and Troyer [2017] G. Carleo and M. Troyer, “Solving the quantum many-body problem with artificial neural networks,” Science 355, 602–606 (2017).
  * Schätzle  _et al._ [2023] Z. Schätzle, P. B. Szabó, M. Mezera, J. Hermann, and F. Noé, “DeepQMC: An open-source software suite for variational optimization of deep-learning molecular wave functions,” J. Chem. Phys. 159, 094108 (2023).
  * Pfau  _et al._ [2020] D. Pfau, J. S. Spencer, A. G. D. G. Matthews, and W. Foulkes, “Ab-Initio solution of the many-electron schrödinger equation with deep neural networks,” Phys. Rev. Res. 2, 033429 (2020).
  * Cassella  _et al._ [2023] G. Cassella, H. Sutterud, S. Azadi, N. D. Drummond, D. Pfau, J. S. Spencer, and W. M. C. Foulkes, “Discovering quantum phase transitions with fermionic neural networks,” Phys. Rev. Lett. 130, 036401 (2023).
  * Schmitt and Heyl [2020] M. Schmitt and M. Heyl, “Quantum many-body dynamics in two dimensions with artificial neural networks,” Phys. Rev. Lett. 125, 100503 (2020).
  * Reh, Schmitt, and Gärttner [2021] M. Reh, M. Schmitt, and M. Gärttner, “Time-dependent variational principle for open quantum systems with artificial neural networks,” Phys. Rev. Lett. 127, 230501 (2021).
  * Ullah and Dral [2022] A. Ullah and P. O. Dral, “Predicting the future of excitation energy transfer in light-harvesting complex with artificial intelligence-based quantum dynamics,” Nat. Commun. 13, 1930 (2022).
  * Luo  _et al._ [2022] D. Luo, Z. Chen, J. Carrasquilla, and B. K. Clark, “Autoregressive neural network for simulating open quantum systems via a probabilistic formulation,” Phys. Rev. Lett. 128, 090501 (2022).
  * Hartmann and Carleo [2019] M. J. Hartmann and G. Carleo, “Neural-network approach to dissipative quantum many-body dynamics,” Phys. Rev. Lett. 122, 250502 (2019).
  * Vicentini  _et al._ [2019] F. Vicentini, A. Biella, N. Regnault, and C. Ciuti, “Variational neural-network ansatz for steady states in open quantum systems,” Phys. Rev. Lett. 122, 250503 (2019).
  * Nagy and Savona [2019] A. Nagy and V. Savona, “Variational quantum monte carlo method with a neural-network ansatz for open quantum systems,” Phys. Rev. Lett. 122, 250501 (2019).
  * Chen and Chen [1995] T. Chen and H. Chen, “Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its applications to dynamic systems,” IEEE Trans. Neural Netw. 6, 911 – 917 (1995).
  * Kovachki  _et al._ [2023] N. Kovachki, Z. Li, B. Liu, K. Azizzadenesheli, K. Bhattacharya, A. Stuart, and A. Anandkumar, “Neural operator: Learning maps between function spaces with applications to pdes,” J. Mach. Learn. Res. 24, 1–97 (2023).
  * Lu  _et al._ [2022] L. Lu, X. Meng, S. Cai, Z. Mao, S. Goswami, Z. Zhang, and G. E. Karniadakis, “A comprehensive and fair comparison of two neural operators (with practical extensions) based on fair data,” Comput. Methods Appl. Mech. Eng. 393, 114778 (2022).
  * Rosofsky, Majed, and Huerta [2023] S. G. Rosofsky, H. A. Majed, and E. A. Huerta, “Applications of physics informed neural operators,” Mach. Learn.: Sci. Technol. 4, 025022 (2023).
  * Lu  _et al._ [2021] L. Lu, P. Jin, G. Pang, Z. Zhang, and G. E. Karniadakis, “Learning nonlinear operators via deeponet based on the universal approximation theorem of operators,” Nat. Mach. Intell. 3, 218–229 (2021).
  * [52] Z. Li, N. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. Stuart, and A. Anandkumar, “Fourier neural operator for parametric partial differential equations,” Preprint at https://arxiv.org/abs/2010.08895 (2021).
  * Kovachki, Lanthaler, and Mishra [2021] N. Kovachki, S. Lanthaler, and S. Mishra, “On universal approximation and error bounds for fourier neural operators,” J. Mach. Learn. Res. 22, 1–76 (2021).
  * [54] J. Pathak, S. Subramanian, P. Harrington, S. Raja, A. Chattopadhyay, M. Mardani, T. Kurth, D. Hall, Z. Li, K. Azizzadenesheli, P. Hassanzadeh, K. Kashinath, and A. Anandkumar, “FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators,”  10.48550/ARXIV.2202.11214, preprint at https://arxiv.org/abs/2202.11214 (2022).
  * [55] P. Jiang, N. Meinert, H. Jordão, C. Weisser, S. Holgate, A. Lavin, B. Lütjens, D. Newman, H. Wainwright, C. Walker, and P. Barnard, “Digital twin earth – coasts: Developing a fast and physics-informed surrogate model for coastal floods via neural operators,”  10.48550/ARXIV.2110.07100, preprint at https://arxiv.org/abs/2110.07100 (2021).
  * [56] J. Guibas, M. Mardani, Z. Li, A. Tao, A. Anandkumar, and B. Catanzaro, “Adaptive fourier neural operators: Efficient token mixers for transformers,”  10.48550/ARXIV.2111.13587, preprint at https://arxiv.org/abs/2111.13587 (2021).
  * Mizera [2023] S. Mizera, “Scattering with neural operators,” Phys. Rev. D 108, L101701 (2023).
  * Adolphs and Renger [2006] J. Adolphs and T. Renger, “How proteins trigger excitation energy transfer in the fmo complex of green sulfur bacteria,” Biophys. J. 91, 2778–2797 (2006).
  * Duan  _et al._ [2022] H.-G. Duan, A. Jha, L. Chen, V. Tiwari, R. J. Cogdell, K. Ashraf, V. I. Prokhorenko, M. Thorwart, and R. J. D. Miller, “Quantum coherent energy transport in the fenna–matthews–olson complex at low temperature,” Proc. Natl. Acad. Sci. U.S.A. 119, e2212630119 (2022).
  * Kim  _et al._ [2020] Y. Kim, D. Morozov, V. Stadnytskyi, S. Savikhin, and L. V. Slipchenko, “Predictive first-principles modeling of a photosynthetic antenna protein: The fenna–matthews–olson complex,” J. Phys. Chem. Lett. 11, 1636–1643 (2020).
  * Milder  _et al._ [2010] M. T. W. Milder, B. Brüggemann, R. van Grondelle, and J. L. Herek, “Revisiting the optical properties of the fmo protein,” Photosynth. Res. 104, 257–274 (2010).
  * Engel  _et al._ [2007] G. S. Engel, T. R. Calhoun, E. L. Read, T.-K. Ahn, T. Mančal, Y.-C. Cheng, R. E. Blankenship, and G. R. Fleming, “Evidence for wavelike energy transfer through quantum coherence in photosynthetic systems,” Nature 446, 782–786 (2007).
  * Mulvihill  _et al._ [2021] E. Mulvihill, K. Lenn, X. Gao, A. Schubert, B. Dunietz, and E. Geva, “Simulating energy transfer dynamics in the Fenna–Matthews–Olson complex via the modified generalized quantum master equation,” J. Chem. Phys. 154, 204109 (2021).
  * Riva, Romaniello, and Berger [2023] G. Riva, P. Romaniello, and J. A. Berger, “Multichannel dyson equation: Coupling many-body green’s functions,” Phys. Rev. Lett. 131, 216401 (2023).
  * Benavides-Riveros  _et al._ [2022] C. L. Benavides-Riveros, L. Chen, C. Schilling, S. Mantilla, and S. Pittalis, “Excitations of quantum many-body systems via purified ensembles: A unitary-coupled-cluster-based approach,” Phys. Rev. Lett. 129, 066401 (2022).
  * Tuovinen  _et al._ [2023] R. Tuovinen, Y. Pavlyukh, E. Perfetto, and G. Stefanucci, “Time-Linear Quantum Transport Simulations with Correlated Nonequilibrium Green’s Functions,” Phys. Rev. Lett. 130, 246301 (2023).
  * Sangalli  _et al._ [2017] D. Sangalli, J. A. Berger, C. Attaccalite, M. Grüning, and P. Romaniello, “Optical properties of periodic systems within the current-current response framework: Pitfalls and remedies,” Phys. Rev. B 95, 155203 (2017).
  * Yin  _et al._ [2022] J. Yin, Y. Chan, F. H. da Jornada, D. Y. Qiu, S. G. Louie, and C. Yang, “Using dynamic mode decomposition to predict the dynamics of a two-time non-equilibrium green’s function,” J. Comput. Sci. 64, 101843 (2022).