## Non-markovian neural quantum propagator and its application to the simulation of ultrafast nonlinear spectra

Jiaji Zhang and Lipeng Chen Zhejiang Laboratory, Hangzhou 311100, China

a) b)

The accurate solution of dissipative quantum dynamics plays an important role on the simulation of open quantum systems. Here we propose a machine-learning-based universal solver for the hierarchical equations of motion, one of the most widely used approaches which takes into account non-markovian effects and nonperturbative system-environment interactions in a numerically exact manner. We develop a neural quantum propagator model by utilizing the neural network architecture, which avoids time-consuming iterations and can be used to evolve any initial quantum state for arbitrarily long times. To demonstrate the efficacy of our model, we apply it to the simulation of population dynamics and linear and two-dimensional spectra of the Fenna-Matthews-Olson complex.

## I. INTRODUCTION

Ultrafast nonlinear spectroscopy provides a versatile tool to reveal the electronic structure and chemical reaction mechanism. 1-5 Two-dimensional electronic spectroscopy (2DES), in particular, has been widely employed to monitor electronic excitation dynamics in polyatomic molecules. 6-9 By utilizing multiple UV-Vis pulses, one can measure the correlation between different electronic states via off-diagonal peaks of 2DES. In addition, extracting dynamical information from the evolution of 2DES enables the direct visualization of chemical reaction processes. 10-14

learning based surrogate models have been developed to provide universal solvers for PDE. 30-33 In contrast to iterative methods, those surrogate models solve the PDE by defining a functional that describes the mapping between an arbitrary initial condition and its corresponding solution at some subsequent time. This functional is then parameterized as a deep neural network and optimized with a prepared dataset. The state-of-the-art surrogate models, such as Fourier Neural Operator (FNO) and DeepONet, have shown their effectiveness over conventional methods on a set of PDEs of the classical dynamical problems. 34-36

Theoretical simulation of nonlinear spectra is based on the response function theory, which requires the accurate description of system dynamics upon interaction with external laser pulses. 15,16 As molecular systems inevitably interact with their surrounding environment, a commonly used strategy is to treat the environmental degrees of freedom as a heat bath, and derive the equations of motion for the reduced system after tracing out bath degrees of freedom. 17,18 The hierarchical equations of motion (HEOM) is one of the best-known quantum dynamics approaches, which takes into account non-markovian effects and non-perturbative system-environment interactions in a numerically exact manner. 19-21 As a typical partial differential equation (PDE), one recursively solves HEOM by conventional iterative solvers such as fourth-order Runge-Kutta (RK4) and split-operator methods. 22-24 Despite their straightforward implementation, the main drawbacks of iterative methods are the large computational cost and longtime numerical instabilities. While improvements have been proposed to alleviate the numerical issues of iterative methods, an efficient universal solver is still to be proposed. 25-27

Over recent years, the fast development of machine learning technique offers new possibilities to circumvent aforementioned difficulties. 28,29 Avariety of machine

a) Electronic mail: jiaji.zhang@zhejianglab.com

b) Electronic mail: chenlp@zhejianglab.com

In this work, we extend the surrogate models to the non-markovian quantum dynamics by developing a so-called neural quantum propagator (NQP) model for HEOM. As the quantum analogue of the universal PDE solver, the NQP model directly generates dynamics of system without invoking the tedious, expensive iterations. Following our previous work, we adopt the FNO architecture as the core neural network structure. 37 We test its performance by comparing with the conventional RK4 method in various computational scenarios. In addition to the simulation of population dynamics, we also employ the NQP model to compute linear and nonlinear spectra.

Similar to other neural network architectures, training NQP model requires a large amount of high precision data. Those data can only be generated by conventional iterative solvers with a small enough time step, which in turn leads to a large computational cost in the data preparation stage. To address this issue, we introduce a super-resolution algorithm, which only relies on the lowresolution data to construct high-resolution NQP. The intrinsic error in the dataset is systematically improved by utilizing the physics-informed loss function (PILF), which is defined directly from the HEOM. The optimization of PILF does not rely on any prepared dataset, which significantly improves the overall computational performance of the NQP model.

The rest of the paper is organized as follows. In Section II, we introduce the HEOM approach and linear and nonlinear response functions. In Section III, we present the NQP model, including the FNO architecture, the train-

ing setup, and the super-resolution algorithm. Numerical demonstrations on the Fenna-Matthews-Olson (FMO) system are presented in Section IV. Finally, conclusions are drawn in Section V.

## II. METHODOLOGY

## A. HEOM approach

We consider an electronic system interacting with a set of heat baths. The total Hamiltonian can be written as

$$\hat { H } _ { t o t } = \hat { H } _ { s } + \hat { H } _ { b } + \hat { H } _ { s - b }. \quad \quad ( 1 ) \quad \partial _ { t } \hat { \rho }$$

Here, the first term ˆ H s is the Hamiltonian of the electronic system,

̸

