series · 6 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
- 06 The Geometry of Attention Is a Choice of Kernel An attention head is only ever evaluated at the tokens of the sequence, but its formula accepts any query vector at all. So between the tokens there is a whole continuous space, and the head silently carves it into territories: for every possible query, some key takes the top weight. This post draws that map for the two attention laws this series has trained head to head. The dot-product law slices space into infinite wedges of sky that all meet at the origin, a gravity that reads only bearings, where a query's length is a temperature and a key inside the others' hull can never win. The Yat kernel carves bounded neighborhoods around bodies that always own their ground, a gravity that reads places, where length is an address. One theorem per map, one live map per theorem; the trained transformers appear only in a coda, to confirm they were never free to draw anything else. JAX companion Measuring Attention's Geometry in JAX/Flax NNX