Emulating Quantum Dynamics with Neural Networks via Knowledge Distillation

In Figure 1, we show our proposed training framework for preparing machine learning-based quantum dynamics emulators. The main idea is based on the concepts of knowledge distillation ²⁴ and curriculum learning ²³. However, instead of extracting information from a larger machine learning model, our target is a simple, physically informed simulator. In detail, the framework consists of the following steps. First, the simulator samples the initial conditions and generates time-trajectories describing an evolution of the physical system of interest (cf. Fig. 1a). Next, we construct training examples from these recorded simulations (cf. Fig. 1b). We select a diverse set of the training examples, making sure that all important phenomena of interest are represented (cf. Fig. 1c). This balanced curriculum of examples is consequently used as the input for the machine learning model. We train the model, and then we validate it using some novel examples. When testing the model, we include cases that were not directly represented during the training, to measure whether the model can combine and generalize the acquired knowledge (cf. Fig. 1d).

We illustrate our approach by training a neural network based emulator to predict the quantum dynamics of a one-dimensional system. To demonstrate that the emulator can extract knowledge from simple examples and generalize it in a non-trivial way, we restrict the physical simulator to cases with only a single rectangular potential barrier. The emulator need to learn from these simple examples the basic properties of the wave function propagation, namely dispersion, scattering, tunneling, and quantum wave interference. Next, we test whether the emulator can predict the time-evolution of wave packages in a more general case: e.g., for packages of different shapes propagating through an arbitrarily complex potential landscape.

In detail, a single training simulation depicts a propagation of a Gaussian modulated wave packet. The initial conditions consists of: the center-of-mass position ($X_{0}$), the spread ($S_{0}$), and the energy ($E_{0}$) of the packet. Due to the translational symmetry of the problem, we can assume that the rectangular barrier is located at the center of the system, without any loss of generality. Consequently, the environment is fully described by two numbers: the height ($H_{b}$), and the width ($W_{b}$) of the rectangular barrier (cf. Fig. 1a, again; see also more details in the Supplementary Information).

The architecture of our machine learning emulator is depicted in Figure 2a. We start by stacking four time steps of the original simulation data. To make the input independent of the size of the emulated system, we use windows of fixed width (represented by red rectangles in Fig. 2b). Each windowed chunk of data contains local information about the real and the imaginary part of the wave function as well as the local information about the potential landscape. Altogether, the data can be represented in the form of three channels (similar to the RGB channels when representing visual data), as depicted in on the left side of Fig. 2a. Next, we feed the input data to our neural network. The first few layers transform each time step into a hidden vector (represented graphically in our diagram by the red array). The role of the dense layer in our architecture is to allow the network to mix the information between different spatial points (non-local operation) and between channels (both local and non-local operations). In the next step, the hidden vectors pass through a set of gated recurrent units (GRU) ⁴⁴. These units are responsible for extracting time-dependent information from the data. In the last step we reconstruct the wave function (the real and the imaginary part, represented by two channels, as depicted on the right side of Fig. 2a).

We generated a training data set with combinations of 189 different sets of initial Gaussian wave packets and 14 different rectangular potential barriers (in total, 2646 configurations). We kept the widths of the barrier fixed at $W_{b}=7\,a.u.$ (71 pixels), the size of the environment at $N_{x}=1024$ pixels (with periodic boundary conditions enabled), and the window width at $W=23$ pixels.

We trained the emulator for five epochs using the AdamW⁴⁵ optimizer with the MSE loss functions as the training objective. For more details regarding the architecture of our neural network-based emulator, the training procedure, and the composition of our training sets, see the Methods section.

We present the results of the emulation and the comparison with the ground truth in Figs. 2d–g. To show, that the emulator can generalize to out-of-domain situations, we have chosen two challenging shapes for the potential barrier: a step pyramid and a smooth half-circle. Since the emulator was trained only on a single rectangular barrier of a fixed width (much wider than the width of the step in the pyramid), those testing cases can not be reduced to any examples that were seen during the training. Instead, to make valid predictions, the neural networks must recombine the acquired knowledge in a non-trivial way.