$$\hat { H } _ { s } = \sum _ { j = 1 } ^ { N } \varepsilon _ { j } | j \rangle \langle j | + \sum _ { j \neq j ^ { \prime } } \Delta _ { j, j ^ { \prime } } | j \rangle \langle j ^ { \prime } |, \quad \ ( 2 ) \quad \text{$$

where ε j is the energy of the j -th electronic state | j ⟩ , and ∆ j,j ′ is the interstate coupling. The second term is the Hamiltonian of harmonic heat baths,

$$\hat { H } _ { b } = \sum _ { j = 1 } ^ { N } \sum _ { \nu } \left ( \frac { \hat { p } _ { j, \nu } ^ { 2 } } { 2 } + \frac { \omega _ { j, \nu } ^ { 2 } \hat { x } _ { j, \nu } ^ { 2 } } { 2 } \right ), \quad \ \ ( 3 ) \quad \text{open}$$

where ˆ p j,ν , ˆ x j,ν , and ω j,ν are the dimensionless momentum, coordinate, and frequency of ν -th oscillator of j -th bath. The last term is the system-bath interaction Hamiltonian,

$$\hat { H } _ { s - b } = - \sum _ { j = 1 } ^ { N } \hat { V } _ { j } \sum _ { \nu } g _ { j, \nu } \hat { x } _ { j, \nu }, \quad \ \ ( 4 ) \quad \ ^ { \text{me} }$$

where ˆ V j = | j ⟩⟨ j | , and g j,ν is the coupling constant between the j -th state and the ν -th oscillator which can be specified by a spectral density,

$$J _ { j } ( \omega ) = \sum _ { \nu } g _ { j, \nu } ^ { 2 } \delta ( \omega - \omega _ { j, \nu } ). \quad \quad ( 5 ) \quad \text{ive}$$

The influence of the j -th heat bath on the electronic system is characterized by the bath correlation function, 17,18

$$& C _ { j } ( t ) \\ = & \frac { 1 } { \pi } \int _ { 0 } ^ { \infty } \mathrm d \omega J _ { j } ( \omega ) \left [ \coth \left ( \frac { \beta \hbar { \omega } } { 2 } \right ) \cos ( \omega t ) - i \sin ( \omega t ) \right ],$$

where J j ( ω ) is the spectral density of the j -th bath, β = 1 /k B T is the inverse temperature with k B being the Boltzmann constant. We model the bath by the Drude spectral density,

$$J _ { j } ( \omega ) = \frac { 2 \lambda _ { j } \gamma _ { j } \omega } { \gamma _ { j } ^ { 2 } + \omega ^ { 2 } },$$

where λ j is the reorganization energy, and γ j is the inverse of the bath correlation time. In this paper, we consider the high-temperature approximation ( β ℏ γ j &lt; 1), and express Eq. (6) as C j ( ) = t c e j -γ j | t | , where

$$c _ { j } = \frac { 2 \lambda _ { j } } { \beta \hbar { ^ { 2 } } } - i \, \frac { \lambda _ { j } \gamma _ { j } } { \hbar }.$$

To go beyond this approximation, one can include socalled low-temperature correction terms. 38,39 The time evolution of the reduced density matrix can be described by the HEOM approach, which is written as 19,40

$$\i a \, \text{set} \quad \text{by the RETURN approach, winicn is written as} \overset { \cdot \cdot \cdot \cdot \cdot } { \cdot \cdot } \\ \text{en as} \quad \text{(1)} \quad \partial _ { t } \hat { \rho } _ { \vec { n } } ( t ) & = - \left [ \frac { i } { \hbar } \hat { H } _ { s } ^ { \times } + \sum _ { j = 1 } ^ { N } n _ { j } \gamma _ { j } \right ] \hat { \rho } _ { \vec { n } } ( t ) - i \sum _ { j = 1 } ^ { N } \hat { V } _ { j } ^ { \times } \hat { \rho } _ { \vec { n } + \vec { \epsilon } _ { j } } ( t ) \\ \cdot \, \text{e} \, \text{e} \, - \, i \sum _ { j = 1 } ^ { N } \left [ c _ { j } \hat { V } _ { j } \hat { \rho } _ { \vec { n } - \vec { \epsilon } _ { j } } ( t ) - c _ { j } ^ { * } \hat { \rho } _ { \vec { n } - \vec { \epsilon } _ { j } } ( t ) \hat { V } _ { j } \right ] \,, \\ \quad \text{(2)} \quad \text{where} \, \vec { n } = \left \{ n _ { 1 }, n _ { 2 }, \dots, n _ { N } \right \} \text{ denotes the index vector with} \,.$$

where ⃗ n = { n , n 1 2 , ..., n N } denotes the index vector with n j being the non-negative integer, and we have introduced abbreviated notations, ˆ A B × ˆ = ˆ ˆ AB -ˆ ˆ . BA The density operator with all indexes equal to zero, ˆ ( ) ρ ⃗ 0 t with 0 = ⃗ { 0 0 , , ..., 0 } , corresponds to the density operator of the reduced electronic system, while all other density operators are introduced to describe non-markovian and non-perturbative effects.

## B. Linear and nonlinear response functions

The linear and nonlinear spectra are evaluated within the framework of response function theory. 1,16,41 The linear response function is defined as

$$R ^ { ( 1 ) } ( t ) = \left ( \frac { i } { \hbar } \right ) \text{Tr} \left \{ \hat { \mu } \mathcal { G } _ { t o t } ( t ) \hat { \mu } ^ { \times } \hat { \rho } _ { t o t } ( 0 ) \right \}, \quad ( 1 0 )$$

where ˆ is the transition dipole operator, and µ ˆ ρ tot and G tot ( ) t = exp( -iH ˆ × tot / ℏ t ) are the density operator and the quantum propagator of the total system, respectively. The linear absorption spectrum is obtained by the Fourier transformation

$$R ^ { ( 1 ) } ( \omega ) = \text{Im} \int _ { 0 } ^ { \infty } \text{d} t e ^ { i \omega t } R ^ { ( 1 ) } ( t ), \text{ \quad \ \ } ( 1 1 )$$

where Im denotes the imaginary part. The third-order response function is defined as

$$& ( 6 ) \quad \ \ R ^ { ( 3 ) } ( t _ { 3 }, t _ { 2 }, t _ { 1 } ) = \left ( \frac { i } { \hbar } \right ) ^ { 3 } \\ & \text{with},$$

The rephasing and non-rephasing parts of 2D spectrum are defined by

$$R ^ { ( 3, R ) } & ( \omega _ { 3 }, \omega _ { 1 } ; t _ { 2 } ) = \text{Im} \\ & \int _ { 0 } ^ { \infty } \text{d} t _ { 3 } \int _ { 0 } ^ { \infty } \text{d} t _ { 1 } e ^ { i \omega _ { 3 } t _ { 3 } - i \omega _ { 1 } t _ { 1 } } R ^ { ( 3 ) } ( t _ { 3 }, t _ { 2 }, t _ { 1 } ),$$

$$R ^ { ( 3, N R ) } ( \omega _ { 3 }, \omega _ { 1 } ; t _ { 2 } ) & = \text{Im} \\ & \int _ { 0 } ^ { \infty } d t _ { 3 } \int _ { 0 } ^ { \infty } d t _ { 1 } e ^ { i \omega _ { 3 } t _ { 3 } + i \omega _ { 1 } t _ { 1 } } R ^ { ( 3 ) } ( t _ { 3 }, t _ { 2 }, t _ { 1 } ),$$

$$\int _ { 0 } ^ { \infty } d t _ { 3 } \int _ { 0 } ^ { \infty } d t _ { 1 } e ^ { i \omega _ { 3 } t _ { 3 } + i \omega _ { 1 } t _ { 1 } } R ^ { ( 3 ) } ( t _ { 3 }, t _ { 2 }, t _ { 1 } ),$$

Within the HEOM formalism, Eqs. (10) and (12) can be evaluated by replacing ˆ ρ tot and G tot ( ) with ˆ t ρ ⃗ n (0) and Eq. (9), respectively. The final trace is only taken for the zeroth order element of ˆ ρ ⃗ n ( ), i.e., t ˆ ( ). ρ ⃗ 0 t

## III. NEURAL QUANTUM PROPAGATOR

We introduce the abbreviated index, x = ( j, j ′ , n 1 , n 2 , ..., n N ), and align the matrix entries ρ x, t ( ) = ⟨ j ρ | ˆ ( ) ⃗ n t | j ′ ⟩ as the column vector

$$\vec { \rho } _ { t } = \{ \rho ( x _ { 0 }, t ), \rho ( x _ { 1 }, t ), \dots \}. \quad \quad ( 1 5 ) \quad \text{tra}$$

The HEOM (Eq. (9)) can be recast to a matrix-vector form as

$$\partial _ { t } \vec { \rho } _ { t } = L \, \vec { \rho } _ { t }, \quad \quad \ \ ( 1 6 ) \quad \underset { \mathbb { m } { 1 } } { \ln }$$

where the matrix entries of L can be inferred from the right-hand side of Eq. (9). The propagator of HEOM is then defined through the integration form as G t = exp( t L ), which satisfies the composition property,

$$\vec { \rho } _ { t } = G _ { t - t _ { 0 } } \vec { \rho } _ { t _ { 0 } } = e ^ { ( t - t _ { 0 } ) L }. \quad \quad ( 1 7 ) \quad _ {. } \, \mathbb { I }$$

To facilitate the description of later sections, we also introduce uniform time grid as t m = mδ t for m = 1 ∼ N t , where δ t = t max /N t is the time step with N t and t max being the total number of time steps and the fixed upper time limit, respectively.

## A. Model's architecture

The NQP model serves as an extension of FNO to the quantum dynamics. Its mathematical foundation is the extended universal approximation theorem, which states that large enough neural network can approximate the functional, representing the mapping between input and output pairs. 42 To construct the NQP model, we follow our previous work 37 and parameterize the HEOM propagator as a deep neural network, G t m [ θ ], where θ represents all the trainable parameters. The architecture of the NQP model is shown in Fig. 1. In Fig. 1(a), P in and P out are the linear projections between physical and latent Fourier spaces. They are parameterized as the point-wise convolution network with one hidden layer and a Gaussian Error Linear Unit (GeLU) activation function. The rest parts are called the Fourier layers with their structure presented in Fig. 1(b).

To process the input of the l -th layer ⃗ v l , two different routes are adopted. On the upper route, F and F -1 denote the Fourier and its inverse transform. The point-wise convolution W l serves as the learnable weight

Figure 1. The architecture of (a) the NQP model, and (b) the l -th Fourier layer. Here, F and F -1 denote the Fourier transform and its inverse. + and σ represent the elementwise sum and the GeLU activation function. The learnable parameters are those in P in , P out , and W l .

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00222v2-figure-061db134-8b52-4ec0-b7da-84c66b47773d.jpg)

