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 → What Attention Weights Can Explain- 01 What Attention Weights Can Explain Self-attention has the normalized weighted-average form of a compatibility smoother. That exposes exact routing arithmetic, but it does not make attention weights causal explanations or guarantee a Mercer kernel on tokens. JAX companion Self-Attention as Kernel Regression in JAX/Flax NNX
- 02 What a Finite Kernel Buys an MLP Replace the activation with a finite bank of learned kernel sections. The resulting MLP exposes prototypes, exact layer-local contributions, measurable geometry, and the conditions those claims require, then tests the construction on arithmetic and Fashion-MNIST. JAX companion The Yat-Kernel MLP in JAX/Flax NNX
- 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. 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