The MLP Block Is a Representer Theorem

If you watched the 3Blue1Brown attention video, you came away able to read half of a transformer. You can point at the attention pattern and say what it does: this token is looking at that one, the verb is reaching back for its subject, the pronoun is fetching its referent. The picture is legible because you can see who is talking to whom.

Then the same series reaches the other half, the MLP block, and the tone changes. This is where the network stores what it knows, and nobody is quite sure how it stores it. The neurons are polysemantic, the features are superposed, the block is a wall of numbers with no handle on it. So you finish the series able to read one half of the machine and staring at the other. Why is attention the readable half and the MLP the black box?

What makes attention readable?#

The asymmetry is not about depth or width; it is about the primitive. Attention is built from a similarity. Each token forms a query, every other token a key, and the output is a weighted average of values where the weights are how similar the query is to each key. That is a kernel, and a vote by similarity is something you can always read: the weights are the explanation, one number per source token, all of them adding to one. We made this exact in Attention Is Explainable Because It Is a Kernel. The legibility was never a happy accident of attention; it is what any kernel gives you, because a kernel says “here is how much each reference counts,” out loud.

So the real question sharpens. The MLP is opaque, but is its opacity fundamental, or is it just that the MLP was built from the wrong primitive?

And what makes the MLP opaque?#

It is the wrong primitive. A standard hidden unit computes $\sigma(\langle w, x\rangle + b)$, and the weight $w$ in there is not a reference you can point at; it is a direction, a hyperplane normal, and the unit fires across an entire half-space of inputs. There is no center, no “this is what I am looking for,” nothing to put your finger on, which is exactly where a weight cannot live: off in a space of directions, joined to the data only by a shadow. Superposition and polysemanticity are downstream of that one fact. The primitive carries no geometry, so the block it builds carries none either, and there is nothing for an explanation to grab.

Attention got a metric for free, because a similarity has one built in. The MLP threw the metric away when it chose an activation. So the fix writes itself: give the MLP a kernel too.

Replace the activation with a kernel#

Concretely, the swap is one move: the Yat kernel in place of the activation turns the hidden unit from a direction into a similarity to a learned center $W_u$:

Now the unit has a reference, the center $W_u$, a point in the same space as the input, and it fires when the input is near and aligned with it. The hidden layer is no longer a wall of numbers; it is a vector of similarities to a handful of named prototypes. The same swap that makes attention’s weights readable has been applied to the other half of the block.

None of what follows is a diagram: the panels on this page replay a real trained model, a small transformer with a Yat feed-forward block trained on Shakespeare by scripts/yat_ffn_whitebox.py, whose exported weights run live in your browser.

The block is a representer theorem#

But a hidden layer you can read is only half the block; the linear projection stacked on top, $y = W_2 \,\mathrm{yat}(x)$, could still scramble whatever the kernel made legible. It does the opposite. Write that sum out, with $\alpha_{u} = (W_2)_{:,u}$ the output row for unit $u$, and the whole block collapses to

a signed, weighted sum of kernel evaluations against the centers. That is exactly the form the representer theorem promises. The theorem says the optimal weight in a reproducing kernel Hilbert space is always a combination of kernel evaluations at a set of centers, $f(\cdot) = \sum_i \alpha_i\, k(x_i, \cdot)$, which is the whole subject of Where Does a Weight Live?. The kernelized MLP block is that form, made the architecture.

It is worth being exact about the difference, because the difference is the point. The representer theorem fixes its centers to the training data and proves that the optimum lands in their span. The Yat block does not wait for the optimum and does not fix the centers to the data: it bakes the representer form in and lets the centers $W_u$ be learned. So every MLP block, at every step of training, is already a representer-theorem vote, a sum over prototypes you can read, where the prototypes themselves are what learning moves around. The decomposition is not an after-the-fact attribution method bolted onto a black box, the way representer-point selection wraps a trained network; it is what the block literally computes.

Both halves, one principle#

So what happened to the asymmetry we started with? Side by side, the two halves of the transformer are now the same equation. Attention produces a token’s output as $\sum_j a_j\, v_j$, a vote over the other tokens’ values with weights $a_j$ that are similarities. The Yat MLP produces its output as $\sum_u \alpha_u\, k(W_u, x)$, a vote over learned prototypes with weights that are similarities. One votes over the tokens present in the context, the other over concepts the model has stored, but both are kernels, and both wear their explanation on the outside: a list of references and how much each one counted.

What a white box buys you#

Legible is not the end of it; legible is the beginning. Once the block is $\sum_u k(W_u, x)\, v_u$, with $v_u = (W_2)_{:,u}$ the vector it writes into the residual stream, the feed-forward layer is a key-value memory, the Geva et al. reading that the readout post traced through trained language models, made exact here rather than statistical. And an exact memory is one you can operate on.

You can read a slot: decode its value $v_u$ through the unembedding and it tells you, in tokens, what it writes (“this one ends lines,” “this one starts capitals”). You can attribute an output: since it is a sum, you split any FFN output into the slots that produced it, with no gradients and no saliency guesswork, because the decomposition is the arithmetic. That is what the memory map and the attribution panel above already showed, on the real model.

The next two go further, from reading the memory to operating on it. You can edit a single slot: turn one $v_u$ up or down and, because the output is linear in $k(W_u, x)$, the generation changes the way its label predicts, model surgery with a screwdriver instead of a sledgehammer.

And you can tell when the memory is out of its depth: the Yat kernel is local, so an input near nothing it has stored lights up no slot and the block writes almost nothing, the abstention signal Your Neuron Is a Picture read off image prototypes, now in a transformer.

None of that is available when the primitive is a direction with no reference. All of it falls out the moment the primitive is a kernel.

That is the whole argument. The MLP was never doomed to be a black box; it was built from a primitive with no metric, and the opacity followed from that and nothing else. Give it a kernel and the block becomes a representer-theorem vote over prototypes, a key-value memory you can read, attribute through, edit, and trust to abstain, as legible as the attention beside it. The panels on this page are the feed-forward half of “the whole transformer explains itself”; the attention half was made exact in Attention Is Explainable Because It Is a Kernel. Between the two posts, both halves now read, by the same principle: a similarity to a reference you can name was available for the whole machine all along.

Want to run it? The JAX/Flax NNX companion builds a real transformer with a Yat feed-forward block, trains it on tinyshakespeare, and then reads, attributes, edits, and probes its memory directly, the four white-box payoffs in code.

The attention and MLP framings are from 3Blue1Brown’s deep learning series; the representer theorem is Schölkopf, Herbrich and Smola (2001); representer-point attribution is Yeh et al. (2018); feed-forward layers as key-value memories is Geva et al. (2021); superposition in MLP blocks is Elhage et al. (2022); the Yat kernel is Bouhsine (2026). It builds on Attention Is a Kernel and What a Finite Kernel Buys an MLP.