As we can see in Figs. 2d–g, in both cases the predictions of our emulator match well the ground truth. It is worth of noticing, that the emulator makes its predictions in a recurrent manner (by using the predictions of the previous steps as input of the next step). Therefore, it is expected that predicting the evolution over hundreds of steps will cause some error accumulation. The fact that even after $350$ steps the error is negligible, speaks about the quality of the individual (step by step) predictions. This result indicates also, that long-term predictions of the dynamics are possible in our framework.

As a conclusion, the proposed neural network emulator successfully learns the classical aspects of the wave dynamics, such as dispersion and interference. It also captures the more complex quantum phenomena, such as quantum tunneling. To further show that the emulator can generalize the acquired knowledge to make both in- and out-of-distribution predictions, we hand-designed a test data set with 12 freely dispersing cases, 11 rectangular barriers (with randomly chosen width and height), as well as 14 more challenging test cases depicting both multiple and irregularly shaped barriers (in total, 37 test instances). In all cases, the results were satisfactory, confirming the ability of the emulator to generalize to novel (and notably, more challenging) situations (see the detailed results in Figs. S3 and S4 in the Supplementary Information).

In this section, we aim to provide some justification for the specific architecture of the emulator, that we introduced in the previous sections. In order to verify, whether all the introduced complexity is necessary, we ask a question whether there exists a simpler machine learning model capable of learning quantum dynamics with a similar accuracy. Consequently, we compare our recurrent-network based approach with other popular network architectures. For consistency, in each case we use the same window-based scheme (to keep the input constant), while emulating the wave packet propagation for 400 time steps. When comparing with the ground truth, we measured the average mean absolute error (MAE) and the average normalized correlation $\mathcal{C}$.

The results are presented in Table 1, where the values represent the average performance over all our 37 test cases. For the comparison, we used three other models: a linear model, a densely connected feedforward model, and a convolutional model (for a detail description of each test case and each network architecture, see the Method section). As evident in our results, the proposed architecture (utilizing the gated recurrent units, GRU) outperforms all other, simpler architectures by a large margin.

We present a details comparison for one of our test set in Figure 3. It is evident that the recurrent architecture provides the best results, whereas the simpler architectures fail to capture the complex long-time evolution of the wave packets. Notably, a failure of each simpler model can be used to justify different aspect of our final design. For example, the failure of the linear model might indicate the importance of the nonlinear activation functions included in our final model. The convolutional architecture captures quantum wave dispersion, but does not capture correctly the interaction between waves and the potential barrier. It might suggest, that to correctly capture the tunneling and scattering phenomena, we must mix the information not only between different spatial points but also between different channels. The dense architecture is able to capture both the reflection and the tunneling phenomena, but yields significant errors comparing to the recurrent-based model. This indicates, how important the temporal dimension is – something what recurrent architectures are designed to explore, as they are capable of (selectively) storing the memory of the previous steps in their internal states – and retrieving them when needed ^{44, 46}.

The usefulness of a neural network is mainly determined by its generalization capability. In this section, we provide a systematic analysis of both the in-domain and the out-of-domain predictive performance of our emulator.

In the training process, the raw emulation data is broken up into small windows, and those windows are re-sampled to build the curriculum. One of the reason for doing so, was to balance cases featuring different distinct phenomena, e.g., free propagation vs. tunneling through the potential barrier. During the training, we artificially restricted our training instances only to those featuring Gaussian packets and single rectangular potential barriers. This allows us later to test the out-of-distribution generalizability potential – namely, whether our emulator can correctly handle packets of different modulation or barriers of complex shape.