in Fourier space. Only the lowest k max modes are explicitly included in the weight tensor, while others with higher frequencies are truncated to control the size of the model and avoid the numerical instabilities. The lower route is similar to the residual network. The results of two different routes are summed and activated by GeLU before passing to the next layer.

The NQP model takes all the entries of initial condition ⃗ ρ 0 and a chosen time t as the input, and outputs ⃗ ρ t = G t [ θ ⃗ ] ρ 0 satisfying Eq. (16). It should be noted that no restrictions are a priori made on the explicit forms of ⃗ ρ 0 . The NQP model can be directly applied to the simulation of response function by taking the field interaction form ˆ µ × ˆ ρ ⃗ n as the input. Since the composition property is also retained during the parameterization, the time evolution up to arbitrarily long times can be obtained by recursively applying Eq. (17).

## B. Training objective

The NQP model is trained by minimizing an objective function L , defined as

$$\mathcal { L } = \alpha \mathcal { L } _ { d a t a } + ( 1 - \alpha ) \mathcal { L } _ { p h y s }, \text{ \quad \ \ } ( 1 8 )$$

where L data and L phys are referred to as the data and physics-informed loss functions, respectively. The hyperparameter α ∈ (0 , 1) serves as a weight factor, which will be dynamically adjusted in the training stage.

