Attention is actually all you need
On brilliance through selective ignorance
2017-12-20 — 2022-08-05
Wherein the Structure of Transformer Stacks and Self‑attention Layers Is Described, the Role in Processing Sequential Data Such as Text Is Examined, and Recent Optimizations Such as FlashAttention Are Noted.
Bookmarked: A short essay in taking attention absolutely seriously.
To a first order approximation, we can think of the transformer architecture as one that chooses its outputs based on regressing against certain features of its inputs, which is to say, we can think of the attention architecture as “choosing which predictors” and then regressing off them to predict the logits of the next token. That is to say, it looks to me like we can think of next token prediction as logit regression with learned input feature selection (This leave the residual stream somehow under-regarded, but let’s shelve that for the moment.) And insofar as regression is, approximately, smoothing over input samples, I wonder if this gives us an angle of attack to understand how to interpret what kind of learning such models are capable of. Some claim LLMs are “mere interpolators”), but I wonder if this claim is kinda degenerate , i.e. strictly true but vacuous, if we understand them as interpolating in this special sense where they choose how they interpolate. As I have given it to you, this word vomit is a little vague. I want you to search for some good operationalizations of it, however, and see if you can devise a minimal model of transformers as optimal-context-choosing-regressors. We can be very minimal here, as long as we resemble the transformer architecture. in particular, we can skip out the training process, treat NNs as generic function approximators. We can discretize attention if we want or skip out untidy details of KVQ operations in practice. But I wonder, does this get us an interesting way of analyzing these algorithms.
A softmax attention head is a Nadaraya–Watson kernel smoother over the context, with a learned, asymmetric kernel. The values enter linearly; all the nonlinearity is in the weights, which makes a transformer block a varying-coefficient regression whose coefficients are chosen per query. Pretraining learns the kernel; inference regresses over the context. “Mere interpolation” is a theorem for one softmax head in value space, false for the block, and in any case mislocated, since the thing interpolated is the prompt, not the corpus. The constraints that do bind are three: values are local, selection is pairwise, aggregation is convex. What I want to measure is the gap between the best selection the architecture could make and the one it learned.
1 Operationalizations already on the shelf
The hunch is not new, and several people have made it precise in different ways. These are pointers, not a review.
The oldest and most literal is attention as kernel regression. Write a head as \[ \operatorname{Att}(q; K, V) = \sum_j \frac{\exp(q^\top k_j/\sqrt d)}{\sum_{j'} \exp(q^\top k_{j'}/\sqrt d)}\, v_j . \] That is the Nadaraya–Watson estimator with kernel \(\kappa(q, k) = \exp(q^\top k/\sqrt d)\): the keys are the training inputs, the values the training targets, the query a test input. (Tsai et al. 2019) wrote this down and asked which other kernels are worth swapping in; (Santos et al. 2026) extend it to the sparse attentions (sparsemax, entmax, top-\(k\)) as compact-support kernels; the D2L textbook teaches attention this way from the first page. The kernel is asymmetric and bilinear rather than a function of distance, and it is learned, which is where the “choosing” lives.
Second, the varying-coefficient model. \[ y = \sum_j \beta_j(x)\, v_j, \] with \(\beta\) a nonlinear function of the input and the \(v_j\) entering linearly, is a varying-coefficient model in the sense of (Hastie and Tibshirani 1993). Attention is one with the coefficients constrained to the simplex and computed pairwise.
Third, the QK/OV decomposition of Elhage et al. (2021). For a one-layer attention-only model the logits at position \(i\) are \[ z_i = \underbrace{W_U h_i}_{\text{direct}} + \sum_{j\le i} a_{ij}\,\underbrace{W_U W_O W_V h_j}_{\ell_j}, \] a mixture of per-token votes \(\ell_j\) with mixture weights set by the query. Their QK circuit is the selection; their OV circuit is the regression coefficient attached to each predictor. I had not noticed before that their cut and the regression cut are the same cut.
Fourth, RASP, from Weiss, Goldberg and Yahav (2021) (Weiss, Goldberg, and Yahav 2021), which is the discretized version. Its two primitives are select (a boolean predicate on query–key pairs) and aggregate (average the selected values). Programs in it compile to transformers, so whatever we can say about select-then-average we can say about the architecture.
Fifth, in-context learning as regression, where the context is literally a dataset of \((x, y)\) pairs. (Garg et al. 2023) set the task up; (Akyürek et al. 2023) and (von Oswald et al. 2023) show that linear attention can implement a step of gradient descent on least squares over the context, so that depth is iterations. (Wang, Shi, and Fox 2025) go furthest: they recast softmax attention, linear attention, state-space models and fast-weight programmers as test-time regression over the context, differing only in the regressor class, the sample weights and the optimizer. Softmax attention is a local-constant regressor in their taxonomy; linear attention and its relatives are parametric regressors fitted by gradient steps. That paper is the closest existing thing to the model I want, and some of what follows re-derives things it contains, with a different question in mind.
Sixth, the critique the hunch is aimed at. Balestriero, Pesenti, and LeCun (2021) take the convex-hull definition of interpolation seriously and show that in more than a hundred dimensions a test point is almost never inside the hull of the training set, so by that definition nothing high-dimensional ever interpolates. (Richter and Wattenhofer 2020) see the same hull from the other side, as a cage on what a softmax head can output in value space, and propose dropping the normalization to escape it.
2 A minimal model
Here is the smallest thing I can write that keeps the shape of the architecture.
A context is \(x_1, \dots, x_n\). Each position carries a feature vector \(h_j \in \mathbb{R}^d\), initialized from the token and position embeddings; this is the residual stream. One layer, at position \(i\), does \[ h_i \leftarrow h_i + \sum_{j \le i} a_{ij}\, v(h_j), \qquad a_{i\cdot} = \sigma\big(s(h_i, h_1), \dots, s(h_i, h_i)\big), \] with \(a_{i\cdot}\) a probability vector over the positions \(j \le i\), and then \(h_i \leftarrow h_i + m(h_i)\). Logits at the end are \(W_U h_n\).
Three learned pieces: a selection score \(s\), a value map \(v\), a pointwise map \(m\). We take all three as generic function approximators and skip the question of how they get trained. One fixed piece: \(\sigma\), which turns scores into weights on the simplex. Softmax at temperature \(\tau\) gives the standard model; argmax gives hard attention; uniform-over-the-argmax-set gives RASP’s aggregate. Multi-head is \(H\) copies of the first line, summed.
Everything the architecture forbids is in three constraints. Values are local: \(v\) sees one position. Selection is pairwise: \(s\) sees a query and one key. Aggregation is convex: the weights live on the simplex. Everything else is free. If we relax locality and let \(v\) see the whole context, the model is universal in one line, since \(v_n\) can compute anything and \(a\) can select it. The “interpolation” of the hunch is the third constraint only, and I will argue it is the least binding of the three.
3 Who interpolates what
The model has two timescales and the critique conflates them. Pretraining fits \(s\), \(v\) and \(m\) to the corpus. Inference runs a kernel regression over the context, with that \(s\) as the kernel and that \(v\) as the targets. Nothing in the forward pass interpolates the training corpus; what gets interpolated is the prompt.
What pretraining learns, in this view, is a similarity. The score \(s(h_i, h_j)\) says how much position \(j\) counts as the same situation as position \(i\) for the purpose of guessing what comes next. That is metric learning for kernel regression, as in (Weinberger and Tesauro 2007), with the metric chosen to minimize next-token loss averaged over the corpus’s contexts. The regression itself is dumb; the similarity is where whatever the model knows about the world sits.
The induction head is the worked example. After one layer has written the previous token into each position’s features, a second-layer head can learn \(s(h_i, h_j) \approx \mathbb{1}[x_{j-1} = x_i]\) and a \(v\) that copies the token. The kernel is relational rather than metric: it tests a relation between two context positions and is defined for tokens the model has never seen. The regression is trivial; the selection is a program. “Interpolating the context” with a relational kernel is how transformers copy, count and match brackets, and RASP is the language for saying what else it reaches.
This also sorts few-shot from many-shot in-context learning. With a match-type kernel a single precedent is enough, and the regressor is a lookup. With a smooth kernel the regressor averages, and we need many precedents to average over. Which regime we are in depends on the kernel the model learned, not on the prompt.
4 The hull
One softmax head’s output is a convex combination of its values, so it lies in \(\operatorname{conv}\{v(h_j)\}_{j \le i}\). That is the one place where “interpolation” is a theorem. It stops being one as soon as we leave the head.
With \(H\) heads the output is a sum of \(H\) convex combinations, which is the Minkowski sum of \(H\) hulls. Two heads that attend to values \(a\) and \(b\) output \(a + b\), outside \(\operatorname{conv}\{a, b\}\). The pointwise map \(m\) leaves the hull, since it is nonlinear, but only along directions the training distribution taught it; and by the Balestriero argument its input is almost surely outside the hull of the inputs it saw in training, so \(m\) is the component doing classical extrapolation, with the usual hazards. The unembedding is linear and unconstrained. Linear attention drops the normalization, so its weights are signed and unbounded and its output is \((\sum_j v_j k_j^\top)\, q\), a fast-weight memory rather than a smoother; that is why it can implement least squares and a softmax head cannot without help.
So: interpolate across positions, extrapolate within them. Attention moves information between positions and can only average it; the MLP transforms information at a position and can only do so pointwise. “Mere interpolator” is true of the first and false of the second, and the architecture is their composition.
5 An oracle and two gaps
Fix the value side of the one-layer vote model: the votes \(\ell_j\) and the direct term \(\ell_0\). At each position let an oracle choose the weights \(a\) on the simplex to minimize expected loss under the true next-token distribution \(p(\cdot \mid x_{1:i})\): \[ a^\star_i = \arg\min_{a \in \Delta} \; \mathbb{E}_{y \sim p(\cdot\mid x_{1:i})}\Big[\operatorname{LSE}\big(\ell_0 + \textstyle\sum_j a_j \ell_j\big) - \big(\ell_0 + \textstyle\sum_j a_j \ell_j\big)_y\Big]. \] Log-sum-exp is convex and its argument is affine in \(a\), so this is a convex program over the simplex, and cheap. Call the model run with oracle weights the oracle-selection regressor. Its loss sits between two others, \[ \text{Bayes} \;\le\; \text{oracle selection} \;\le\; \text{trained model}, \] and the two differences deserve names.
- hull gap
- The oracle-selection loss minus the Bayes loss. The price of being confined to mixtures of the votes the context makes available.
- selection error
- The trained model’s loss minus the oracle-selection loss. The price of not having learned the oracle’s choice.
The decomposition separates what the architecture cannot do from what training did not teach it. For copying tasks the hull gap is zero, since the right vote is in the context. For “output the sum of two numbers in the context”, with values that are the numbers, it is positive, because the sum is outside the hull; a second head closes it directly and an MLP closes it by rescaling the average. On natural text I do not know the Bayes loss, but the oracle that sees the realized next token is computable, and it bounds the loss from below for any selection rule with that value side, including rules that cheat. Running it on a trained one-layer model would tell us how much of that model’s deficit against a large one is the hull and how much is the selection.
My guess, which is the thing to test, is that with two or more heads, an MLP and an unembedding, the hull gap on text is near zero. If so the architecture is not an interpolator even in the convex sense, and the critique has nothing to hold on to. For deeper models the clean version fixes everything but the last attention layer and asks the same question of it; how to do it layer by layer I have not worked out.
6 Three things the view predicts for free
A kernel smoother has well-known behaviours, and three of them are known transformer facts.
Iterated smoothing oversmooths. Each attention layer multiplies the values by a positive row-stochastic matrix, and products of positive row-stochastic matrices converge to a rank-one matrix, every row the same. Dong, Cordonnier and Loukas (2021) (Dong, Cordonnier, and Loukas 2023) show that attention-only networks without skip connections lose rank doubly exponentially in depth, and that the skip connection is what prevents it. In smoother language the skip connection makes the stack a boosted ensemble: each layer fits a smoother to what is left over, which is the reading of residual networks in (Huang et al. 2018).
Bandwidth must shrink with sample size. A Nadaraya–Watson estimator is consistent only if the bandwidth goes to zero as \(n\) grows. The transformer’s bandwidth is of order \(\sqrt d / (\|q\|\,\|k\|)\), a function of the weights, fixed at training time. So a trained model is a smoother tuned for the context lengths it saw, and at other lengths it is mis-smoothed. The arithmetic: if one position has score gap \(g\) over \(n - 1\) distractors, its weight is \(e^g/(e^g + n - 1)\), which goes to zero. (Chiang and Cholak 2022) multiply the logits by \(\log n\), which turns \(e^g\) into \(n^g\) and keeps that weight away from zero whenever \(g > 1\). They motivate it from constructions for PARITY and FIRST, but it is the bandwidth correction a statistician would write down.
Averaging is counting. Hard attention picks one position, which makes it a nearest-neighbour regressor: no bias when an exact precedent exists, and no variance reduction. Soft attention averages, and averaging is how a transformer counts. The formal-language results say the same thing in their own terms: transformers with unique hard attention sit in \(\mathrm{AC}^0\) ((Hao, Angluin, and Frank 2022)), while the averaging ones are bounded instead by first-order logic with majority quantifiers ((Merrill and Sabharwal 2025)). The select step is the first-order part; the aggregate step is the majority quantifier. Those results are theorems about this regression class.
7 What kind of learning, then
Three senses of “interpolation” are in play and they come apart.
The first is the one from the benign-overfitting literature: a model that fits its training data exactly. That is a statement about the optimizer and the loss surface and has nothing to do with the architecture; it is as true or false of a transformer as of a kernel machine.
The second is the convex-hull sense. It is a theorem about one softmax head, false for a block, and by Balestriero’s count false for every high-dimensional learner at the level of the training set anyway.
The third is the informal one, that the model blends things it saw in training. That one is mislocated. What gets blended is the context; what training supplied is the rule for blending. Whether that rule is itself “interpolated” from the corpus is a question about the function \(s\), and it is the same question we would ask of any neural network, with the same unsatisfying answers.
What is left, once the word is retired, are the three constraints: local values, pairwise selection, convex aggregation, iterated \(L\) times. The expressivity people have measured how far those reach in the worst case, and the answer is roughly first-order logic with counting. The statistical version I have not seen asked in these terms: which similarity structures are learnable from next-token loss, and at what sample cost? A bilinear score has \(d^2\) parameters per head, which is cheap, and composition over depth is what makes the class of selection rules rich. In nonparametric statistics, adapting the bandwidth to the unknown function costs a logarithmic factor over knowing it; what it costs to learn the adaptation, when the adaptation is a relational program, I do not know.
One cheap interpretive consequence. A smoother with the wrong metric pools the wrong precedents and returns something smooth, plausible and wrong. That is a description of confabulation that needs nothing beyond kernel bias, and it predicts that confabulations should look like averages over nearby precedents rather than like noise.
8 The residual stream, still shelved
In this reading the residual stream is the design matrix. Attention appends columns that are context-smoothed versions of existing columns; the MLP appends columns that are pointwise nonlinear functions of existing columns; the unembedding regresses the logits on the final columns. That is a stack of smoothers with nonlinear feature construction between them, which is also a description of boosting with smoothers as weak learners, and of a deep additive model. I suspect the useful next step is to treat the columns as the object and ask which ones each layer adds, but that is a different post. TODO.
9 Experiments
None of this needs a big model.
- The two gaps on a one-layer attention-only model trained on small text, with the cheating oracle solved by projected gradient on the simplex at each position. Report the hull gap and the selection error as fractions of the deficit against a large reference model.
- Learned bandwidth against oracle bandwidth on in-context denoising: the context holds \(m\) noisy copies of a hidden value among \(n - m\) distractors, the oracle weight is uniform over the copies, and we check whether attention entropy tracks \(\log m\) as \(m\) and \(n\) vary.
- Length transfer: does the trained model degrade with \(n\) the way a fixed-bandwidth smoother does, and does the \(\log n\) rescaling recover it?