Transformers Learn In-Context by Gradient Descent

A gradient step is a change of targets, and one attention layer makes it (von Oswald et al.). Push the update from the weights into the data: for $\Delta W=-\frac\eta N\sum_i(Wx_i-y_i)x_i^\top$, $$L(W+\Delta W)=\frac1{2N}\sum_i\big\|Wx_i-(y_i-\Delta Wx_i)\big\|^2,$$ so a layer that subtracts $\Delta Wx_j$ from every token's $y$-entry, the query included, has taken the step without ever storing $W$. With $W_K^\top W_Q$ picking out $x_i^\top x_j$, $W_V$ sending $e_i$ to $(0,W_0x_i-y_i)$ and $P=\frac\eta NI$, the attention sum is exactly that update; starting the query at $(x_{\rm query},-W_0x_{\rm query})$ and negating its $y$-entry at the end reads out $(W_0+\Delta W)x_{\rm query}$. From $W_0=0$ the prediction is $\frac\eta N\sum_iy_ix_i^\top x_{\rm query}$, Garg's averaging estimator up to $\eta$.

Stacking gives more steps, and one more block adds curvature correction. After a layer the context tokens hold the residuals $y_i-\Delta W_1x_i$, the data the updated model faces, so $K$ layers are $K$ GD steps. GD++ also moves every input, the query included, by $x\leftarrow(I-\gamma XX^\top)x$, which replaces the Hessian $H=XX^\top$ by $$H^{++}=(I-\gamma H)H(I-\gamma H),\qquad\lambda\mapsto\lambda(1-\gamma\lambda)^2 .$$ The gain depends on $\gamma$: for $\gamma\le1/(3\lambda_{\max})$ the map is increasing on the spectrum and the condition number falls, larger $\gamma$ maps the top of the spectrum to values below the peak at $1/(3\gamma)$, which can help more, and $\gamma$ near $1/\lambda_{\max}$ makes it worse. Repeated over layers it is an iterative preconditioner computed from the prompt.

An MLP supplies features, and an earlier head builds the tokens. After a residual MLP $m$ the same layer takes GD steps on kernel regression with $k(x,x')=m(x)^\top m(x')$. When inputs and targets sit in separate tokens, a head whose key-query product selects the successor's positional encoding copies each $y_j$ into its $x_j$ token, as an induction head does, and the next layer takes the GD step, as trained two-layer models do. Through a softmax the copy is only approximate (weight $e/(e+2N)$ on the successor among $2N+1$ tokens), so the key-query product must be scaled up.

Trained models match the construction (experiments). A trained single LSA layer matches one tuned GD step in loss, predictions and input sensitivities, its rescaled matrices visibly match the construction, and the midpoint of trained and constructed weights has the same loss. The two fail together: for inputs from $U(-\alpha,\alpha)$ with $\alpha\gt1$ both blow up, because a step size tuned for $\alpha=1$ overshoots the higher curvature. The match needs task diversity (it fails at 128 training tasks and is nearly perfect above 2048). Deeper models drift toward GD++, but only with fitted $\eta,\gamma$ per layer, which is weaker evidence.