The barrier used in our training instances had a variable height, but a fixed width of $7\,a.u$. Notably, the width of the window, that we used to cut chunks of the input data for our neural network, was $2.25\,a.u.$ – i.e., it was smaller than the width of the potential barrier itself. As a result, the network was exposed during the training to three distinct situations: (1) freely dispersing quantum waves in zero potential; (2) quantum wave propagation with non-zero constant potential; (3) quantum wave propagation with a potential step from zero to a constant value of the potential (or vice versa). While it is obvious that our emulator should handle rectangular potential barriers wider than the width of the window, since such situations are analogous to those already encountered in the training process, it is not all so obvious that the same should happen with barriers of a smaller width, yet alone with barriers of different shapes than rectangular. However, as it was already presented in Fig. 2, those cases do not present a challenge for our emulator.

In Figure 4, we present a more systematic study of that phenomena. We have evaluated how the accuracy of the emulation is affected when changing the following parameters: (a) the initial wave packet spread, (b) the initial wave packet energy, (c) the rectangular barrier width, and (d) the rectangular barrier height. In the experiments, we alter one of the parameters, while holding all other parameters at a constant value. We report the average correlation calculated across all the spatial grid points and all the emulated time steps. As we can see, the performance of the emulator remains effectively unchanged for all the tested values of the initial wave spread and the rectangular barrier width. This demonstrates that the network generalizes well with respect to those parameters. Remarkably, the performance does not deteriorate even when the barrier width becomes smaller than the size of the window, i.e., $2.5\,a.u$. This explains why our emulator was able to correctly emulate the evolution with presence of continuously-shaped barriers (that, due to the discrete nature of our input, can be seen as a collection of adjacent 1-pixel-wide rectangular barriers).

With respect to the value of the initial wave packet energy, the correlation curve has a “$\cap$” shape, indicating that the accuracy deteriorates as we leave the energy-region sampled during the training stage. However, this is as much a failure to generalize as the limitation of the way how we represent the data and the environment. Lower energy of the packet means that the evolution of the quantum wave is slower – and we have to predict more steps of the simulation to cover the same spatial distance as for wave packets with a higher energy. Due to the recurrent process of our emulation, this means larger error accumulation. On the other end, when the energy is high, the wave functions change rapidly within each discretized time unit. This means a larger potential for error each time we predict the next time step. Similar situation also happens when the barrier height increases. A steeper barrier means a more dramatic changes to the values of the wave function happening between two consequent time steps. Those both effects can be mitigated by increasing either the spatial or temporal resolution of our emulation.

In this section, to get some insight into how our emulator makes the prediction, we measure the feature attribution (the contribution of each data input to the final prediction). We employ the direct gradient method. Namely, for a given target pixel in the output window, we calculate its gradients with respect to every input value. We follow the intuition that larger gradient is associated with higher importance of measured input ⁴⁷. Since all the target values are complex numbers, this method gives us two sets of values, one for the real and one for the imaginary value of the target pixel,

In Figure 5, we show a sample of the results. Namely, we selected one test example, and as the target we took the central pixel of the predicted frame. Next, we calculated the gradient with respect to each input value. In this concrete example, we can see that the importance of the input pixels increases with each consecutive time steps. However, how much exactly the early steps matter depends on the region where the predictions are made. Far from the potential barrier (cf. Fig. 5a), the gradient measured with respect to the pixels of the first two steps (no. 176 and 177) is close to zero, indicating that those steps carry relative little importance. However, when predicting the evolution of the wave function in the vicinity of the potential barrier (cf. Fig. 5c), the importance of the third-from-the last step (no. 177) increases noticeable. This indicates, that when predicting free propagation, the network simply performs a linear extrapolation based on the latest two time-steps. However, when predicting tunneling and scattering, the emulator reaches beyond the last two time-steps for additional information, likely due to the non-linear character of those predictions.

We consider the one-dimensional time-dependent Schrödinger equation in atomic units,

where $x\in[0,L_{x}),t\in[0,T)$, and the Hamiltonian operator $\mathcal{H}$ is defined as

with $T$ and $V$ being the kinetic and potential energy operators, respectively.

We represent the quantum waves as complex-valued functions, $\psi(x,t)\in\mathbb{C}$, on a $N_{x}\times N_{t}$ mesh grid, and we impose periodic boundary conditions: $\psi(x+L_{x},t)=\psi(x,t)$, where $L_{x}$ is the size of our 1-dimensional environment.

