Fourier Neural Operators for Learning Dynamics in Quantum Spin Systems

TL;DR

The paper addresses the challenge of efficiently simulating quantum spin dynamics by leveraging Fourier Neural Operators (FNOs) to learn the time-evolution operator. It introduces two architectures (energy-domain and time-domain) that operate on the full wavefunction, plus a Hamiltonian-observable variant that compresses inputs/outputs to poly$(n)$ terms, enabling extrapolation to longer times and larger systems. The results show long-time extrapolation with errors around a few percent (e.g., ~5.8% for a 20-qubit system) and substantial inference speedups (up to ~$10^4$×) over exact solvers, with zero-shot super-resolution demonstrated on finer grids. This work suggests a scalable, data-efficient pathway for predicting quantum dynamics beyond coherence and tensor-network limits, with potential applications to noisy devices and large-scale quantum simulations.

Abstract

Fourier Neural Operators (FNOs) excel on tasks using functional data, such as those originating from partial differential equations. Such characteristics render them an effective approach for simulating the time evolution of quantum wavefunctions, which is a computationally challenging, yet coveted task for studying quantum systems. In this manuscript, we use FNOs to model the evolution of quantum spin systems, so chosen due to their representative quantum dynamics. We explore two distinct FNO architectures, examining their performance for learning and predicting time evolution on both random and low-energy input states. We find that standard neural networks in fixed dimensions, such as U-Net, exhibit limited ability to extrapolate beyond the training time interval, whereas FNOs reliably capture the underlying time-evolution operator, generalizing effectively to unseen times. Additionally, we apply FNOs to a compact set of Hamiltonian observables ($\sim\text{poly}(n)$) instead of the entire $2^n$ quantum wavefunction, which greatly reduces the size of our FNO inputs, outputs and model dimensions. Moreover, this Hamiltonian observable-based method demonstrates that FNOs can effectively distill information from high-dimensional spaces into lower-dimensional spaces. Using this approach, we perform numerical experiments on a 20-qubit system and extrapolate Hamiltonian observables to twice the training time with a relative error of $5.8\%$. Relative to numerical time-evolution methods, FNO achieves an inference speedup of approximately $10^{4}\times$ for 20-qubit systems. The extrapolation of Hamiltonian observables to times later than those used in training is of particular interest, as this stands to fundamentally increase the simulatability of quantum systems past both the coherence times of contemporary quantum architectures and the circuit-depths of tractable tensor networks.

Figures (4)

• Figure 1: (a) Schematic representation of quantum spin system used to model its dynamics using FNOs, where $J_x$, $J_y$, and $J_z$ represent the coupling constants and $h$ denotes the external driving field. (b) Diagram illustrating the FNO framework, which consists of multiple FNO blocks performing spectral convolution over the input in the latent space. (c) We investigate two distinct architectures for learning dynamics in quantum spin systems. The first is the energy-domain architecture (see Section \ref{['subsection:energy']}), where the input is the wavefunction at the initial time $t=0$, and the output is the wavefunction evolved to time $t=T$. (d) The second is the time-domain architecture (see Section \ref{['subsection:time']}), which takes the wavefunction evolved over the initial time interval $[0,\frac{3}{2}T]$ (discretized on a grid with width $\Delta t$) as input and produces an output wavefunction over the time interval $[T,\frac{5}{2}T]$.
• Figure 2: (a) Prediction of temporal dynamics in an 8-qubit Heisenberg spin system using the energy-domain architecture in Section \ref{['subsection:energy']}. At time $T=\pi$, the predictions made by the FNO are on unseen data but still within the time range it was trained on. The extrapolated time refers to future time predictions beyond the range on which the FNO was trained. VTI (Various Time Interval) indicates training on multiple time intervals, i.e., $[0, T]$, $[T, 2T]$, and $[2T, 3T]$, instead of training only on the first interval. (b) The results show the fidelity at specific time steps for 4 and 8 qubits, including times within the training range as well as extrapolated future times beyond the training interval. Superscripts indicate whether the particular time step is within the training range ($\mathrm{train}$) or an extrapolated time ($\mathrm{ext}$).
• Figure 3: (a) Prediction of temporal dynamics in a four-qubit Heisenberg spin system using the time-domain architecture in Section \ref{['subsection:time']}. The model is trained on the time interval $[0, 3\pi/2]$ as illustrated in Figure \ref{['fig:model']}d, while the extrapolated time refers to future predictions beyond the trained range. Performance for both the random and low-energy wavefunction initial states using the FNO is benchmarked against a U-Net baseline. The shaded regions around each curve indicate one standard deviation. The FNO accurately captures the underlying operator and extrapolates reliably to later times, while the U-Net exhibits poor extrapolation, especially for low-energy input states. (b) Corresponding results for an eight-qubit Heisenberg spin system.
• Figure 4: (a) Mean Relative Error (MRE) for predicted Hamiltonian observables in a 20-qubit system using FNO, utilizing the time-domain architecture. The model is trained on the input time interval $[0,\frac{3}{2}\pi]$ as shown in Figure \ref{['fig:model']} d) and predicts the output interval $[\frac{3}{2}\pi,\frac{5}{2}\pi]$. Additionally, it extrapolates future unseen dynamics $[\frac{5}{2}\pi,\frac{7}{2}\pi]$, which is twice the length of the output interval. The MRE is $5.8\%$ over the extrapolated time predictions on ground truth (GT). This is particularly significant in quantum simulations, where the primary objective is to extend these simulations. Additionally, based on the time interval the model was trained on, the MRE for extrapolated time predictions is $9.6\%$ denoted by orange dashes. We additionally calculate the mean MSE loss for the training time interval as $2.16 \times 10^{-9} \pm 2.15 \times 10^{-9}$ and for future time predictions over the same interval on ground truth as $5.34 \times 10^{-9} \pm 4.13 \times 10^{-9}$. (b) The results present the MRE at specific time steps for both the train time (i.e., the time intervals that FNO encountered during training) and the extrapolated time [$\frac{5}{2}\pi, \frac{7}{2}\pi$] for systems with 8 and 20 qubits for hamiltonian observables.