The one-layer loss is a sum of convex functions of rank-one matrices, and a convex relaxation solves it (Ahn et al.). For Gaussian inputs the global minimizer predicts $x_{\rm query}^\top\Gamma^{-1}\frac1N\sum_iy_ix_i$, one step from zero of GD preconditioned by $\Gamma^{-1}$; for $\Lambda=I$, $\Gamma^{-1}=\frac N{N+d+1}I$, plain GD with a tuned step, von Oswald's construction. The loss is nonconvex, so the proof relaxes it. With the query masked, the prediction is $-b^\top GAx_{\rm query}$, where $G$ is the second-moment matrix of the context tokens, $b^\top$ the last row of the value-projection product and $A$ the first $d$ columns of the key-query product, and for $\Lambda=I$ the loss splits over the columns $a_j$ of $A$: $$f(b,A)=\sum_{j=1}^df_j(ba_j^\top),\qquad f_j(X)=\mathbb E\big[\langle G,X\rangle+w[j]\big]^2 .$$ Each $f_j$ is a convex quadratic in $X$; minimize it over all $X$ by Gaussian moments, and since the minimizers share one $b=e_{d+1}$, they solve the original problem (general $\Lambda$ follows by rotating). The $\frac{\mathrm{tr}\Lambda}N$ in $\Gamma$ comes from the fourth moment $\mathbb E[\|x\|^2xx^\top]$, the variance of $\frac1N\sum_iy_ix_i$ from having only $N$ samples, and it shrinks directions with $\lambda_i\lesssim\mathrm{tr}(\Lambda)/N$ as ridge would. The minimal loss is $\sum_i\lambda_i(\lambda_i+\mathrm{tr}\Lambda)/((N+1)\lambda_i+\mathrm{tr}\Lambda)$, which is $d(d+1)/(N+d+1)$ for $\Lambda=I$ (a short computation, not in the paper, checked by simulation), so one layer never interpolates.

With sparse layers the forward pass is preconditioned GD, and the loss is a matrix polynomial. Restrict each layer to a value product that writes only the label row and a key-query product $-\mathrm{diag}(A_\ell,0)$. The inputs never change, the mask makes the query's label enter only additively, and every context label follows the same recursion as the query's, so by induction it holds the residual of the current model: $$\theta_{\ell+1}=\theta_\ell-\tfrac1NA_\ell^\top\bar X\bar X^\top(\theta_\ell+w),\qquad w_\ell^{\rm gd}=-\theta_\ell,$$ with $\bar X$ the context inputs: $L$ steps of GD with per-layer preconditioners. Putting the true label in the query slot turns the output into the prediction error, so the loss is an explicit matrix polynomial in the $A_\ell$, built from $X\prod_\ell(I+\mathcal MX^\top A_\ell X)$ with $\mathcal M$ the mask; the remaining results analyze that polynomial.

For two layers some global minimizer is diagonal, by rotation and random sign flips. For $L=2$, isotropic data and symmetric $A_1,A_2$, some pair of diagonal matrices is a global minimizer, so the best two-layer model is GD with coordinatewise step sizes. Take a global minimizer, diagonalize $A_1^*$ by rotation invariance, and conjugate $A_2$ by a random diagonal sign matrix $S$: the loss is unchanged ($SD_1S=D_1$, the data are sign-invariant) and convex in $A_2$, because the output is affine in $A_2$. Since $\mathbb E_S[SA_2S]=\mathrm{diag}(A_2)$, Jensen gives $$f(D_1,\mathrm{diag}A_2^*)\le\mathbb E_Sf(D_1,SA_2^*S)=f(A_1^*,A_2^*) .$$

At symmetric points the gradient can be matched inside the symmetric set, so it contains critical points. With prior $w\sim N(0,\Lambda^{-1})$, Theorem 3 says the set $S=\{A_\ell=a_\ell\Lambda^{-1}\}$, GD preconditioned by the Newton matrix, contains critical points, possibly approached only in the limit; Theorem 4 adds input updates $b_\ell I$, which give GD++ when $\Lambda=I$; the input update commutes with the symmetry used below only when it is a multiple of $I$. The data law is invariant under $X\mapsto U_\Lambda X$, $U_\Lambda=\Lambda^{1/2}U\Lambda^{-1/2}$ with $U$ Haar orthogonal, which fixes points of $S$, so perturbing a layer in direction $R$ changes $f$ as perturbing it in direction $U_\Lambda^\top RU_\Lambda$, and by linearity in the direction one may average: $$\mathbb E_U[U_\Lambda^\top RU_\Lambda]=\tfrac1d\mathrm{tr}(\Lambda^{1/2}R\Lambda^{1/2})\,\Lambda^{-1}.$$ Every directional derivative at a point of $S$ is thus matched by one inside $S$, so a flow constrained to $S$ decreases $f$ at rate at least the full squared gradient norm, and since $f\ge0$ the gradient cannot stay away from zero. This shows existence and says nothing about optimality, but Adam finds these points: in a trained three-layer model $\Lambda^{1/2}A_\ell\Lambda^{1/2}$ converges to multiples of $I$, and with input updates the whitened step sizes grow across layers (about $-0.38,-2.40,-9.59$) as the conditioning improves.