For the data part L data , we prepare a dataset by randomly sampling a set of initial condition { ⃗ ρ 0 } , and then evaluating their time evolution { ⃗ ρ t } up to t ∈ [0 , t max ] using conventional RK4 method. The data loss function is defined as follows,

$$\mathcal { L } _ { d a t a } = \sum _ { p = 1 } ^ { N _ { d a t a } } \sum _ { m = 1 } ^ { N _ { t } } \frac { \left \| G _ { t _ { m } } [ \theta ] \vec { \rho } _ { 0 } ^ { ( p ) } - \vec { \rho } _ { t _ { m } } ^ { ( p ) } \right \| _ { F } } { \left \| \vec { \rho } _ { t _ { m } } ^ { ( p ) } \right \| _ { F } }, \quad ( 1 9 )$$

where || · || F denotes the Frobenius-norm, N data is the number of individual samples in the dataset, and ⃗ ρ ( p ) 0 and ⃗ ρ ( p ) t m are the initial condition and the corresponding evolution for the p -th sample, respectively.

To ensure the universality of G t [ θ ] that is applicable to any ⃗ ρ 0 , one needs a large number of samples N data , which leads to even more computational cost in the data preparation stage. We introduce a physics-informed loss function to reduce the effective number of samples N data while keeping the universality of G t [ θ ]. 43 The physicsinformed loss function is defined by minimizing the difference between leftand right-hand sides of Eq. (16) as

$$\mathcal { L } _ { p h y s } = \sum _ { p ^ { \prime } = 1 } ^ { N _ { p h y s } } \sum _ { m = 1 } ^ { N _ { t } } \left | \left | \partial _ { t } G _ { t _ { m } } [ \theta ] \bar { \rho } _ { 0 } ^ { ( p ^ { \prime } ) } - L G _ { t _ { m } } [ \theta ] \bar { \rho } _ { 0 } ^ { ( p ^ { \prime } ) } \right | \right | _ { F } ^ { \text{ reg} } \quad \text{$\text{neg}$} \\ \text{where $N$}, \quad \text{is the number of samnices in the nthvisc} \quad \text{absences}.$$

where N phys is the number of samples in the physics dataset. The time derivative ∂ ⃗ t ρ t is evaluated by the finite difference method. It should be mentioned that the calculation of L phys involves far less samples as compared to that of L data . In addition, we adopt the onthe-fly sampling approach by re-generating the physics dataset at each training epoch to further improve the performance of the trained model.

## C. Super resolution algorithm

To further reduce the computational cost in the data preparation stage, we introduce a super resolution algorithm, which allows the construction of the high resolution NQP model from a lower resolution dataset. As illustrated from the previous subsection, the lower resolution dataset is prepared by integrating the HEOM with a larger time step Kδ t ( K &gt; 1) for a set of { ⃗ ρ 0 } using the RK4 method. The obtained data is then embedded into the finer grid { t m = mδ t } by interpolating the missing value using the linear interpolation scheme. L data is evaluated on this interpolated dataset in the training stage.

On the other hand, the physics-informed loss function L phys is evaluated directly on the finer time grid and serves as the correction over the deviation from the dataset. The super resolution algorithm is then completed by dynamically adjusting the weight factor α in Eq. (18) during the training process. At the begining, we set α = 1 and randomly initilize all the model's parameters. During the training process, α is gradually decreased to a small enough value such as ∼ 0 01, and . L phys gradually becomes the dominant contribution term. The minimization of L phys allows the improvement of the resolution over the intrinsic deviation of the dataset.

At the end of this subsection, we briefly discuss the possibility of data-free training, which is achieved by fixing α = 0 and using only L phys during the training process. From a theoretical point of view, training with or without L data results in the same model as long as L phys becomes the dominant contribution of Eq. (18). In practice, however, training with only L phys requires longer epochs for convergence when all the learnable parameters are randomly initialized. In this case, a prepared dataset, even with low resolution, serves as a well-performed guidance for the training.

