Transformers as Statisticians

One attention layer takes one gradient step, one head per ReLU piece of the loss derivative. The query carries the current iterate $w$ (stored in every token), the key carries $(x_j,y_j)$, and the value the direction $x_j$: head $m$ scores $a_m\langle w,x_j\rangle+b_my_j+d_m$, with a key term that pushes the score below zero on the query token, and writes $-\eta c_mx_j$ into the $w$-slot. Summed over heads the ReLUs rebuild the derivative, $$\sum_m c_m\,\sigma\big(a_m\langle w,x_j\rangle+b_my_j+d_m\big)\approx\partial_s\ell(\langle w,x_j\rangle,y_j),$$ so the residual update is $w-\frac\eta N\sum_j\partial_s\ell(\langle w,x_j\rangle,y_j)\,x_j$ up to an error of norm $\varepsilon\eta B_x$. This determines the head count. The squared loss needs two heads, since $s-y=2\sigma((s-y)/2)-2\sigma(-(s-y)/2)$, an $\ell_2$ penalty one more, and a smooth loss as many as its sum-of-ReLUs approximation. A final 2-head layer reads out $\langle\hat w,x_{N+1}\rangle$ via $t=\sigma(t)-\sigma(-t)$.

Errors add linearly across layers because convex gradient descent is non-expansive. Stacking gives a perturbed iteration $\hat w^{\ell+1}=T_\eta(\hat w^\ell)+\varepsilon_\ell$ with $T_\eta(w)=w-\eta\nabla f(w)$, and co-coercivity of a convex $L_f$-smooth $f$ shows that $T_\eta$ never moves two points apart: $$\|T_\eta w-T_\eta w'\|^2\le\|w-w'\|^2-\eta\big(\tfrac2{L_f}-\eta\big)\|\nabla f(w)-\nabla f(w')\|^2\le\|w-w'\|^2\quad(\eta\le2/L_f).$$ Each layer passes the earlier error on without amplifying it and adds its own, so $\|\hat w^L-w^L_{GD}\|\le L\varepsilon$; the same argument around the minimizer keeps the iterates in the ball where the ReLU approximation is valid. Without convexity this fails: for GD on a two-layer network the construction needs two layers per step and the error compounds exponentially.

Ridge, least squares and GLMs are gradient descent run for enough steps, and their risks follow by conditioning on a well-conditioned prompt. On a well-conditioned prompt the ridge objective is $(\alpha+\lambda)$-strongly convex and $(\beta+\lambda)$-smooth, so with $\eta=1/(\beta+\lambda)$ gradient descent contracts like $e^{-t/\kappa}$ and $2\kappa\log(B_xB_w/2\varepsilon)$ steps reach accuracy $\varepsilon/B_x$; the squared loss is exactly a sum of ReLUs, so the transformer adds no error. GLMs are the same argument with $g(s)-y$ approximated by ReLUs, which is where the $\varepsilon^{-2}$ heads come from. Designs with $N\gtrsim d$ are well conditioned except with polynomially small probability, where the clipped output adds little risk, so the classical risks transfer.

Lasso uses proximal gradient, with the MLP doing the soft threshold. Proximal gradient alternates a gradient step on the squared loss (two heads) with the $\ell_1$ prox, and soft thresholding is piecewise linear, $$\mathrm{st}_\tau(u)=\sigma(u-\tau)-\sigma(-u-\tau)\quad\text{coordinatewise},$$ so a token-wise MLP of width $2d$ implements it exactly. That gives the generic $\beta B_w^2/\varepsilon$ steps without strong convexity, which fails when $d\gt N$. For a depth that does not grow with $1/\varepsilon$, restricted strong convexity and smoothness of the Gaussian design on sparse directions make proximal gradient converge linearly to within the statistical error (as in Agarwal et al.), and the Lasso risk transfers.

Post-ICL validation takes three layers, which score, rank and mix the candidates. Tokens carry $t_i=1$ for training, $-1$ for validation and 0 for the query, and the $K$ base algorithms run in parallel on the training split. An attention layer scores all $K$ on the validation tokens, with $MK$ heads approximating the loss as a sum of ReLUs. An MLP ranks them: $c_k=\sum_{l\ne k}\sigma(L_k-L_l)$ is at most $\gamma$ only when $L_k\le\min_lL_l+\gamma$, and $u_k=\sigma(1-c_k/\gamma)\in[0,1]$ flags near-minimizers. A last attention layer mixes, $$\hat f=\sum_k\lambda_kf_k,\qquad \lambda_k=\sigma\Big(1-\sum_{l\lt k}u_l\Big)-\sigma\Big(1-\sum_{l\le k}u_l\Big),$$ so $\lambda_k$ is the part of a total weight of one left after the earlier flagged predictors have received theirs. The weights are nonnegative and sum to one, a soft version of the rule "take the first near-minimizer", and convexity of the risk turns this into the oracle inequality.

Validation over $K$ ridge penalties is nearly Bayes-optimal for mixed noise levels, with excess risk of order $N^{-1/3}$. Since ${\rm BayesRisk}_\pi\ge\mathbb E_k{\rm BayesRisk}_k$, it suffices to compete with the oracle that knows $k$: fit all $K$ ridges, $\lambda_k=d\sigma_k^2/N$, on $N-N_{\rm val}$ points and pick one. The Bayes predictor averages the same ridges, but its weights $P(k\mid D)$ involve $\sigma_k^{-(N-d)}\det(X^\top X+d\sigma_k^2I)^{-1/2}$ and a data-dependent exponent, which attention does not compute easily. The excess comes from the lost training points, $O((\sigma_k^2+1)N_{\rm val}/N)$, and from the uniform deviation of $K$ validation losses, $O(B_w^2\sqrt{\log K/N_{\rm val}})$, and balancing them sets the split: $$\frac{N_{\rm val}}N+\sqrt{\frac{\log K}{N_{\rm val}}}\quad\text{is minimized at}\quad N_{\rm val}\asymp N^{2/3}(\log K)^{1/3},\ \text{giving}\ (\log K/N)^{1/3}.$$ The transformer commits to the noise level that validates best instead of carrying a posterior over it.

Pre-ICL testing computes a statistic of the prompt first and branches on it. To tell classification from regression, average a label test over the prompt, $$\Psi_{\rm binary}(D)=\frac1N\sum_i\psi(y_i),$$ with $\psi$ a pair of tents of width $\varepsilon$ at 0 and 1. Since $\psi$ is a combination of six ReLUs of $y$, one 6-head attention layer computes $\Psi_{\rm binary}$ exactly at every token. Another layer thresholds it into a gate in $[0,1]$, and the network outputs the gated combination of logistic regression and least squares: for binary labels the gate is exactly 1, and for labels not concentrated near $\{0,1\}$ it is 0 with high probability. This is Proposition 19.

Pretraining needs polynomially many prompts because transformers are Lipschitz in their weights. With clipping after each layer, the output on bounded inputs is Lipschitz in $\theta$ with constant $LB_H^{L-1}B_\Theta$, where $B_H,B_\Theta$ are polynomial in $B$ and $R$, and the norm ball has log covering number $L(3MD^2+2DD')\log(1+2r/\delta)$. Chaining for the bounded loss process then gives $$\sup_\theta|\hat L_{\rm icl}-L_{\rm icl}|\lesssim B_y^2\sqrt{(L(MD^2+DD')\iota'+\log(1/\xi))/n};$$ the Lipschitz constant is exponential in $L$ but enters only through the log, giving the extra factor $L$ in Theorem 20. The construction only certifies that the comparator $\inf_\theta L_{\rm icl}(\theta)$ is small, and the result says nothing about what SGD finds.

Trained transformers behave as if they select an algorithm for each prompt. A 12-layer ReLU-attention encoder matches the best baseline on noiseless, noisy and sparse regression and on linear classification. Trained on a mixture of two noise levels, or of regression and classification, one model approaches the best baseline on each component, while single-task models fall short on the other task; a GPT-2 decoder does the same. This is behavioral evidence only; the trained weights are not probed for validation or testing.

Consequences and limits. "The transformer runs gradient descent" becomes quantitative: on well-conditioned convex problems the depth needed grows like the logarithm of the inverse accuracy, the number of heads sets how accurately the loss derivative is represented, polynomially many prompts suffice to find the weights, and selection by validation is nearly Bayes-optimal on a noise mixture without representing the posterior. The limits are that the constructions are existence results for ReLU attention, and the paper does not show that softmax transformers trained by SGD find them; the in-context guarantees need a well-conditioned prompt ($N\gtrsim d$ for least squares), so the interpolation regime $N\lt d$ of the benign-overfitting papers is handled only by Lasso and ridge; the validation route's excess decays only like $N^{-1/3}$; and the pretraining bound is for constrained ERM and grows with $L^2$.