What can linear attention learn from nonlinear teachers in-context?

arXiv:2610.10761 · stat.ML, cs.LG · Submitted 2026-10-07 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "What can linear attention learn from nonlinear teachers in-context?".

Jane: The gist: Linear attention extracts only the linear Hermite component of nonlinear functions in-context, while the remaining nonlinear structure contributes to generalization error as effective noise,

Tom: First, who's behind it and why it matters.

Paper summary: Tom: We started by setting up this model for single-index targets, y = f(x w) + epsilon, where epsilon is Gaussian label noise with variance rho. The paper then uses an orthonormal Hermite basis to expand that function, f(x) = sum a m h m(x), which is the first step in looking at the function's structure.

Jane: And they introduce the reduced linear-attention predictor, approximating the next-token prediction y test using a formula involving and H, where H is constructed from training examples, specifically using all of those training inputs and labels.

Lu: The central claim of the paper is the nonlinearity–noise equivalence. They show that if there's a nonzero linear mode, a one not equal to zero then the nonlinear in-context learning error equals the error from a purely linear model with adjusted task covariances and label noise, rho nonlinear <ref:2610.10761#pg1>.

Meng: So they are essentially saying that as long as the function has *some* linear part, we can predict its performance using a simpler linear framework, just by changing how we scale our input covariances.

Tom: But if there is no linear mode at all, a one equals zero, then the error isn't related to that original label noise rho, but instead it becomes this new rho nonlinear value <ref:2610.10761#pg1>.

Jane: This is important because it means that the structure of the function itself dictates what kind of generalization error you're going to face, not just how much data you throw at the problem.

Lu: They also provide a condition on tau: if a one is zero, then the error depends on tau, which is defined as n over d squared, and whether that tau value is less than one or greater than one <ref:2610.10761#pg1>.

Tom: And they show that even in the limit where all those parameters like alpha, tau, and n go to infinity, if there's a nonzero linear mode, the error still equals the linear ICL error with those adjusted variances.

Jane: But if a one is zero, then in that infinite limit, the test error settles at this new rho nonlinear value <ref:2610.10761#pg1>. That’s when it becomes hard to improve performance by just adding more pretraining data or task diversity.

Meng: From an engineering standpoint, it’s a warning that you need to be careful about what kind of nonlinear structure you're training your model on if you expect good generalization on unseen tasks.

Lalam: This paper really gives us a framework to understand why some models might struggle with complex in-context learning: they are just picking up noise instead of learning the useful linear signal.

Conclusion: Tom: So, thinking about the title "What can linear attention learn from nonlinear teachers in-context?" it points to a fundamental question about what this specific mechanism is actually capable of extracting when it’s faced with non-linear context.

Jane: It’s not just about whether the model learns a single function; it’s about *which* components of that function are successfully extracted by the linear attention block versus which ones are left behind as noise.

Lu: The implication is that for any nonlinear task, we can decompose the performance into a learnable linear part and an unlearnable noise part dictated by the function's specific structure.

Meng: For practical AI development, this means we shouldn't just chase higher accuracy numbers on complex tasks if we aren't sure if those gains are coming from learning a new signal or just better noise management.

Lalam: The limitation they point out is that increasing pretraining data or task diversity won't remove that fundamental error floor created by the nonlinear components invisible to this predictor.

Tom: It confirms that this block is fundamentally linear in its extraction capability, and it only pulls out the linear Hermite component of the link function.

Jane: So, we can use this as a starting point for studying how models handle nonlinearity in contexts where we want to understand what signal is actually being learned.

Mary Letey, Arman Rysmakhanov, Yue M. Lu, Cengiz Pehlevan, Jacob A. Zavatone-Veth

Harvard University

stat.ML, cs.LG

Submitted: 2026-10-07

Updated: 2026-10-07

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 88/100

The gist: The gist: Linear attention extracts only the linear Hermite component of nonlinear functions in-context, while the remaining nonlinear structure contributes to generalization error as effective

Key concepts

Hermite Basis Expansion
This technique decomposes a function $f(x)$ into an infinite sum of Hermite polynomials, $f(x) = igoplus_{m=0}^{\infty} a_m h_m(x)$. This allows researchers to separate the function's behavior into different modes, where the linear mode ($a_1$) is distinct from higher-order nonlinear modes.
Linear Attention Predictor
This is a simplified model used during training that approximates next-token prediction. It uses a reduced linear attention mechanism to learn only the first Hermite mode of the target function, effectively ignoring the more complex nonlinear parts of the teacher's function.
Effective Noise Floor ($ ho_{ ext{nonlinear}}$)
When nonlinear components are present, they contribute as noise that cannot be learned by linear attention. This results in a new error floor for in-context learning. Increasing training data or task diversity does not eliminate this floor because the nonlinear signal is invisible to the linear predictor.

Terminology

Summary

The gist: Linear attention extracts only the linear Hermite component of nonlinear functions in-context, while the remaining nonlinear structure contributes to generalization error as effective noise, which allows results from linear theory to be transferred to nonlinear tasks.

Setup for Single-Index Tasks

The model is trained across regression tasks and must use the examples in a prompt to predict the response to a new covariate A prompt is a sequence with yi = f(x⊤i w) + εi, ytest = f(x⊤testw) + εtest. (1) We take xi, xtest i.i.d. ∼ N (0, Id/d) and independent Gaussian label noise ε ∼ N (0, ρIl) and εtest ∼ N (0, ρ). The parameters are set from n such samples of these sequences, each with a potentially different defining task w For the training sequences, we use k tasks sampled independently from N (0, Ctrain); these k tasks are split uniformly across the n sequences in the training batch At test time, the model is tested on sequences with tasks sampled independently from N (0, Ctest)