## IV. NUMERICAL EXPERIMENTS

In the following, we use seven-sites FMO-complex (apo-FMO) as model system and adopt the Adolph and Renger's Hamiltonian. 44,45 We restrict the discussion to the one-exciton manifold. The electronic state | j ⟩ ( j = 1 , · · · , 7) corresponds to the state where only j -th pigment is excited, and | g ⟩ denotes the ground state without any excitation. The system-bath interaction is chosen as ˆ V j = | j ⟩⟨ j | and ˆ V g ≡ 0. The heat bath parameters are chosen as λ j = 35cm -1 , γ j = 200cm -1 , and T = 300K, respectively. The HEOM is truncated at the hierarchy level of ∑ n j ≤ 2 after adopting the filtering algorithm. 46 Within our choice of parameters and under the high-temperature limit, we found that it is accurate enough for the testing of our NQP model. We set the upper time limit as t max = 30fs with a time step of δ t = 0 6 fs, which results in . N t = 50 time points.

## A. Training and validation test

We first introduce the model's hyper-parameters and training setup. In order to train the NQP model, we prepare the low-resolution training dataset by randomly sampling N data = 3000 initial conditions ⃗ ρ 0 . The low resolution dataset is prepared by integrating HEOM with a larger time step of 3 δ t . The missing values are linearly interpolated when embedded into the finer grid with the time step of δ t . To test the accuracy of the model, we also prepare a high-resolution validation set with 500 samples, following the same setup but using a smaller time step of δ t . It should be noted that the high-resolution validation set is never referred in the training stage. In the training process, the physics dataset is prepared using the on-thefly sampling algorithm by randomly generating N phys = 2000 initial conditions at each epoch.

The other hyper-parameters of the NQP model are chosen as follows. We set the hidden channel of projections P in and P out as 512. We use 4 Fourier layers, each of which has a hidden channel of size 64, and the total number of trainable parameters is around 10 million. The model is trained for 10 5 epochs using the Adam optimizer. The learning rate is initially set to 10 -4 , and then halved every 500 epochs until reaching 10 -6 . The weight factor α in Eq. (18) is initialized as α = 1, and halved every 100 epochs until reaching ∼ 10 -2 . All the tasks are performed on the Nvidia A40 GPU with 48 GB memory.

Figure 2. The relative error of L data for the validation set. The horizontal axis represents the index of the samples, while the vertical axis is the relative error.

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00222v2-figure-e296e28d-150c-43c5-8761-71effa90bf02.jpg)

To test our model, we present the validation test by showing the relative error of L data for each sample in the validation set in Fig. 2. For all samples, the relative error is around 0 5%. This error can be further reduced by . using more samples in the data and physics sets, extending the training to longer epochs, and increasing the size of the NQP model. It should be pointed out that this error corresponds to the overall deviation of all the entries of ˆ , including those deep hierarchy elements that have ρ ⃗ n much smaller magnitude as compared to ˆ . ρ ⃗ 0

## B. Population dynamics

̸

By using the composition property, ⃗ ρ t 1 + t 2 = G t 2 [ θ ⃗ ] ρ t 1 , our NQP model can infer truly long-time dynamics well beyond the training time limit t max . To test the accuracy of the long-time dynamics predicted by the NQP model, we compute population dynamics up to 40 t max ( ∼ 1 2 ps). . The reference results are obtained from the RK4 method with an integration time step of δ t . Here, we consider two initial conditions: (a) ˆ (0) = ρ ⃗ 0 | 1 ⟩⟨ 1 , and | ( b ) ˆ (0) = ρ ⃗ 0 | 6 ⟩⟨ 6 , which correspond to the excitation lo-| calized at the first and sixth pigment, respectively. All other hierarchy elements ˆ ρ ⃗ n (0) ( ⃗ n = 0) are set to zero for ⃗ the factorized bath initial condition.

In Fig. 3, we show the time evolution of populations p n ( ) = t ⟨ n ρ | ˆ ( ) ⃗ 0 t | n ⟩ for sites (a) n = 1, 2, and 3, and (b) n = 4, 5, and 6, respectively, following the experimentally demonstrated energy transfer pathways. In both cases, our NQP model yields results in perfect agreement with those from the reference RK4 method up to 10 t max . While model-predicted long time dynamics deviates slightly from the exact results due to the accumulation of errors in the training stage, our NQP model still infers the accurate dynamics even far beyond the training time (40 t max ).

Figure 3. Population dynamics computed using NQP model (solid lines) and reference RK4 method (dashed lines) for two typical initial conditions up to t = 1 2 ps (40 . t max ).

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00222v2-figure-0c5ba0f2-a1b8-4c56-bc89-61e30856b4e6.jpg)

Figure 4. The linear spectrum R (1) ( ω ) evaluated from the NQP model (blue line) and the RK4 method (red line), respectively.

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00222v2-figure-93ef2876-dcdb-416d-942e-17ee96ad236c.jpg)

## C. Linear spectra

Next, we apply our NQP model to simulate the linear and third-order response functions as defined in Eqs. (10) and (12). We choose the transition dipole operator as

$$\hat { \mu } = \sum _ { j = 1 } ^ { 7 } \mu _ { j } \left ( | j \rangle \langle g | + | g \rangle \langle j | \right ), \text{ \quad \ \ } \left ( 2 1 \right )$$

