The MLP Block Can Be a Kernel Memory

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 exposed arithmetic. Attention forms a score for every query-key pair, normalizes each row, and averages named value vectors. Those weights show which values were routed through that layer, but they are not automatically causal explanations of the model’s decision. Separate query and key maps also mean ordinary attention need not define a symmetric PSD kernel on tokens. The useful affordance is the explicit score-and-average structure.

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 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.

A symmetric PSD kernel supplies a feature-space pseudometric, while an arbitrary compatibility score does not. The construction below gives the MLP named centers and a declared PSD comparison at each unit, then asks what that explicit structure buys.

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.

The panels on this page replay a real trained model: a small transformer with a Yat feed-forward block, trained on Shakespeare, whose exported weights run live in your browser.

The block is a kernel expansion#

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. This has the same algebraic shape as a finite RKHS expansion, $f(\cdot)=\sum_u \alpha_u k(W_u,\cdot)$. That resemblance is useful, but it is not itself a representer theorem. A classical representer theorem starts with a regularized optimization problem in a fixed RKHS and proves that an optimizer lies in the span of kernel sections at the training inputs, $f(\cdot)=\sum_i \alpha_i k(x_i,\cdot)$. Here the finite expansion is chosen as the architecture.

The distinction matters. The Yat block learns its centers $W_u$, so training changes the feature dictionary and generally changes the induced function class. The classical theorem does not justify that learned-center optimization or transfer fixed-kernel generalization guarantees to it. What the architecture guarantees is narrower and concrete: at every step, the block is a finite sum over prototypes that learning moves around. Representer-point selection constructs an attribution around a trained network after the fact; here the slot decomposition is simply the arithmetic the block already runs.

Both halves, one principle#

Side by side, the two halves now share a weighted-sum shape. Attention produces $\sum_j a_j v_j$ with nonnegative normalized routing weights. The Yat MLP produces $\sum_u k(W_u,x)v_u$ with nonnegative kernel scores but signed effects after decoding through each value vector. One indexes context tokens; the other indexes learned slots. The common benefit is a list of references and exact layer-local terms, not a guarantee that the terms are semantic concepts or complete explanations.

Shape is not the same as regime, and the difference sets how far the attribution can be pushed. Attention’s weights are nonnegative and sum to one, which puts it in the convex corner of the readout taxonomy: its output cannot leave the hull of the values, and “this token supplied 30 percent” is well posed. The block here keeps the kernel scores but leaves the coefficients $\alpha_u$ signed and unnormalized, which is the linear corner of that same taxonomy, the one with the widest reachable set and the weakest attribution guarantee. What survives is exactness: the decomposition is the model’s actual arithmetic rather than a fitted surrogate, so every edit below lands where the algebra says it will. What does not survive is contribution mass. A slot’s term can be large and negative, two slots can cancel, and a row reads as a signed push toward the answer rather than a share of it. Normalizing the coefficients would recover the convex reading and turn the block into a different model, which this post does not train.

What a white box buys you#

Reading the block was the easy half. Once it 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 the finite-kernel MLP read off image prototypes, now in a transformer.

Read, attribute, edit, abstain: four operations, none of them available when the primitive is a direction with no reference, all four of them just arithmetic once the primitive is a kernel.

The MLP was never doomed to be a black box. Give it an explicit kernel expansion and the block becomes a key-value memory you can read, attribute through, and edit. The raw peak kernel score can also be tested as an abstention signal, but it needs calibration before it should be called confidence. These operations make this feed-forward block more inspectable; they do not make every transformer decision self-explanatory, and attention weights alone are not causal explanations. The useful claim is the exact one: this architecture exposes a named set of references and the signed term each one contributes.

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.