Nonlinearity and Expansion

For g ∼ N (0, 1), expand f in an orthonormal Hermite basis: f(x) = X∞ m=0 amhm(x), am = Eg∼N(0,1)[f(g)hm(g)] We assume f has at most polynomial growth, i.e., f(x) < A(1 + x s) for some power s and constant A We take d → ∞ with κ = k/d, α = l/d, τ = n/d squared, all of order one, as in prior work (Lu et al., 2025; Letey et al., 2026b)

Reduced Linear-Attention Predictor

The next-token prediction is approximated by ybtest = ⟨Γ, H⟩F, Γ ∈ R d×(d+1), where for one prompt, H = xtest " dl Xl i=1 yix⊤i 1l Xl i=1 y2i (5) For n training prompts, write γ = vec(Γ) and hµ = vec(Hµ). Ridge regression gives γ∗ = n/d λId(d+1) + Xn µ=1 hµh⊤ µ!−1 Xn µ=1 yµtesthµ (6) The test error is EICL(Γ∗) = Enew h(y new test − ⟨Hnew, Γ∗⟩F) 2i (7)

Theoretical Results and Analysis

Result 1 states that if the link function f has a nonzero linear mode, a1 ≠ 0, then e nonlinear ICL (Ctrain, Ctest, ρ) = e linear ICL (a 2 1Ctrain, a2 1Ctest, ρnonlinear) If instead a1 = 0, then e nonlinear ICL = ρnonlinear if τ 1 if τ > 1 This implies that linear attention learns the linear mode, while the other modes appear as noise The new effective noise is computed as (a 2 1 + ρnonlinear)/a 2 1 to balance the a 2 1 scaling of the task covariances in Result 1

Phase Transition in Task Diversity

If we take Ctrain = Ctest = Id for simplicity, when κ > 1 the error approaches the new noise floor (so long as a1 ≠ 0) eICL(Γ∗) −→ ρnonlinear (11) Nonlinear functions raise the limiting ICL test error, even in the best possible data case of α, κ, τ → ∞ due to nonlearnable signal captured as noise

Conclusion

We have extended the asymptotic theory of in-context learning by a reduced linear-attention block from linear regression tasks to nonlinear single-index targets The resulting equivalence shows that the model uses the target function only through its linear Hermite component, while the remaining Hermite modes contribute as effective noise Increasing the amount of pretraining data or the diversity of pretraining tasks cannot remove the error floor created by nonlinear components that are invisible to this predictor This limitation persists beyond the task-diversity transition, where the error approaches ρnonlinear rather than the original label-noise level ρ

How it works

The core mechanism relies on decomposing the target function f into its Hermite expansion to show that linear attention cannot learn beyond the linear Hermite mode of the link function The derivation shows that E[y 2 i] ≈ ρf, (12) The conditional expectations lead to a selfconsistent asymptotic expression for the extended resolvent G(π) ≈ τ 1 + χπ ρnonlinear 0 0 Id ⊗ B + π 0 0 0 Id ⊗ B + τλI−1 χπ = (see Section C) The final error formula EICL → ρnonlinear is obtained by substituting the derived selfconsistent expression into the formula of Letey et al. (2026b)

Population Averages

Lemma 1 provides conditional expectations of order-two quantities, such as E[y 2 i] ≈ ρf, (12) The proof involves decomposing the input x into a component aligned with w and its orthogonal component The calculation of the scalar block E s squared = ρ + X∞ m=0 a 2 m!2 + 1/l [Ef 4(g) + 6ρX∞ m=0 a 2 m + 3ρ squared - ρ + X∞ m=0 a 2 m!2] is derived The derivation of the ICL formula follows directly from comparing equations (12), (13), and (14) with Lemma 2 of Letey et al. (2026b)

Gaussian Equivalent for w⊤x

To expand in the basis of Hermite polynomials, we must first understand the distribution of this random variable We know that, conditional on fixed task w, w⊤x d= √σwg for σw = 1/d w squared, g ∼ N (0, 1) independent of x, w We will assume two conditions on C: Scaling freedom and Sufficient effective rank

Task Ratio Transition

Previous work has found that taking α, τ → ∞ with τ ∝ α gives a sharp transition at κ = rank(Ctrain)/d This corresponds to a transition between the model memorising the pretraining tasks and actually generalising to the correct task structure for good performance on Ctest The error approaches ρnonlinear rather than the noise floor ρ obtained for a linear target when κ > 1

Conclusion

The paper concludes that increasing pretraining data or task diversity cannot remove the error floor created by nonlinear components that are invisible to this predictor This limitation persists beyond the task-diversity transition, where the error approaches ρnonlinear rather than the original label-noise level ρ The study provides a tractable starting point for studying nonlinear in-context learning

Acknowledgments and Disclosure of Funding

M.L.

Improvements for AI systems

  1. System can be adapted to extract only linear structure from nonlinear in-context learning targets by treating all non-Hermite modes as effective noise. This limits the model's learned function to only its linear Hermite component when no linear mode is present, effectively preventing the model from learning beyond a simple regression task.

  2. The system can be used to rigorously quantify the generalization error floor for models using linear attention by relating it to a modified noise level, as shown by the equivalence: linear attention extracts only the linear Hermite component of f, while the remaining nonlinear structure contributes to the generalisation error as effective noise.

  3. For scenarios where a model is trained on insufficient or low-diversity data, this framework predicts that increasing pretraining data or task diversity cannot remove the error floor created by nonlinear components that are invisible to this predictor, suggesting a fundamental limitation in current in-context learning architectures.

Abstract

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)+ epsilon. 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.

Sources

Related papers