MLP
The transformer MLP as a kernel and a prototype readout: reading W_out as output prototypes and the convex/conic/affine/linear regimes.
-
A White-Box FFN: the Representer Theorem in JAX/Flax NNX
A runnable companion: build a transformer whose feed-forward block is a Yat kernel, so the FFN is exactly a representer sum over learned key-value memory slots. Train it on tinyshakespeare, then do four things you cannot do to an opaque ReLU FFN: read each memory slot, attribute an output to the slots that wrote it, edit one slot and watch generation change, and read off when the memory is out of its depth.
-
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.
-
Your Neuron Is a Picture, in JAX/Flax NNX
A runnable companion: build the prototype MLP from the post in Flax NNX, train it on Fashion-MNIST, and watch the neurons. Pull the prototypes out as images, read a prediction as a vote over pictures, see the model abstain on out-of-distribution digits, check that random-init prototypes classify but stay noise, and track the prototypes migrating through a UMAP fit on the dataset as they train.
-
Your Neuron Is a Direction. It Should Be a Picture.
Why should a neuron store a direction when it could store a thing? A direction is not a referent you can point at, which is why MLPs are opaque. Put the Yat kernel where the activation was, train on Fashion-MNIST, and every neuron becomes a prototype that lives in pixel space, literally a picture, so the network reads its own predictions: this looks like that, no saliency method required.
-
The Yat-Kernel MLP in JAX/Flax NNX
A runnable companion to What a Finite Kernel Buys an MLP: build a layer whose unit is the Yat kernel instead of a linear map plus an activation, assert it is positive definite and nonnegative, write down its exact finite feature map, train it end-to-end on two moons with no activation function, and measure the lazy-loading sparsity, the bounded off-distribution response, the RKHS capacity, and the force field that pulls each prototype onto its data.
-
What a Finite Kernel Buys an MLP
Replace the activation function with a finite, explicit, positive-definite kernel, the Yat kernel, and an MLP stops being a stack of linear maps glued by a nonlinearity. It becomes a kernel machine, with locality, attribution, geometry, capacity control, and a feature map you can write down.
-
The Prototype Readout in JAX/Flax NNX
A runnable companion to The Readout is a Convex Combination of Prototypes: read the columns of W_out as output prototypes in Flax NNX, measure the convex/conic/affine/linear regimes numerically, then build a Nadaraya–Watson kernel readout that is convex by construction (nonnegative weights that sum to one, a point that never leaves the prototype hull), with the nonnegativity-vs-positive-definiteness distinction checked in code.
-
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.
-
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.