We use an equal mesh spacing. With $\Delta x=L_{x}/N_{x}$ and $\Delta t=T/N_{t}$, the spatial and temporal coordinates become $x_{i}=i\Delta x$ and $t_{j}=j\Delta t$, with $i\in\{0,1,\ldots,N_{x}-1\}$ and $j\in\{0,1,\ldots,N_{t}-1\}$, respectively. For brevity, we denote the discrete wave function values as $\psi_{i}^{j}=\psi(x_{i},t_{j})$, and the discrete potential steps as $v_{i}=V(x_{i})$.

Given the discretizations described above, we consider the neural network-based emulator as a parameterized map from an input, constructed from the $H$ consecutive time steps, $\{t_{j},t_{t+1},\ldots,t_{j+H-1}\}$, to the output that relates to the next time step, $t_{j+H}$. To make the input and the output of the neural network independent from the size of the system, we construct a slices (we call them windows through the text) of a fixed width $W$. Namely, a portion of the input representing the values of the quantum wave at the time step $t_{j}$, can be written as

Similarly, a portion of the input representing the potential landscape at the same time step, can be written as

As shown in Fig. 2, the input of raw simulation data can be decomposed into a series of windows. The proposed neural network-based emulator maps each input data window to an output data window, that can be subsequently re-combined to form a complete time step. More specifically, we look for a network to learn the mapping

During the training, we try to find parameters $\bm{\Theta}$ of the map $f_{\bm{\Theta}}$, that minimize the training objective function (a loss function $\mathcal{L}$) on each window output, defined as

where $\bm{\hat{\omega}}_{i}^{j+H}$ denotes the predicted values of the quantum wave at the time step $t_{j}$ in the considered window, and $\bm{{\omega}}_{i}^{j+H}$ denotes the ground truth values.

To make the emulator independent of the environment size, we predict only the local evolution of the wave packet $\bm{\omega}_{i}^{j}$ within a given window (cf. Eq. 6). However, to recreate the entire wave function $\psi(x,t_{j})$ for a given time step $j$, we need to assemble the individual predictions from different, overlapping windows. Let us denote the $\psi_{m}^{j}$ evaluated from window $\bm{\omega}_{m+k}^{j}$ as $\psi_{m}^{j}\left(\bm{\omega}_{m+k}^{j}\right)$. We average those predictions by giving a larger weight to those predictions for which the grid-point $x_{m}$ is located closer to the center of the window. For the weight, we use the Gaussian modulation,

where $A=1/\sum_{k=-\lfloor(W-1)/2\rfloor}^{\lfloor W/2\rfloor}\exp\!{\left(-k^{2}\middle/2\delta^{2}\right)}$ is the normalization constant, and $\delta$ denote the Gaussian averaging spread.

A complete prediction of the next time-step is a collection of all predicted points $[\langle\psi_{0}^{j}\rangle,\langle\psi_{1}^{j}\rangle,\ldots,\langle\psi_{N_{x}}^{j}\rangle]$. Fortunately, all the intermediate predictions $\psi_{m}^{j}\left(\bm{\omega}_{m+k}^{j}\right)$ can be obtained independently, therefore the entire algorithm is easily parallelizable (e.g., by performing a batch inference on a modern GPU).

We predict the next time-step of the system evolution in a recurrent manner, i.e., the predictions of the previous time steps are used to form the input of the current time-step. By iteratively predicting next time-steps, we can obtain a sequence of snapshots, portraying the wave function evolution, of an arbitrary length.

As an initial condition, the trained models are fed in with first four time steps of the system evolution.

Our data generation implements a two-step procedure. First, we generate raw simulation data that represents the state of the system. Then we slice them into smaller windows to feed into the neural network.