where µ j is the transition dipole moment of j -th pigment. The system is initially in the electronic ground state before the photoexcitation, i.e., ˆ (0) = ρ ⃗ 0 | g ⟩⟨ g | .

In Fig. 4, we show the linear spectrum evaluated from the NQP model and the RK4 method. For each case, we set the time window of t to [0 , 40 t max ], and normalize the peak intensities with respect to their own maximum value. Overall, the NQP model yields spectrum in good

Figure 5. 2D spectra at t 2 = 0fs evaluated from (a,b) the NQP model and (c,d) the RK4 reference, respectively. 'R' in (a,c) and 'NR' in (b,d) denote the rephasing and nonrephasing part, respectively.

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00222v2-figure-c9d0deca-e532-4f57-ad44-5695a9c2cc13.jpg)

agreement with that from the reference RK4. The small deviations of some peak intensities may be attributed to the model's architecture. The adaption of Fourier transform in the model's architecture generates some artificial aliasing modes, the magnitudes of which are increased after recurrent evaluation of long time dynamics. This systematic error could be resolved by carefully finetuning the truncation level of Fourier modes in NQP model, or by replacing the Fourier transform with methods such as wavelet transform or spatial convolutions.

## D. Two-dimensional spectra

We further apply the NQP model to compute 2D spectra at different time t 2 . Since we restrict the discussion to the one-exciton manifold, the two-exciton states, which contribute to the excited-state absorption of 2D spectra, are not included in the Hamiltonian. The simulated 2D spectra thus only involve the contribution of groundstate bleach and stimulated emission. The inclusion of two-exciton states can be easily established by adding the electronic states | jj ′ ⟩ , which represent the simultaneous excitation at two pigments, to the Hamiltonian, and increasing the length of column vector ⃗ ρ t defined in Eq. (15). The NQP model can handle this expanded system by using more layers and increasing the size of each layer. However, this will require more learnable parameters and lead to a dramatic increase of computational cost especially in the training stage, which is outside the scope of the present work.

In Fig. 5, we show the rephasing (a, c) and nonrephasing (b, d) parts of 2D spectra at t 2 = 0 fs evaluated from the NQP model (a, b) and the RK4 reference (c, d),

Figure 6. The rephasing (a, c) and non-rephasing (b, d) parts of 2D spectra at t 2 = 50fs evaluated from (a, b) the NQP model and (c, d) the RK4 reference, respectively.

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00222v2-figure-8698d83b-30f5-4ee9-8f23-c3968f3844aa.jpg)

Figure 7. The rephasing (a, c) and non-rephasing (b, d) parts of 2D spectra at t 2 = 100fs evaluated from (a, b) the NQP model and (c, d) the RK4 reference, respectively.

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00222v2-figure-3618a053-54e7-449c-9d87-d4045c541c69.jpg)

respectively. We set the time window of both t 1 and t 3 to [0 , 40 t max ]. For better illustration, we adopt the widely used arcsinh scaling after the normalization of peak intensities. 47-49 While the simulation of 2D spectra at t 2 = 0 fs requires the propagation up to t 1 + t 3 ≤ 80 t max , far beyond the training time limit, both rephasing and non-rephasing parts of 2D spectra predicted by the NQP model are again in good agreement with those from the RK4 reference, demonstrating the long-time reliability of our NQP model. The rephasing and non-rephasing parts of 2D spectra at t 2 = 50 100 500 fs evaluated from the , , NQPmodel and the RK4 reference are shown in Figs. 6-8. The NQP model again yields accurate 2D spectra, illustrating the energy relaxation process from higher exciton states to lower exciton states as t 2 increases.

Figure 8. The rephasing (a, c) and non-rephasing (b, d) parts of 2D spectra at t 2 = 500fs evaluated from (a, b) the NQP model and (c, d) the RK4 reference, respectively.

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00222v2-figure-eca2c25c-59b9-422c-a9b0-41d13a9154ce.jpg)

Figure 9. The relative error between the NQP model and the RK4 reference as a function of t 2 .

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00222v2-figure-ba8b97ea-d12c-4dfc-ba50-7bab93d7e5c8.jpg)

To quantify the difference between model-predicted and reference spectra more specifically, we define the averaged point-wise relative deviation as

$$\Delta ( t _ { 2 } ) & = \int _ { \Omega _ { 3 } } \mathrm d \omega _ { 3 } \int _ { \Omega _ { 1 } } \mathrm d \omega _ { 1 } \, \left | 1 - \frac { R _ { N Q P } ^ { ( 3 ) } ( \omega _ { 3 }, \omega _ { 1 } ; t _ { 2 } ) + \varepsilon } { R _ { R K 4 } ^ { ( 3 ) } ( \omega _ { 3 }, \omega _ { 1 } ; t _ { 2 } ) + \varepsilon } \right |, \\, \quad \circ \, \circ \, \quad - \, \nu \, \cos \infty \, \quad - 1 \, \nu \, \cos \infty \, \quad - 1 \, \nu \, ( 3 ) \, \quad ( 2 2 ) \, \quad J i a _ { \text{lead} } \, \right \rangle$$

