Organizations: PMI Lab, School of Physics, Sun Yat-sen University, Guangzhou 510275, People’s Republic of China · Guangdong Provincial Key Laboratory of Magnetoelectric Physics and Devices, Sun Yat-sen University, Guangzhou 510275, People’s Republic of China
Pretrained Transformer architectures display a surprising in-context learning ability without additional training. How the learning process guides the network to achieve this intriguing capability remains unsolved. In this study, we focus on a single-layer, single-head linear attention model, where empirical risk minimization is defined for learning. Given a finite amount of training examples, we determine the exact solution of network parameter dynamics, which is linked to the training error and generalization error, characterizing how memorization and generalization emerge from learning and how these transitions are related to curvatures of the empirical risk landscape. Furthermore, we calculate a data-dependent free energy, from which the training and generalization errors in the long-time limit of learning dynamics can be well predicted. Our work adds insights toward understanding the mechanism of in-context learning, especially from a learning dynamics perspective.
Figures & tables
Figure 1: Eigenspectra of the rescaled Hessian matrix. Blue histograms: three independent datasets for each α , with D=N=100 (a-c) and D=N=20 (d-f). Red curves: Eq. ( 36 ). The α=0.5 histogram excludes the zero eigenvalues but is normalized by all eigenvalues, so its area is 0.5 . The zero atom is stated separately. The α=1 density diverges at the origin; its vertical axis is truncated to show the rest of the spectrum.
Figure 2: Stationary training and test losses as a function of the number of training prompts P for different context lengths N , with fixed input dimension D=10 . Solid curves denote our theoretical predictions from the free-energy calculation in Sec. IV , while markers represent numerical estimates from Langevin dynamics simulations [see Eq. ( 25 )]. In the simulations, we set β=103,λ=0.01 , and γ=1 . After an equilibration period of 105 steps, we sample 20 points separated by 5000 steps and average them to estimate the final loss. (a) As the task diversity increases, the stationary training loss increases for all values of N . The smaller the value of N , the faster the training loss increases. (b) As the task diversity increases, the stationary test loss decreases for all values of N . However, for a given task diversity, smaller values of N consistently yield a higher test loss.
Figure 3: Stationary training and test losses as a function of the number of training prompts P for different input dimensions D , with fixed context length N=100 . Solid curves denote the theoretical predictions from the free-energy calculation in Sec. IV , while markers represent numerical estimates from Langevin dynamics simulations [Eq. ( 25 )]. In simulations, we set β=103 and λ=0.01 , and γ=1 . (a) The training loss increases with P and is larger for a higher-dimensional input. (b) The test loss decreases as P grows, but a larger D requires more task diversity to reach a comparable generalization performance.
Figure 4: Gradient-flow dynamics of the training and test losses for task diversity P=10 and 100 , with D=10 and N=100 . The regularization parameter λ=0.01 . Solid curves represent the analytical gradient-flow solution with the shaded regions indicating the theoretical uncertainty (see Eq. ( 16 ) in Sec. III in the main text), while markers represent gradient descent simulations. Asymptotic error means the stationary limit of the error predicted by the zero-temperature free energy analysis in Sec. IV . (a) Training losses rapidly converge to its stationary value for both small and large P . (b) Test losses exhibits qualitatively different behavior depending on the task diversity: for the small P , the model exhibit memorization-like behavior and the test loss increases, whereas for the larger P , the test loss decreases monotonically, indicating the emergence of in-context generalization.
Figure 5: Dependence of the relative generalization gap ΔEr and the contrast ratio C on task diversity P . The results are obtained by taking the long-time limit of the analytical gradient-flow solution derived in Sec. III with the regularization strength λ=0.01 . The main panels show the generalization gap, while the insets show the contrast ratio. Shaded bands are ±1 s.d. over 100 training sets per P . Different color intensities correspond to different values of N (or D ). The relative gap can be slightly negative at large P because Et includes the weight-decay term. (a) Results for different context lengths N at fixed D=10 . (b) Results for different input dimensions D at fixed N=100 . The dashed line of ΔEr=5 identifies a rough estimate of the value of P separating memorization from generalization.
Figure 6: Dependence of the relative generalization gap ΔEr and the contrast ratio C on task diversity P . We focus on W21=0 , and λ=0 . The simulations of gradient-flow dynamics are run up to 105 steps for equilibrium. The main panels show the generalization gap, while the insets show the contrast ratio. Different color intensities correspond to different values of N (or D ). Lines are averages over 50 training realizations. Error bars are shown only for specific values of P . Dotted lines mark P=D2 . (a) Results for different context lengths N at fixed D=10 . (b) Results for different input dimensions D at fixed N=100 . The dashed line of ΔEr=5 identifies a rough estimate of the value of P separating memorization from generalization.
Modern distributed networks, notably transformers, acquire a remarkable ability (termed `in-context learning') to adapt their computation to input statistics, such that a fixed network can be applied to data from a broad range of systems. Here, we provide a complete mechanistic characterization of this behavior in transformers trained on a finite set S of discrete Markov chains. The transformer displays four algorithmic phases, characterized by whether the network memorizes and generalizes, and whether it uses 1-point or 2-point statistics. We show that the four phases are implemented by multi-layer subcircuits that exemplify two qualitatively distinct mechanisms for implementing context-adaptive computations. Minimal models isolate the key features of both motifs. Memorization and generalization phases are delineated by two boundaries that depend on data diversity, K=∣S∣. The first (K1∗) is set by a kinetic competition between subcircuits and the second (K2∗) is set by a representational bottleneck. A symmetry-constrained theory of a transformer's training dynamics explains the sharp transition from 1-point to 2-point generalization and identifies key features of the loss landscape that allow the network to generalize. Put together, we show that transformers develop distinct subcircuits to implement in-context learning and identify conditions that favor certain mechanisms over others.
Cole Gibson, Wenping Cui, Gautam Reddy
Lewis-Sigler Institute for Integrative Genomics, Princeton University, Princeton, NJ 08544, USA · Department of Physics, Princeton University, Princeton, NJ 08544, USA
In-context learning (ICL) enables a pretrained model to infer a task from demonstrations without updating its parameters. While much of the existing theory focuses on linear target functions, in this paper we study nonlinear cases by comparing two one-layer attention architectures on the same family of single-index tasks. A kernel learner first maps inputs through a fixed nonlinear feature map and then applies linear attention, whereas a feature learner applies attention to the original input, followed by a learned nonlinear readout. We derive predictions for their memorization and generalization errors using the replica method, retaining the effects of pretraining size, task-pool diversity, and training and inference context lengths. The resulting predictions closely match numerical experiments across a broad range of regimes. Our analysis yields phase diagrams that characterize when each architecture is advantageous as the amount of pretraining data, task diversity, and context lengths vary. We further identify qualitatively different context-length scalings for the two learners. Together, these results clarify how architectural choices interact with the dataset and govern nonlinear in-context learning.
Haotian Gu, Yizhou Xu, Lenka Zdeborová
Statistical Physics of Computation Laboratory, École Polytechnique Fédérale de Lausanne (EPFL) · University of Chinese Academy of Sciences · Information, Learning and Physics Laboratory, École Polytechnique Fédérale de Lausanne (EPFL)
Linear attention is a tractable model for understanding the mechanisms governing in-context learning in transformers. For linear regression tasks, recent asymptotic analyses have characterised its learning and generalisation behaviour. We extend this theory to nonlinear single-index targets, y=f(x⊤w)+ε. Our main result establishes a nonlinearity-noise equivalence: linear attention extracts only the linear Hermite component of f, while the remaining nonlinear structure contributes to the generalisation error as effective noise. This reduction allows results from the corresponding linear theory to be transferred to nonlinear tasks. We illustrate its implications for finite pretraining data and for the transition from task memorisation to task generalisation as task diversity increases. These results identify a limitation of the reduced linear-attention model and provide a tractable starting point for studying nonlinear in-context learning.
Mary Letey, Arman Rysmakhanov, Yue M. Lu +2
Applied Mathematics, Harvard University Kempner Institute, Harvard University · Williams College · Center for Brain Science, Harvard University Society of Fellows, Harvard University