series · 5 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
- 05 The Kernel Between the Roles The QK post in this series ended on a construction it refused to build: keep the query and key roles, but replace the bilinear-then-exponential score with a genuine Mercer kernel between them. This post builds it, with the kernel's full form: a per-head learned bias inside the square, the term the universality theorem requires, and a per-head learned softening. Because the kernel is nonnegative by construction, attention needs no softmax at all and routing becomes a literal Nadaraya-Watson smoother. Trained head to head at matched parameters and per-variant swept learning rates, the kernel transformer lands within 1.1 percent of softmax on character-level Shakespeare, and the differences that survive are the interesting part: no gauge, no max-trick, a mass channel softmax cannot represent, and two of our own assumptions measured dead. JAX companion Softmax-Free Attention in JAX/Flax NNX