Source-linked AI summary
What can linear attention learn from nonlinear teachers in-context?
Mary Letey, Arman Rysmakhanov, Yue M. Lu, Cengiz Pehlevan, Jacob Zavatone-Veth
TL;DR
The paper asks which components of nonlinear tasks linear attention can learn in-context, beyond the linear-regression setting already analyzed. It studies a reduced linear-attention predictor through the link function's Hermite modes and finds that only the linear mode is learned, while the others act as effective noise.
Problem
Prior asymptotic analyses characterize linear-attention ICL for linear regression, leaving which components of nonlinear tasks it can extract in-context as an open question.
Source: §1 · pp. 1–2Method
Results
Takeaways & Limitations
Beyond the task-diversity transition, error approaches the nonlinear-noise floor rather than the original label-noise level.
Source: §4 · p. 4The analysis assumes pretraining and test tasks are sampled from Gaussian task distributions, with label-noise variance ρ.
Source: §3 · p. 3
Abstract
from arXiv · showhide
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^\top w)+\varepsilon $. 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.
1 Introduction
Linear attention offers a tractable setting for analyzing in-context computations, but which parts of nonlinear tasks it can extract remains a question. This work extends linear-regression ICL analysis to single-index targets and formalizes a limitation: the model learns only the linear Hermite mode.
1 Introduction
- Linear attention provides a useful setting for detailed analysis of prompt-data computations, and recent asymptotic analyses describe generalization in linear-regression tasks (Zhang et al., 2023; Von Oswald et al., 2023; Lu et al., 2025; Bordelon et al., 2025; Letey et al., 2026b,a).Source: §1 · p. 1
- The paper asks which components of a nonlinear task linear attention can extract in-context and which must come from the surrounding architecture.The question is motivated by the increasing prevalence of linear-attention mechanisms in real-world models.Source: §1 · pp. 1–2
- The analysis extends a reduced linear-attention regression ICL model to nonlinear single-index targets, while focusing on linear attention's limitations rather than adding a nonlinear readout as concurrent work does (Gu et al., 2026).The paper analyzes the link function through its Hermite expansion (Joshi et al., 2026; Eftekhari et al., 2021; Dudeja and Hsu, 2018).Source: §1 · p. 2
- Linear attention cannot learn beyond the link function's linear Hermite mode; without that mode, all signal is indistinguishable from noise.The paper relates this intuition to results on learnability of single-index targets under student-link incompatibility.Source: §1 · p. 2
2 Setup
The setup studies a model trained across regression tasks that uses prompt examples to predict a response for a new covariate. It specifies task and noise distributions, a reduced linear-attention predictor, and ridge-regression evaluation.
2 Setup
- The model trains across regression tasks and uses prompt examples to predict the response to a new covariate.Training and test tasks are sampled from potentially different distributions, with training tasks shared across sequences.Source: §2 · p. 2
- For each prompt sequence, covariates and label noise are Gaussian, while task vectors are sampled from separate training and test task distributions.Training uses k independently sampled tasks distributed across n sequences; test tasks are sampled independently.Source: §2 · p. 2
- The setup expands the link function in an orthonormal Hermite basis and assumes it has at most polynomial growth.The passage states the growth condition as |f(x)| < A(1 + |x|^s) for some power s and constant A.Source: §2 · p. 2
- The next-token prediction is approximated by a reduced linear-attention predictor, whose parameters are fit by ridge regression on n training prompts and evaluated with test error.The provided passages introduce the predictor and test-error expressions without giving their formulas.Source: §2 · p. 3
3 Theoretical results and analysis
The analysis reduces nonlinear single-index ICL to the target’s linear Hermite component, with remaining modes acting as effective noise. Under stated conditions, this produces an error floor above the linear-target noise floor, and Figure 1 illustrates equivalence to a higher-noise linear target.
- Linear attention uses only f’s linear Hermite mode; the remaining modes act as effective noise.Source: §3 · p. 3
- Figure 1’s direct-data simulation shows nonlinear ICL at baseline noise is empirically equivalent to linear ICL at higher noise.Source: §3 · p. 3, Figure 1 · p. 4
- With identity covariances and κ>1, error approaches a new floor when a1≠0, above the linear-target floor ρ; Figure 2 shows raised limiting error even as α, κ, and τ grow.Source: §3 · p. 4, Figure 2 · p. 4
4 Conclusion
The paper extends reduced linear-attention ICL theory from linear regression to nonlinear single-index targets. It finds that only the linear Hermite component is used, leaving nonlinear components as effective noise that creates an error floor.
- The asymptotic ICL theory for a reduced linear-attention block is extended from linear regression to nonlinear single-index targets.Source: §4 · p. 4
- The model uses the target function through its linear Hermite component, while the remaining Hermite modes contribute as effective noise.Source: §4 · p. 4
- More pretraining data or task diversity cannot remove the nonlinear-component error floor, which persists beyond the task-diversity transition.Source: §4 · p. 4
B Population averages
The population-average calculations treat w⊤x as asymptotically distributed according to g and use Hermite expansion. They evaluate conditional moments by separating the component of x aligned with w from its orthogonal component and computing three blocks.
B Population averages
- The moment calculation treats w⊤x as asymptotically distributed according to g and writes ∥w∥2 as dω, with ω asymptotically treated as 1.The calculation proceeds using Hermite expansion.Source: B Population averages · p. 8
- For fixed ∥w∥2 = dω with ω ≈1, the section sets up conditional expectations of order-two quantities over x and y.Source: B Population averages · p. 8
- The calculation decomposes x into components aligned with w and orthogonal to w, then combines their contributions.Source: B Population averages · p. 8
- The final moment calculation computes three blocks directly and neglects the last 1/ℓ term in the final expression.The resulting leading terms in the three blocks give equation (14).Source: B Population averages · p. 9
C Derivation of the ICL formula
The derivation transfers the corresponding linear-theory calculation to nonlinear links by expressing their effect through transformed second moments. When a1 = 0, those moments lose task randomness, yielding a simpler resolvent calculation.
C Derivation of the ICL formula
- When a1 ≠ 0, the derivation obtains the stated ICL-error result by matching the nonlinear second moments to the corresponding linear-theory expressions.The link function f affects the second moments of y2, yH, and vec(H)vec(H)⊤ only through transformations.Source: C Derivation of the ICL formula · p. 10
- Setting a1 = 0 removes task randomness from the conditional moments and simplifies the resolvent calculation.The resulting extended resolvent has a similar self-consistent asymptotic expression to that in Letey et al. (2026b).Source: C Derivation of the ICL formula · p. 10
- Taking the ridgeless limit λ → 0 and substituting ˜λ = λ(1 + χ0) and C → 0 into Letey et al. (2026b)’s formula gives exactly (10).Source: C Derivation of the ICL formula · p. 10