Attention in Krylov Space

TL;DR

The paper addresses the difficulty of obtaining long Lanczos coefficient sequences $b_n$ due to numerical instability and memory limits, which hinders accurate reconstruction of operator dynamics via Krylov space. It reframes the problem as causal sequence forecasting and trains a decoder-only transformer to predict future $b_n$ from a short prefix using $ abla b_n = b_n - b_{n-1}$, ensuring stable training. The key contributions are: (i) a substantial accuracy boost over traditional asymptotic fits in both $b_n$ extrapolation and observable reconstruction of $K(t)$ and $C(t)$, (ii) successful zero-shot transfer to larger system sizes, and (iii) insights into which history segments drive predictions through attention analysis. This approach enables reliable probing of long-time operator dynamics from limited data, with potential experimental links via spectral moments and broad applicability to diverse Krylov-space problems.

Abstract

The Universal Operator Growth Hypothesis formulates time evolution of operators through Lanczos coefficients. In practice, however, numerical instability and memory cost limit the number of coefficients that can be computed exactly. In response to these challenges, the standard approach relies on fitting early coefficients to asymptotic forms, but such procedures can miss subleading, history-dependent structures in the coefficients that subsequently affect reconstructed observables. In this work, we treat the Lanczos coefficients as a causal time sequence and introduce a transformer-based model to autoregressively predict future Lanczos coefficients from short prefixes. For both classical and quantum systems, our machine-learning model outperforms asymptotic fits, in both coefficient extrapolation and physical observable reconstruction, by achieving an order-of-magnitude reduction in error. Our model also transfers across system sizes: it can be trained on smaller systems and then be used to extrapolate coefficients on a larger system without retraining. By probing the learned attention patterns and performing targeted attention ablations, we identify which portions of the coefficient history are most influential for accurate forecasts.

Figures (9)

• Figure 1: Schematic illustration of a typical sequence of Lanczos coefficients. The dashed line indicates an asymptotic linear fit, while the red segments depict residual deviations that are not captured by simple asymptotic fits. The deviations are systematic subleading structures and can strongly affect reconstructed dynamics approx_Greens_extrap_and_cutoff_spinYates_2020Yates2_staggeringstaggering2Yates_Floquet_staggeringsubleadingsubleading2pseudomode_extrapolate_spindynamics_oneandtwo_dstaggering_cft.
• Figure 2: (a) Illustration of the Decoder-only transformer architecture. The scalar inputs $\{ \Delta b_1 \dots\Delta b_{m} \}$ are mapped to $d_\text{model}$-dimensional token embeddings and augmented by positional encodings to become hidden states $\{\mathbf{h}_1^1 \dots \mathbf{h}_m^1\}$, which pass through $L$ decoder blocks. The final hidden state is projected to the next token prediction $\Delta b_{m+1}$ by the unembedding (output) layer. (b) Internal structure of the $l$-th decoder block, showing the masked multi-head self-attention mechanism, residual connections, and the feed-forward network.
• Figure 3: Transformer performance on the classical XYZ model. (a) Extrapolation of Lanczos coefficients from the input interval $n \leq 10$ to $n=100$. The transformer's predictions remain closer to the ground truth, whereas the linear-fit baseline exhibits a systematic bias as $n$ increases. Inset: Residuals $b_n - b_n^{\text{pred}}$ deep in the extrapolation regime. The transformer's error remains a few times smaller than linear extrapolation. (b) Root Mean Squared Error (RMSE) averaged over 100 unseen instances of $H_{\text{XYZ}}$. The transformer achieves an error a few times smaller compared to the linear baseline, effectively capturing the subleading dynamics of the coefficients.
• Figure 4: Transformer performance on the quantum Transverse-Field Ising Model (TFIM). (a) Extrapolation of Lanczos coefficients for an unseen $L=8$ system. Inset: Residual $b_n-b_n^{\text{pred}}$ for the transformer and asymptotic fit at later $n$. (b) Root Mean Squared Error (RMSE) averaged over $100$ unseen test sequences, demonstrating that the transformer systematically outperforms the asymptotic fit by capturing subleading, non-linear deviations from the asymptotic ramp.
• Figure 5: Root-mean-squared error (Eq. \ref{['eq:deltaKt']}) in the reconstruction of Krylov complexity and autocorrelation function, using extrapolated Lanczos sequences by transformer and asymptotic fit, compared to reconstructions from the exact coefficients. Errors are averaged over $N_T = 100$ test sequences. While both methods are comparably accurate at short times $tJ \lesssim 1$, transformer-based extrapolation leads to substantially smaller errors at later times by orders of magnitude.