where Ω 1 , Ω 3 ∈ (12000cm -1 , 12900cm -1 ), R (3) NQP and R (3) RK 4 are the spectra evaluated from the NQP model and the RK4 reference, respectively, and ε ∼ 10 -6 is employed to avoid the numerical divergence. In Fig. 9, we display ∆( t 2 ) up to t 2 = 500fs with a time step of 20 fs. With the increase of t 2 , the relative error also rapidly increases due to the fast relaxation of the system. After the system reaches the equilibrium, the error also reaches the plateau at around 4% ∼ 5%. This again demonstrates the long-time reliability of our NQP model.

## V. CONCLUSION

In this work, we develop a NQP model for the HEOM approach to treat the non-markovian dynamics. We use the FNO as the model's architecture, and design a superresolution algorithm to reduce the computational cost in the data preparation stage. In the training stage, we employ both data loss function and an extra PILF to improve the numerical performance. The accuracy of the NQP model is tested by computing the population dynamics and linear and two-dimensional spectra. The NQP model yields results in good agreement with the conventional RK4 method, demonstrating its potential applicability in various computational scenarios.

In the current NQP model, the number of learnable parameters scales exponentially with both the size of the reduced system and the number of hierarchy elements. To alleviate the computational cost, future work may extend the NQP model to deal with the decomposed form of ˆ ( ). ρ ⃗ n t For example, ˆ ( ) can be expressed as a matrixρ ⃗ n t product-state form by using the twin-space formulation and tensor-train decomposition, which requires far less entries than the original density matrix. 50,51 In addition, our focus here is on the time-independent Hamiltonian. Future development may extend to the driven dynamics, where the system Hamiltonian contains time-dependent external fields. A typical application is the spectroscopic equations of motion approach which was employed to the simulation of strong-field nonlinear spectra. 52,53 Work in these directions is in progress.

## ACKNOWLEDGMENTS

J.Z. and L.P.C. acknowledge support from the starting grant of research center of new materials computing of Zhejiang Lab (No. 3700-32601).

## AUTHOR DECLARATIONS

## Author Contributions

Jiaji Zhang : Data curation (lead); Formal analysis (lead); Investigation (equal); Supervision (equal).

Lipeng Chen : Conceptualization (lead); Funding acquisition (lead); Investigation (equal); Supervision (equal).

## Conflict of Interest

The authors have no conflicts to disclose.

## DATA AVAILABILITY

The data that support the findings of this study are available from the corresponding author upon reasonable request.

## CODE AVAILABILITY

The code that support the findings of this study are available from the corresponding author upon reasonable request.

## REFERENCES

