series · 4 parts
Attention Is a Kernel
Attention, read as kernel regression: what the softmax is really computing, why that makes it explainable, and what happens when you make the kernel cheap.
Start reading → Attention is Explainable Because it is a Kernel- 01 Attention is Explainable Because it is a Kernel Self-attention in transformers is a Nadaraya–Watson kernel smoother. That fact, and not "we visualize the matrix", is why attention heads are readable while MLPs are not. JAX companion Self-Attention as Kernel Regression in JAX/Flax NNX
- 02 What an MLP Knows, When It's a Kernel The transformer MLP is illegible because its primitive does not carry a kernel. Give it one and the four objects that make attention legible follow for free, for the whole network.
- 03 Cheap Attention: Linear-Time Kernel Approximation A 128K-token context creates billions of pairwise questions per attention head. But the N×N matrix is not the essence of attention; it is the receipt for an infinite feature map we never wrote down. Approximate that feature map with random features, reassociate the sum, and softmax attention becomes linear-time kernel attention. The whole argument is built from live in-browser visualizations. JAX companion Cheap Attention in JAX/Flax NNX
- 04 Why Attention Needs Q and K Projections The dot product in attention is not enough by itself. Without learned query and key projections, attention can only compare tokens in the residual stream’s native geometry. With a shared projection it learns a symmetric metric. With separate Q and K projections, the score becomes a learned bilinear form x_iᵀW_QW_Kᵀx_j: directional, role-aware, low-rank, and different per head. That bilinearity is what lets attention ask one kind of question and let tokens advertise another kind of answer. JAX companion Q and K Projections in JAX/Flax NNX