The input of the neural network has dimension $4\times 23\times 3$ (four time steps, window size $23$, three channels). First the input tensor is reshaped into a $4\times 69$ tensor, and then connected to a time-distributed dense layer (a dense layer that is applied independently to each time step) with the ReLU activation functions. We treat the size of the dense layer, $K$, as a tunable hyper-parameter. We pass the resulting $4\times K$ tensor to a gated recurrent unit, GRU ⁴⁴ with the ReLU activation as well. The output of that step is a single tensor of size $K$. We connect it to the output layer of size $46$, with a linear activation (see again the schematic representation of the network in Fig. 2a).

In addition to the architecture described above, we used the following three benchmark network architectures:

1. 1.

   Linear model. For this model, we used a simplified input. Instead of four time-steps, the linear model takes a single time-step (thus, the input tensor shape is $1\times 23\times 3$). We connect the input directly to the output layer of size $46$. There is no hidden layers. All activations are linear.

2. 2.

   Fully dense model. For the input we take the standard four time-steps (represented by a $4\times 23\times 3$ tensor). First, the input is reshaped into a $4\times 69$ tensor. Next, each time step of the input goes through a time-distributed dense layer of size $K$ with the ReLU activation. The output of the time-distributed layer is subsequently flattened and connected to a regular dense layers of size $K$, also with the ReLU activation. The resulting tensor is finally connected to the output layer of size $46$ with a linear activation.

3. 3.

   Convolutional model. We use the same input as above, four time-steps represented as $4\times 23\times 3$ tensors. We apply 1-dimensional convolution with a filter size of $4$ and with the ReLU activation. The purpose of that layer is to mix the original channels (encoding the real and imaginary part of the wave packet as well as the values of the potential) from different temporal points. The resulting tensor has shape $1\times 23\times F$, where $F$ is the number of the convolution filters. We took $F=\lfloor K/4\rfloor$. The resulting tensor is subsequently flattened and connected to a hidden layer with $K$ neurons, and with the ReLU activation. The last layer is the output of size $46$ with a linear activation.

The outputs of all the models have sizes of 46. In the last step, the outputs are reshaped to form $23\times 2$ tensors, with the second dimension encoding the real and imaginary part of the wave function.

We have trained the neural networks using AdamW ⁴⁵ optimizer with the mean squared error (MSE) loss function. We used a linearly-decayed learning rate with a warm-up. We included more information regarding the training procedure in the Supplementary Information.

To evaluate the model’s performance, we compare the predicted evolution of the system with the ground truth, using the following two metrics:

1. 1.

   Mean absolute error (MAE), calculated per each time step,

   $$\overline{|\epsilon|}=\sum_{i=1}^{N_{x}}\frac{{|\hat{\psi_{i}}-\psi_{i}|}}{N_{x}},$$

   (12)

   where $N_{x}$ is the number of spatial points, while $\hat{\psi_{i}}$ and $\psi_{i}$ are the ground truth and the predicted wave function value, respectively, evaluated in the i-th discrete position, $x_{i}$. MAE remains at low values for a good model. However, a consistent low MAE is not sufficient to determine a good model, as a model that constantly outputs zero or near-zero values might result in a relatively low MAE in some cases (e.g., when the ground truth is a wave packet concentrated only in a small volume of a larger space). This limitation of MAE, leads us to introduce another metric, namely, the normalized correlation.

2. 2.

   Normalized correlation per time step:

   $$\mathcal{C}=\frac{\sum_{i=1}^{N_{x}}{\hat{\psi_{i}}^{*}\psi_{i}}}{|\bm{\hat{\psi}}||\bm{\psi}|},$$

   (13)

   where, $N_{x}$, $\hat{\psi_{i}}$ and $\psi_{i}$ are defined in the same way as above. The symbol ^∗ represents the complex conjugate. Normalized correlation treats the predicted and true wave functions as two vectors, i.e. $\bm{\hat{\psi}}$ and $\bm{\psi}$. To this end, normalized correlation can be understood as the angular similarity between two wave function vectors.

For an overall performance evaluation, we calculate the average over a number of consequitive time steps, denoted by $\langle\overline{|\epsilon|}\rangle$ and $\langle\mathcal{C}\rangle$ respectively.