Attention as a kernel estimator

Step 1 of 6

A weighted average of what you have seen

You have pairs (x_i, y_i) and a new point x. The oldest non-parametric answer to "what is y at x?" is to average the y_i, weighting each by how much its x_i resembles x:

f̂(x) = Σ_i K(x, x_i) y_iΣ_j K(x, x_j)

K is a kernel: any non-negative function that is large when its two arguments are close. The usual choice is Gaussian, K(x, x_i) = exp(−‖x − x_i‖²/2). The denominator makes the weights sum to one, so the estimate is a convex combination of the y_i: it never leaves the range of what it has seen.

Nadaraya, E. A. (1964). On estimating regression. Theory of Probability and Its Applications 9(1), 141-142. Watson, G. S. (1964). Smooth regression analysis. Sankhyā A 26(4), 359-372.

Step 2 of 6

Change the kernel

Expand the Gaussian: ‖x − x_i‖² = ‖x‖² − 2 x·x_i + ‖x_i‖². The ‖x‖² term is the same for every i and cancels between numerator and denominator. What is left depends on the data only through an inner product. So take the inner product as the kernel itself, and let the two sides be different linear views of the inputs:

K(x, x_i) = exp( q · k_i / √d ),   q = W_Q x,   k_i = W_K x_i

The √d keeps the exponent from growing with the dimension. Nothing else about the estimator has changed: it is still a kernel-weighted average, with a kernel the model can learn instead of one you fixed in advance.

This is the reading in Tsai, Bai, Yamada, Morency, Salakhutdinov (2019). Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel. EMNLP. The textbook version of the same derivation is Dive into Deep Learning, chapter 11, Attention Pooling by Similarity.

Step 3 of 6

The normalised weights are a softmax

Put the new kernel back into the Nadaraya-Watson ratio. The weight on the i-th point is

α_i = exp(q · k_i / √d)Σ_j exp(q · k_j / √d) = softmax(qk^T/√d)_i

and the estimate is f̂(x) = Σ_i α_i v_i with v_i = W_V x_i in place of y_i. Written for all positions at once this is softmax(QK^T/√d) V, the attention of Vaswani et al. (2017). In a language model the x_i are the earlier tokens' residual-stream vectors and x is the current one; the causal mask just restricts the sum to i ≤ current position.

Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser, Polosukhin (2017). Attention Is All You Need. NeurIPS.

Step 4 of 6

The kernel is QK. The values are OV.

Look at what each weight matrix touches. The weights α_i depend on W_Q and W_K only through the product W_Q^T W_K: one bilinear form that scores every (current token, earlier token) pair. That is the QK circuit. It answers where to look.

What actually gets added to the residual stream is W_O Σ_i α_i W_V x_i = Σ_i α_i (W_OW_V) x_i. The matrix W_OW_V is the OV circuit. It answers what to move, and it never sees the weights. In the kernel picture: QK is the kernel, OV is the y_i.

head(x) = Σ_i softmax(x^T W_Q^TW_K x_i / √d)  ·  W_OW_V x_i

Two independent linear maps per head, and neither one alone tells you what the head does. The framework paper draws this split and names it; everything on the poster's left-hand side is the QK half, everything in the attribution graph is the OV half plus the MLPs. The same split written out with explicit tensor indices, shapes and einsum, down to virtual heads and the induction head, is on the companion page.

Elhage, Nanda, Olsson et al. (2021). A Mathematical Framework for Transformer Circuits. Anthropic.

Step 5 of 6

An induction head is a kernel with one sharp peak

Suppose the sequence contains [A][B] somewhere and the current token is [A] again. An induction head's QK circuit puts almost all of α on one position: the token right after the previous [A]. (It can do this because an earlier "previous-token" head has already written "I follow an A" into [B]'s residual stream.) Its OV circuit then copies that token's identity forward, so the prediction is [B].

In the kernel language: a kernel that is nearly zero everywhere except at "the position after the last time I saw this token", and values that are just the token identity. Olsson et al. argue this pair of heads is the main mechanism behind in-context learning.

Olsson, Elhage, Nanda et al. (2022). In-context Learning and Induction Heads. Anthropic.

Step 6 of 6

Back to the poster: where, what, does it matter

The kernel view is why the poster splits the question in three. The attention pattern is the kernel: it says where the return position looks, and nothing more. What is moved along those weights, through OV and through the MLPs, is what the attribution graph shows, with every effect named against a logit. And whether a pathway matters is a separate question again, answered by scaling it down and watching the prediction change.

This is also why an attention heatmap on its own is not an explanation: high weight on a token says the kernel is peaked there, not that what was moved decided the answer. Jain and Wallace (2019) make the point at length. On the poster the heatmap never stands alone; it sits beside the graph and the dose-response curve.

Jain, S., Wallace, B. C. (2019). Attention is not Explanation. NAACL.