nucleation

LIT-859: Exact Learning Dynamics of In-Context Learning in Linear Transformers and Its Application to Non-Linear Transformers

3 min read

Properties4
StatusActive — read 2026-10-09 (NOTE-658); worth reading as an exactly solved case of the slow learning that produces a fast in-context learner: for one linear-attention layer trained by gradient flow on in-context vector regression, with tasks and inputs sharing one eigenbasis, each eigenmode of the input covariance learns by its own logistic curve on a timescale ∝ 1/s_α², the trained product of the two weight blocks settles at the inverse of a finite-context-corrected training covariance, so the learned in-context predictor is a preconditioned estimator (the fixed point Ahn et al. and Zhang et al. had found), and ‖p₂‖² − ‖q₁‖² is conserved. The decoupling rests on an assumed spectral alignment and on large N and many tasks. The second half, measures applied to attention-only transformers and a grokking model, is qualitative and correlational: no seeds, error bars or controls are reported, and "mechanistic explanation" overstates it.
Tagslearning-theory, representation-learning, anthology-candidate
Cited as “paper” byNOTE-658 — Exact Learning Dynamics of In-Context Learning in Linear Transformers
Cited as “source” byTHEORY-481 — Under gradient flow averaged over tasks and data, a linear-attention layer trained on in-context regression learns as a two-layer linear network from an aligned start does: when tasks share the input covariance's eigenbasis and the weights start and stay aligned with it, then in the long-context limit each input mode is learned on its own logistic curve whose timescale falls with the square of its variance, so learning is staged by the input spectrum, and the layer converges to the estimator preconditioned by the inverse input covariance; at finite context length the modes decouple exactly only when inputs are isotropic and the modes start equal, and otherwise couple at order 1/N

Graph View

Nischal Mainali and Lucas Teixeira (2025), NeurReps Workshop 2025 (arXiv v2) — ARXIV-2504.12916

Key takeaways

  • The model. One layer, one head, attention without softmax: f(Z) = Z + W^P (ZZᵀ/N) W^Q Z, with tokens (x_i, y_i) stacked as columns and the query padded with zeros. Under task averaging, the blocks p₁ and q₂ have zero expected gradient at zero, so the prediction reduces to ŷ ≈ p₂ W^µ Σ̂_x q₁ x_q (Eq. 8). Inputs x ~ N(0, Σ_x), Σ_x = USUᵀ; tasks W^µ = UΛ^µUᵀ with Gaussian eigenvalues and the same U; MSE on the query; small-step SGD over a fixed set of P tasks, taken to the continuous limit.
  • Exact dynamics, by assumption of alignment. Assuming p₂ and q₁ stay diagonal in U, gradient flow gives, per mode, τ_α dp_α/dt = q_α(1 − p_α q_α s^∞_α) and the same with p, q exchanged, with τ_α = (ηP s_α²)⁻¹ and s^∞_α = ((N+1)s_α + Tr S)/N. From a balanced start, a_α = p_α q_α obeys a logistic equation with closed-form solution (Eq. 13); the expected loss is an explicit sum over modes (Eq. 14), and the predictions match simulation closely (Figs. 1–3A).
  • Staged learning. Modes with larger input variance are learned first, the time to escape a small initialisation ε being about (τ_α/2) log(1/(s^∞_α ε)). This is the Saxe et al. (2014) deep-linear phenomenology, here driven by the data spectrum through the attention product, not by depth.
  • What is learned. At the fixed point p₂q₁ = s^∞(S)⁻¹ the predictor is ŷ = W^µ Σ̂_x [((N+1)/N) Σ_x + (Tr Σ_x/N) I]⁻¹ x_q (Eq. 11): the context’s empirical covariance against an inverse training covariance with a finite-context correction, i.e. a preconditioned estimator. It tends to W^µ x_q only as N → ∞, and the converged loss is non-zero at finite N. The fixed point is the one already found by Ahn et al. (2023) and by Zhang, Frei and Bartlett; what is new is the trajectory to it, for vector-valued outputs.
  • A conservation law. ‖p̄₂‖²_F − ‖q̄₁‖²_F is constant (indeed each p_α² − q_α² is), from the loss’s invariance under p → cp, q → q/c.
  • The non-linear half is suggestive only. Loss-autocorrelation and norm curvature, a “marginalised effective rank” and a subspace distance min_A ‖A M(t) − M(∞)‖ are applied to 1–4-layer attention-only transformers and to a modular-arithmetic grokking model. Dips in QK effective rank and early subspace stabilisation coincide with in-context learning’s emergence in the 2+-layer models; in the grokking model the OV subspace stabilises late, with the test-error drop. These are timing coincidences, with no seeds, error bars or controls reported.

Standing in the record

Filed here at the owner’s request on 2026-10-09, with no stated context, in a batch with arXiv 2509.24914, 2502.09863, 2205.10343 and 2410.17770, and read on its own merits (NOTE-658). The anthology does not hold it (clone of commit d8b5ba5, possibly stale). It is chiefly a result about the training of transformers and in-context learning, which the anthology’s topics can hold, so it carries anthology-candidate.

Here it is the exactly solved account of the outer learning that produces the inner, in-context learner of LIT-792 and THEORY-169: for isotropic inputs the trained fixed point is that construction’s scaled gradient step, and for anisotropic inputs it is a preconditioned one. Its staged, mode-by-mode learning is one more kind of training “phase” for THEORY-039, and the same eigenmode-ordered stepwise learning that LIT-242 reports for self-supervised learning.

Backlinks