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.
Lewis-Sigler Institute for Integrative Genomics, Princeton University, Princeton, NJ 08544, USA · Department of Physics, Princeton University, Princeton, NJ 08544, USA
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)
Applied Mathematics, Harvard University Kempner Institute, Harvard University · Williams College · Center for Brain Science, Harvard University Society of Fellows, Harvard University