With ReLU attention the minimizer is again a multiple of the identity, by sign symmetrization and Stein's lemma. With ReLU applied to the attention scores, isotropic data and the last row of the key-query product zero, the same rotation, sign-flip and Jensen moves bound the loss below by a function of $b_1A$, and Stein's lemma, $\mathbb E[\sigma(x^\top v)x]=v/2$, evaluates every term: $$f_{\rm lower}=\frac{N-1+2(d+2)}{4N}\,\|b_1A\|_F^2+\mathrm{tr}(b_1A)+\text{const},$$ a convex quadratic minimized at $-\frac{2N}{N-1+2(d+2)}I$, where the bound is attained. ReLU changes only the constant.

The prediction is a rank-one factorization, and off the cross entries the loss reduces to two blocks (Zhang et al.). From a small balanced start, $W^{PV}(0)=\sigma\,\mathrm{diag}(0,1)$ and $W^{KQ}(0)=\sigma\,\mathrm{diag}(\Theta\Theta^\top,0)$ with $\|\Theta\Theta^\top\|_F=1$, gradient flow on the population loss converges to a global minimum whenever $\sigma^2\|\Gamma\|_{\rm op}\sqrt d\lt2$ (Theorem 4.1). The prediction is an indefinite quadratic form $u^\top H_\tau u$ in the vector $u$ of the weights that affect it, value weight times key-query weight, so the dynamics are those of a two-layer linear network. The cross entries stay at zero, since every term in their derivatives carries an odd power of the mean-zero task, and $\Gamma$ enters through the Wishart fourth moment of $\hat\Lambda=\frac1N\sum_ix_ix_i^\top$. Up to a constant the loss in the key-query block $U_{11}$ and the value entry $u_{-1}$ is $$\tilde\ell(U_{11},u_{-1})=\tfrac12\mathrm{tr}\big(u_{-1}^2\,\Gamma\Lambda U_{11}\Lambda U_{11}^\top-2u_{-1}\Lambda^2U_{11}^\top\big).$$

Completing the square gives the minima, and conserved balance picks one. Since $\Gamma$ and $\Lambda$ commute, $$\tilde\ell-\min\tilde\ell=\tfrac12\Big\|\Gamma^{1/2}\big(u_{-1}\Lambda^{1/2}U_{11}\Lambda^{1/2}-\Lambda\Gamma^{-1}\big)\Big\|_F^2,$$ so the global minima are exactly the hyperbola $u_{-1}U_{11}=\Gamma^{-1}$, all giving the same predictor. Multiplying the two flow equations by $U_{11}^\top$ and $u_{-1}$ gives $\frac d{dt}\|U_{11}\|_F^2=\frac d{dt}u_{-1}^2$, the conservation law of two-layer linear networks, so balance $u_{-1}^2=\|U_{11}\|_F^2$ holds for all time and pins the limit on the hyperbola at $u_{-1}=\|\Gamma^{-1}\|_F^{1/2}$, $U_{11}=\|\Gamma^{-1}\|_F^{-1/2}\Gamma^{-1}$.

