Transformers for dynamical systems learn transfer operators in-context

TL;DR

The paper investigates how a small transformer trained on a univariate dynamical trajectory can generalize to forecast unseen dynamical systems without retraining, illuminating the mechanism of in-context learning in physics. It demonstrates that the model spontaneously performs time-delay embedding and implicitly learns a transfer operator (Perron-Frobenius) for the underlying system, matching long-timescale dynamics and metastable structures. By comparing the transformer's inferred operator to fully observed operators via Ulam's method, the study shows that in-context learning yields faithful representations of the system's attractor and dominant modes. These findings reveal a concrete mechanism by which pretrained models generalize to unseen physical systems and underscore the role of global attractor information in short-term forecasting.

Abstract

Large-scale foundation models for scientific machine learning adapt to physical settings unseen during training, such as zero-shot transfer between turbulent scales. This phenomenon, in-context learning, challenges conventional understanding of learning and adaptation in physical systems. Here, we study in-context learning of dynamical systems in a minimal setting: we train a small two-layer, single-head transformer to forecast one dynamical system, and then evaluate its ability to forecast a different dynamical system without retraining. We discover an early tradeoff in training between in-distribution and out-of-distribution performance, which manifests as a secondary double descent phenomenon. We discover that attention-based models apply a transfer-operator forecasting strategy in-context. They (1) lift low-dimensional time series using delay embedding, to detect the system's higher-dimensional dynamical manifold, and (2) identify and forecast long-lived invariant sets that characterize the global flow on this manifold. Our results clarify the mechanism enabling large pretrained models to forecast unseen physical systems at test without retraining, and they illustrate the unique ability of attention-based models to leverage global attractor information in service of short-term forecasts.

Figures (4)

• Figure 1: Double descent during in-context learning of dynamical systems. (A) Forecasts from a transformer trained on univariate time series from one dynamical system (Train-ID) then evaluated on its ability to forecast unseen trajectories from the same system (Test-ID, blue), versus forecasts of trajectories from an unseen system (Test-OOD, magenta). Gray curve shows the context, a subset of the total training data. (B) Loss curves on the training data, versus held-out in-distribution and out-of-distribution validation sets. Ranges are standard errors across $100$ replicate experiments, each corresponding to a different randomly-sampled train/test pair of dynamical systems.
• Figure 2: Transformers perform time-delay embedding during inference. (A) Empirical time-delayed next-token probabilities $p_\text{model}(x_{t+1} | x_{t-k})$ averaged across Test-OOD context for a transformer trained on a different system (Train-ID). (B) Exact next-token probabilities $p_\text{true}(x_{t+1} | x_{t-k})$ obtained from fitting a Markov chain on a long sample of Test-OOD. (C) The order of the closest-approximating Markov chain, versus the true intrinsic dimension $d_\mathcal{M}$ of the full attractor from which the univariate Test-OOD trajectory originates. (D) The effective dimension (stable rank) of the transformer's attention rollout matrix versus the true intrinsic dimension of Test-OOD. Each point is an average over $100$ replicate models trained on a different low-dimensional dynamical system.
• Figure 3: Transformers learn in-context transfer operators on reconstructed dynamical manifolds. (A) (Top) The true invariant distribution ($\pi^{(0)}(\mathbf{y})$) and longest-lived metastable distributions ($\pi^{(i)}(\mathbf{y})$$i>0$, almost-invariant sets) estimated from the fully-observed state space as the leading eigenvectors of the fully-observed transfer operator $p(\mathbf{y}_{t+1} | \mathbf{y}_t)$ (ranked by $\left| \lambda_i \right|$). (Bottom) The invariant and metastable distributions of the transformer, based on its next-token prediction probabilities over time-delayed univariate Test-OOD input sequences. Distributions are estimated as the leading eigenvectors of the lag-$K$ conditional probabilities $\hat{p}(x_{t+1} ... x_{t-K+2} | x_{t}...x_{t-K+1})$ derived from the transformer probabilities. Because the transformer is univariate, eigenvectors are plotted on the time-delay embedding $\hat{\mathbf{y}} \equiv [x_{t}, ..., x_{t-K}]$. (B) The eigenvalue spectrum of the transfer operator estimated from the fully-observed attractor, and the implicit transfer operator of the transformer. (C) The KL divergence between the ground-truth invariant distribution of the transfer operator, and transformer empirical operator's distribution, versus both Test-OOD validation loss and training epoch (color). A 95% confidence interval linear fit is underlaid (blue).
• Figure S1: The estimated manifold dimension $d_\mathcal{M}$ of the attractor of the Lorenz-96 system as the number of dynamical variables $D$ varies.