Attention
Attention mechanisms from scratch: scaled dot-product attention, Q/K projections as a bilinear form, attention as a kernel smoother, and linear-time variants.
-
Measuring Attention's Geometry in JAX/Flax NNX
A runnable companion to the geometry-of-attention post: the telemetry that exports a trained head's actual query and key vectors, the offline replay that recomputes every row's winner under rescaled queries, the linear program that asks whether any key sits inside the others' convex hull, and the census of who owns whom through training. Plus the toy laws animated from their real formulas: territories re-carving, a body crossing the hull, the far sky losing its sign.
-
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.
-
Softmax-Free Attention in JAX/Flax NNX
A runnable companion to the compatibility-kernel post: the attention module where one branch computes softmax and the other computes the Yat kernel with no exponential anywhere, the parameter-matched training harness, the telemetry that measured our bounded-scores belief dead, and the checkpointed attention maps. Every number is from the real Kaggle runs.
-
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.
-
The MLP Block Is a Representer Theorem
After the 3Blue1Brown attention video you can read half a transformer: you can see which token attends to which. The other half, the MLP block, stays a black box. But attention is legible because it is a kernel, a vote by similarity, and if you make the MLP a kernel too, its output becomes the same thing: a representer-theorem vote over learned prototypes. Then the whole transformer explains itself.
-
Q and K Projections in JAX/Flax NNX
A runnable companion to Why Attention Needs Q and K Projections: build scaled dot-product attention with separate query and key projections in Flax NNX, pull the bilinear form B = W_Q W_Kᵀ out of the module, split it into a symmetric metric and an antisymmetric directed part, wire a toy induction head, add RoPE, and measure the low-rank budget and the gauge freedom, all in plain JAX.
-
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.
-
The Readout is a Convex Combination of Prototypes
The second linear map in a transformer MLP is not just a projection. If the hidden activations are nonnegative and normalized, W_out reads the active neurons as a convex combination of output prototypes. Two independent constraints, nonnegativity and summing to one, sort the readout into four regimes: convex, conic, affine, and linear. This reframes the MLP readout as the same object that makes attention legible (a weighted sum over named basis elements), connects it to feed-forward key-value memories and modern Hopfield retrieval, and shows when a kernel makes it convex by construction.
-
Cheap Attention in JAX/Flax NNX
A runnable companion to Cheap Attention: implement positive-feature linear attention in JAX and Flax NNX, watch the all-pairs ledger turn into a shared feature state, and see exactly where the N×N matrix disappears.
-
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.
-
Self-Attention as Kernel Regression in JAX/Flax NNX
A runnable companion to Attention is Explainable Because it is a Kernel: build scaled dot-product attention from scratch in Flax NNX, prove in code that it is exactly a Nadaraya–Watson kernel smoother, watch the separate q/k projections break positive-definiteness numerically, swap the exp-dot-product kernel for Gaussian, Yat, and linear kernels to see which keep the weights a convex partition of unity, read the temperature as a kernel bandwidth, and train a single head end-to-end to route to a marked token.
-
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.