A small start avoids the saddle at zero, and then a PL inequality gives convergence. The origin $u_{-1}=0$ is a saddle with $\tilde\ell=0$, and the condition on $\sigma$ makes $\tilde\ell(0)\lt0$ (von Neumann's trace inequality), so $u_{-1}$ never reaches 0; with balance it stays above some $c_0\gt0$. Keeping only the $U_{11}$-gradient and comparing with the completed square, $$\|\nabla\tilde\ell\|^2\ge\mu\,(\tilde\ell-\min\tilde\ell),$$ so the excess loss decays like $e^{-\mu t}$, the flow has finite length, and the limit is a global minimum.

The error is a bias-variance split for a fixed linear smoother. For a test prompt of length $M$ from any distribution with $x\sim N(0,\Lambda)$, $a$ the best linear predictor, $\Sigma=\mathrm{Cov}(xy)$ and $\hat g=\frac1M\sum_iy_ix_i$, $$y_{\rm query}-\hat y_{\rm query}=\big(y_{\rm query}-x_{\rm query}^\top a\big)+x_{\rm query}^\top(I-\Gamma^{-1}\Lambda)a-x_{\rm query}^\top\Gamma^{-1}(\hat g-\Lambda a),$$ with vanishing cross terms. The last term is a variance $\frac1M\mathrm{tr}(\Sigma\Gamma^{-2}\Lambda)$, and $I-\Gamma^{-1}\Lambda=\frac1N\Gamma^{-1}(\Lambda+\mathrm{tr}(\Lambda)I)$ gives a squared bias of order $1/N^2$ (Theorem 4.2). As for ridge, eigendirection $i$ keeps the fraction $N\lambda_i/((N+1)\lambda_i+\mathrm{tr}\Lambda)$ of the signal, a shrinkage level $\mathrm{tr}(\Lambda)/N$ fixed at training, which does not vanish as $M$ grows. For noiseless tasks with $w\sim N(0,I)$ and $\kappa$ the condition number of $\Lambda$, the error is at most $\frac{(d+1)\mathrm{tr}\Lambda}M+\frac{(1+2d+d^2\kappa)\mathrm{tr}\Lambda}{N^2}$. The $1/M$ term stays even for noiseless labels and $M\gg d$, because the model applies a fixed $\Gamma^{-1}$ where least squares would invert $\hat\Lambda$ and be exact once $M\ge d$.

The preconditioner fixed at training determines which distribution shifts the model tolerates. Any labels give the best linear predictor up to $O(1/M+1/N^2)$, so label noise and arbitrary $w$ are tolerated, and for small $M$ and isotropic $\Lambda$ a query orthogonal to all $x_i$ gets 0, as Garg et al. saw. Scaling the test inputs by $c$ gives $\hat y\approx c^2x_{\rm query}^\top w$, Garg et al.'s input-scale failure, while least squares is unaffected. Training on random diagonal covariances, one per prompt, does not fix it, since one layer cannot form $\hat\Lambda^{-1}$ and learns one fixed preconditioner: as $M,N\to\infty$ it multiplies coordinate $i$ of $w$ by $\lambda_{{\rm new},i}\,\mathbb E[\lambda_i^2]/\mathbb E[\lambda_i^3]$, which for Exponential(1) variances is $\lambda_{{\rm new},i}/3$, right only when $\lambda_{{\rm new},i}=3$.

Consequences and limits. For one linear attention layer the analysis is complete: the trained model is $\hat w=\Gamma^{-1}\frac1MX^\top y$, a linear estimator whose risk follows from the spectrum of $\Lambda$. Training length acts as regularization, so test error has a floor of order $1/N^2$ however long the test prompt, and matching least squares in distribution does not make the model least squares: the learned preconditioner matches Garg et al.'s fragility to input scale and covariance, and only computing $\hat\Lambda^{-1}$ from the prompt, which needs depth or nonlinearity, would remove it. The limits are linear attention, noiseless Gaussian training tasks, gradient flow on the population loss from one initialization, critical points without global optimality for depth, and softmax and LayerNorm variants that match less cleanly.