What can linear attention learn from nonlinear teachers in-context?
Authors: Mary Letey, Arman Rysmakhanov, Yue M. Lu, Cengiz Pehlevan, Jacob Zavatone-Veth
Organizations: Applied Mathematics, Harvard University Kempner Institute, Harvard University · Williams College · Center for Brain Science, Harvard University Society of Fellows, Harvard University
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.
Figures & tables
Figure 1 : We simulate ICL test loss directly from data to verify Result 3 . The linear target is y=w⊤x , while the nonlinear target is sin(w⊤x) . d=64 , α=2.0 , κ=2.0 , Ctrain=Id=Ctest .
Figure 2 : We see that nonlinear functions raise the limiting ICL test error, even in the best possible data case of α,κ,τ→∞ due to nonlearnable signal captured as noise. Here d=32,τ=α/2,ρ=0.1 , f(x)=tanh(x).
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)