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.
Foundation models for dynamical systems can forecast unseen partial differential equations (PDEs), ordinary differential equations (ODEs), and time series during inference. However, it remains an open question whether these foundation models truly learn generalizable in-context forecasting strategies. We show that looped transformers are capable of emulating propagators for dynamical systems in-context by first establishing a link between linear attention and (extended) dynamic mode decomposition. We then generalize this insight into the Test-Time Forecasting (TTF) framework in which a one-step forecaster leverages iterations of gradient descent as recurrent depth. We demonstrate that lightweight TTF models trained on ODEs and PDEs generalize to out-of-distribution system dynamics during autoregressive rollout after being trained for one-step prediction. Our results reinforce TTF as a theoretically motivated framework for designing recurrent depth models capable of forecasting unseen dynamical systems at test-time.
Jeffrey B. Lai, John Quinn, Anthony Bao +1
Oden Institute University of Texas at Austin · Department of Physics University of Texas at Austin
Non-stationary sequences arise naturally in control, forecasting, and decision-making. The data-generating process shifts at unknown times, and models must detect the change, discard or downweight obsolete evidence, and adapt to new dynamics on the fly. Transformer-based foundation models increasingly rely on in-context learning for time series forecasting, tabular prediction, and continuous control. As these models are deployed in non-stationary environments, understanding their ability to detect and adapt to regime shifts is important. We formalize this as an in-context change-point detection problem and formally establish the existence of transformer models that solve this problem. Our construction demonstrates that model complexity, in layers and parameters, depends on the level of information available about the change-point location, from no knowledge to knowing exact timing. We validate our results with experiments on synthetic linear regression and linear dynamical systems, where trained transformers match the performance of optimal baselines across information levels. We also show that encoding and incorporating changepoint knowledge indeed improves the real-world performance of a pretrained foundation models on infectious disease forecasting and on financial volatility forecasting around Federal Open Market Committee (FOMC) announcements without retraining, demonstrating practical applicability to real-world regime changes.
Carson Dudley, Yutong Bi, Xiaofeng Liu +1
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
Is it possible to understand the intricacies of a dynamical system not solely from its input/output pattern, but also by observing the behavior of other systems within the same class? This central question drives the study presented in this paper. In response to this query, we introduce a novel paradigm for system identification, addressing two primary tasks: one-step-ahead prediction and multi-step simulation. Unlike conventional methods, we do not directly estimate a model for the specific system. Instead, we learn a meta model that represents a class of dynamical systems. This meta model is trained on a potentially infinite stream of synthetic data, generated by simulators whose settings are randomly extracted from a probability distribution. When provided with a context from a new system-specifically, an input/output sequence-the meta model implicitly discerns its dynamics, enabling predictions of its behavior. The proposed approach harnesses the power of Transformers, renowned for their in-context learning capabilities. For one-step prediction, a GPT-like decoder-only architecture is utilized, whereas the simulation problem employs an encoder-decoder structure. Initial experimental results affirmatively answer our foundational question, opening doors to fresh research avenues in system identification.
Marco Forgione, Filippo Pura, Dario Piga
IDSIA Dalle Molle Institute for Artificial Intelligence USI-SUPSI, Via la Santa 1, CH-6962 Lugano-Viganello, Switzerland.