Transformers for dynamical systems learn transfer operators in-context
Authors: Anthony Bao, Jeffrey Lai, William Gilpin
Organizations: Department of Electrical & Computer Engineering, The University of Texas at Austin, Austin, Texas 78712, USA · The Oden Institute, The University of Texas at Austin, Austin, Texas 78712, USA · Department of Physics, The University of Texas at Austin, Austin, Texas 78712, USA
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 time without retraining, and they illustrate the unique ability of attention-based models to leverage global attractor information in service of short-term forecasts.
Figures & tables
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 test 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: Scaling laws for out-of-distribution generalization in dynamical systems. (A) The test error of the unseen system (cross-entropy, Test-OOD) versus the difference between the training and testing sets (KL divergence between the attractor of Test-OOD and Train-ID). (B) The error in the steady-state invariant distribution of the transformer dynamics relative to the true invariant distribution of Test-OOD, versus the context length C . Dashed line shows the expected scaling for an estimator that counts observed state transitions in-context [ 7 ] . Points indicate 500 transformers trained and tested on distinct ID/OOD dynamical system pairs, and solid lines indicate linear fits with 99% confidence intervals.
Figure 3: Transformers perform time-delay embedding during inference. (A) Empirical time-delayed next-token probabilities p^(xt+1∣xt−k) averaged across Test-OOD context for a transformer trained on a different system (Train-ID). (B) Exact next-token probabilities ptrue(xt+1∣xt−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 dM 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. Stable rank is scaled by the maximum possible rank. Annotated slopes indicate the linear scaling predicted by Takens’ theorem. Each point is an average over 100 replicate models trained on a different low-dimensional dynamical system.
Figure 4: Transformers learn in-context transfer operators on reconstructed dynamical manifolds. (A) The eigenvalue spectrum of the transfer operator estimated from the fully-observed Test-OOD attractor p(yt+1∣yt) , and the time-lagged transfer operator estimated by sampling contiguous length- k sequences from the transformer p^(y^t+1∣y^t)=p^(xt+1,xt,...xt−k∣xt,xt−1,...xt−k−1) . (B) (Top) The invariant distribution π(0)(y) and longest-lived metastable distributions ( π(i)(y) , i>0 ) of the true transfer operator, ranked by ∣λi∣ . (Bottom) The invariant and metastable distributions of the transformer, based on its implied transfer operator p^(y^t∣y^t−1) . Because the transformer is univariate, eigenvectors are plotted on the time-delay embedding y^≡[xt,xt−1]⊤ .
Figure S1: (A) Test-OOD loss curves for replicate models trained with varying context lengths C . (B) The final cross-entropy loss for training data, in-distribution test data, and out-of-distribution test data, as a function of context length C .
Figure S2: The estimated manifold dimension dM of the attractor of the Lorenz-96 system as the number of dynamical variables D varies.
Department of Mathematics, University of Michigan · Department of Electrical Engineering and Computer Science, University of Michigan · Michigan Institute for AI and Data in Society, University of Michigan