- 1 M. F. Gelin, L. Chen, and W. Domcke, Chem. Rev. 122 , 1733917396 (2022).
- 2 M. Nisoli, P. Decleva, F. Calegari, A. Palacios, and F. Mart´ ın, Chem. Rev. 117 , 10760-10825 (2017).
- 3 S. Mukamel, Annu. Rev. Phys. Chem. 51 , 691-729 (2000).
- 4 M. Maiuri, M. Garavelli, and G. Cerullo, J. Am. Chem. Soc. 142 , 3-15 (2019).
- 5 K. E. Dorfman, F. Schlawin, and S. Mukamel, Rev. Mod. Phys. 88 , 045008 (2016).
- 6 E. Fresch, F. V. A. Camargo, Q. Shen, C. C. Bellora, T. Pullerits, G. S. Engel, G. Cerullo, and E. Collini, Nat. Rev. Methods Primers 3 (2023).
- 7 T. A. A. Oliver, R. Soc. Open Sci. 5 , 171425 (2018).
- 8 G. S. Schlau-Cohen, A. Ishizaki, and G. R. Fleming, Chem. Phys. 386 , 1-22 (2011).
- 9 N. S. Ginsberg, Y.-C. Cheng, and G. R. Fleming, Acc. Chem. Res. 42 , 1352-1363 (2009).
- 10 G. D. Scholes, G. R. Fleming, A. Olaya-Castro, and R. van Grondelle, Nat. Chem. 3 , 763-774 (2011).
- 11 M. Kullmann, S. Ruetzel, J. Buback, P. Nuernberger, and T. Brixner, J. Am. Chem. Soc. 133 , 13074-13080 (2011).
- 12 E. A. Arsenault, P. Bhattacharyya, Y. Yoneda, and G. R. Fleming, J. Chem. Phys. 155 (2021).
- 13 J. Kim, J. Jeon, T. H. Yoon, and M. Cho, Nat. Comm. 11 (2020). 14 S. Ruetzel, M. Kullmann, J. Buback, P. Nuernberger, and T. Brixner, Phys. Rev. Lett 110 (2013).
- 15 M. Cho, Coherent Multidimensional Spectroscopy (Springer Singapore, 2019).
- 16 S. Mukamel, Principles of Nonlinear Optical Spectroscopy (Oxford University Press, 1995).
- 17 H.-P. Breuer and F. Petruccione, The Theory of Open Quantum Systems (Oxford University Press, 2007).
- 18 U. Weiss, Quantum Dissipative Systems , 4th ed. (World Scientific, 2012).
- 19 Y. Tanimura, J. Chem. Phys. 153 , 020901 (2020).
- 20 L. Ye, X. Wang, D. Hou, R. Xu, X. Zheng, and Y. Yan, WIREs Comput. Mol. Sci. 6 , 608-638 (2016).
- 21 J. Zhang and Y. Tanimura, J. Chem. Phys. 156 (2022).
- 22 P. E. Kloeden and E. Platen, Numerical Solution of Stochastic Differential Equations (Springer Berlin Heidelberg, 1992).
- 23 Y. Yan, M. Xu, T. Li, and Q. Shi, J. Chem. Phys. 154 , 194104 (2021).
- 24 Y. Ke, J. Chem. Phys. 158 (2023).
- 25 A. Kimura and Y. Fujihashi, J. Chem. Phys. 141 , 194110 (2014). 26 A. W. Schlimgen, K. Head-Marsden, L. M. Sager, P. Narang, and D. A. Mazziotti, Phys. Rev. Lett 127 , 270503 (2021).
- 27 W. Liu, Z.-H. Chen, Y. Su, Y. Wang, and W. Dou, J. Chem. Phys. 159 , 144110 (2023).
- 28 Y. LeCun, Y. Bengio, and G. Hinton, Nature 521 , 436-444 (2015).
- 29 J. Hermann, J. Spencer, K. Choo, A. Mezzacapo, W. M. C. Foulkes, D. Pfau, G. Carleo, and F. No´ e, Nat. Rev. Chem. 7 , 692-709 (2023).
- 30 L. Lu, P. Jin, G. Pang, Z. Zhang, and G. E. Karniadakis, Nat. Mach. Intell. 3 , 218-229 (2021).
- 31 Z. Li, N. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. Stuart, and A. Anandkumar, Eprint arXiv: 2010.08895 (2021).
- 32 N. Kovachki, Z. Li, B. Liu, K. Azizzadenesheli, K. Bhattacharya, A. Stuart, and A. Anandkumar, J. Mach. Learn. Res. 24 , 1-97 (2023).
- 33 J. Guibas, M. Mardani, Z. Li, A. Tao, A. Anandkumar, and B. Catanzaro, Eprint arXiv: 2111.13587 (2021).
- 34 L. Lu, X. Meng, Z. Mao, and G. E. Karniadakis, SIAM Rev. 63 , 208-228 (2021).
- 35 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, Eprint arXiv: 2202.11214 (2022).
- 36 P. Jiang, N. Meinert, H. Jord˜o, a C. Weisser, S. Holgate, A. Lavin, B. L¨tjens, D. Newman, H. Wainwright, C. Walker, u and P. Barnard, Eprint arXiv: 2110.07100 (2021).
- 37 J. Zhang, C. L. Benavides-Riveros, and L. Chen, J. Phys. Chem. Lett. 15 , 3603-3610 (2024).
- 38 A. Ishizaki and Y. Tanimura, J. Phys. Soc. Jpn. 74 , 3131-3134 (2005).
- 39 J. Hu, R.-X. Xu, and Y. Yan, J. Chem. Phys. 133 , 101106 (2010). 40 Y. Tanimura, J. Phys. Soc. Jpn. 75 , 082001 (2006).
- 41
- J. Zhang and Y. Tanimura, J. Chem. Phys. 159 , 014102 (2023). 42 N. Kovachki, S. Lanthaler, and S. Mishra, J. Mach. Learn. Res. 22 , 1-76 (2021).
- 43 S. G. Rosofsky, H. A. Majed, and E. A. Huerta, Mach. learn.: Sci. Technol. 4 , 025022 (2023).
- 44 J. Adolphs and T. Renger, Biophys. J. 91 , 2778-2797 (2006).
- 45 A. Ishizaki and G. R. Fleming, Proc. Natl. Acad. Sci. 106 , 1725517260 (2009).
- 46 Q. Shi, L. Chen, G. Nan, R.-X. Xu, and Y. Yan, J. Chem. Phys. 130 , 084105 (2009).
- 47 S.-H. Yeh and S. Kais, J. Chem. Phys. 141 (2014).
- 48 L. Chen, R. Zheng, Y. Jing, and Q. Shi, J. Chem. Phys. 134 (2011).
- 49 M. Cho, H. M. Vaswani, T. Brixner, J. Stenger, and G. R. Fleming, J. Phys. Chem. B 109 , 10542-10556 (2005).
- 50 R. Borrelli, J. Chem. Phys. 150 , 234102 (2019).
- 51 R. Borrelli and M. F. Gelin, WIREs Comput. Mol. Sci. 11 , e1539 (2021).
- 52 H. Wang and M. Thoss, Chem. Phys. Lett. 389 , 43-50 (2004).
- 53 H. Wang and M. Thoss, Chem. Phys. 347 , 139-151 (2008).
Jiaji Zhang Lipeng Chen
Non-markovian neural quantum propagator and its application to the simulation of ultrafast nonlinear spectra
The accurate solution of dissipative quantum dynamics plays an important role
on the simulation of open quantum systems. Here we propose a
machine-learning-based universal solver for the hierarchical equations of
motion, one of the most widely used approaches which takes into account
non-markovian effects and nonperturbative system-environment interactions in a
numerically exact manner. We develop a neural quantum propagator model by
utilizing the neural network architecture, which avoids time-consuming
iterations and can be used to evolve any initial quantum state for arbitrarily
long times. To demonstrate the efficacy of our model, we apply it to the
simulation of population dynamics and linear and two-dimensional spectra of the
Fenna-Matthews-Olson complex.