Transformer MLPs as learned key-value memories

Yesterday we built a useful picture of the residual stream as the Transformer's shared scratchpad. Attention heads read from that scratchpad, move information between token positions, and write their results back into it.

Today we'll look at the other major component that repeatedly reads and writes the same stream: the MLP, sometimes called the feed-forward network.

The surprising idea is:

An MLP can be viewed as a large bank of learned pattern detectors, each paired with a learned update to the residual stream.

That makes an MLP look less like a generic blob of nonlinear computation and more like a kind of learned associative memory.

A useful approximation is:

\[ \boxed{\text{MLP} \approx \text{detect a pattern} \rightarrow \text{write a useful feature}} \]

This is not the only way to understand MLPs, and real representations are distributed. But it is an extremely productive mental model.

Start with the shape¶

Ignoring gating for a moment, a Transformer MLP looks roughly like:

\[ MLP(x)=W_{down}\,\sigma(W_{up}x). \]

Suppose:

\[ d_{model}=4096 \]

and the intermediate MLP dimension is:

\[ d_{ff}=14336. \]

Then:

\[ x\in\mathbb R^{4096} \]

is projected upward:

\[ W_{up}x\in\mathbb R^{14336}. \]

A nonlinear activation is applied, and the result is projected back down:

\[ W_{down}:\mathbb R^{14336}\rightarrow\mathbb R^{4096}. \]

Finally, as we saw yesterday, that result is added to the residual stream:

\[ x'=x+MLP(x). \]

So the MLP temporarily expands one 4096-dimensional token representation into more than 14,000 intermediate coordinates.

Why?

One useful answer is that the expansion gives the model many different feature detectors it can activate.

Rewrite the MLP one neuron at a time¶

Let the rows of \(W_{up}\) be:

\[ k_1^T,k_2^T,\ldots,k_m^T \]

where:

\[ m=d_{ff}. \]

Then intermediate neuron \(i\) computes something like:

\[ a_i=\sigma(k_i^Tx). \]

Now let the columns of \(W_{down}\) be:

\[ v_1,v_2,\ldots,v_m. \]

Then the MLP output can be written as:

\[ MLP(x)=\sum_{i=1}^{m} a_i v_i. \]

Substituting:

\[ MLP(x) = \sum_{i=1}^{m} \sigma(k_i^Tx)v_i. \]

This equation is the heart of today's lesson.

Each intermediate neuron has two important vectors:

\[ \boxed{k_i} \]

which determines when it activates, and:

\[ \boxed{v_i} \]

which determines what it writes when it activates.

That begins to resemble a key-value memory.

Why call them keys and values?¶

Imagine one detector has a key direction corresponding roughly to:

French city + capital relation + factual completion

If the current residual stream aligns strongly with that direction:

\[ k_i^Tx \gg 0, \]

the neuron activates.

Its associated value vector might point roughly toward:

Paris-like completion

The neuron then contributes:

\[ a_iv_i \]

to the residual stream.

Conceptually:

residual stream
      ↓
does this match my key?
      ↓ yes
activate neuron
      ↓
write my value vector
      ↓
updated residual stream

The model can have thousands of such detector/write pairs in every layer.

A tiny numerical example¶

Suppose our residual stream has only three dimensions:

\[ x= \begin{bmatrix} 2\\ 1\\ 0 \end{bmatrix}. \]

Imagine one MLP neuron has key:

\[ k= \begin{bmatrix} 1\\ 1\\ 0 \end{bmatrix}. \]

Its pre-activation is:

\[ k^Tx = 1(2)+1(1)+0(0) = 3. \]

Suppose for simplicity our activation is ReLU:

\[ \sigma(z)=\max(0,z). \]

Then:

\[ a=\sigma(3)=3. \]

Now suppose its value vector is:

\[ v= \begin{bmatrix} 0\\ 0\\ 2 \end{bmatrix}. \]

The neuron's contribution is:

\[ av = 3 \begin{bmatrix} 0\\ 0\\ 2 \end{bmatrix} = \begin{bmatrix} 0\\ 0\\ 6 \end{bmatrix}. \]

So the updated residual stream becomes:

\[ x' = x+av = \begin{bmatrix} 2\\ 1\\ 6 \end{bmatrix}. \]

The first two dimensions caused the detector to fire.

The neuron then wrote information primarily into the third dimension.

This illustrates the basic pattern:

\[ \boxed{\text{read one direction, write another direction}} \]

Many neurons fire simultaneously¶

Of course, a real MLP does not choose one memory slot.

It computes:

\[ MLP(x) = \sum_i \sigma(k_i^Tx)v_i. \]

So perhaps:

neuron  182 → detects "France"
neuron 1407 → detects "capital relation"
neuron 3901 → detects "factual completion"
neuron 8112 → detects "European geography"

all activate to different degrees.

Their value vectors are summed:

\[ \Delta x = a_{182}v_{182} + a_{1407}v_{1407} + a_{3901}v_{3901} + a_{8112}v_{8112} +\cdots \]

and the combined update is written into the residual stream.

This is why calling an MLP a literal database would be misleading.

It is closer to a soft, distributed associative memory.

Multiple memories can partially match, and their outputs blend.

Connect this to the residual-stream lesson¶

Yesterday we imagined this computation:

prompt:
"The capital of France is"

attention:
finds "France"
      ↓
writes country-related information
into residual stream
      ↓
MLP:
recognizes the resulting feature combination
      ↓
writes information favoring "Paris"

Now we can make that more concrete.

Suppose attention produces:

\[ x'=x+\Delta x_{attn}. \]

The new vector \(x'\) now aligns with some MLP detector:

\[ k_i^Tx' > k_i^Tx. \]

So attention has indirectly caused MLP neuron \(i\) to activate.

That neuron writes:

\[ a_iv_i. \]

If \(v_i\) aligns with the unembedding direction for Paris, then:

\[ w_{Paris}^Tv_i>0. \]

Therefore the neuron increases the Paris logit.

We now have a small circuit:

France token
    ↓
attention head
    ↓
writes feature into residual stream
    ↓
MLP key detects feature
    ↓
MLP value writes new feature
    ↓
Paris logit increases

This is the kind of mechanism interpretability researchers try to uncover.

Why the nonlinearity matters¶

Without the activation function:

\[ MLP(x)=W_{down}W_{up}x. \]

But the product:

\[ W_{down}W_{up} \]

would simply be another linear transformation.

There would be no real need for the expanded hidden dimension.

The nonlinearity allows different detectors to turn on under different conditions.

For ReLU:

\[ a_i=\max(0,k_i^Tx). \]

If:

\[ k_i^Tx<0, \]

the neuron contributes nothing.

If:

\[ k_i^Tx>0, \]

it contributes proportionally.

Modern LLMs often use smoother activations and gated MLPs, but the important idea remains:

The nonlinear intermediate layer allows the model to conditionally apply different learned transformations depending on what features are present.

Modern models often use gated MLPs¶

Many current architectures use variants such as SwiGLU rather than a simple:

\[ \sigma(W_{up}x). \]

A simplified gated form is:

\[ MLP(x) = W_{down} \left[ \operatorname{SiLU}(W_gx) \odot (W_ux) \right]. \]

There are now two upward projections.

One computes something like a gate:

\[ g=\operatorname{SiLU}(W_gx), \]

while another computes content:

\[ u=W_ux. \]

They are multiplied elementwise:

\[ h=g\odot u. \]

Then:

\[ MLP(x)=W_{down}h. \]

Conceptually, each intermediate channel can ask both:

should I activate?

and:

how strongly / with what signed content?

before the down projection writes the result back into the residual stream.

The simple key-value picture becomes less exact mathematically, but the conditional read → write intuition remains valuable.

Why MLPs contain so many parameters¶

This also explains a fact that surprises people when they first inspect a Transformer.

A large fraction of the parameters are often in the MLPs, not attention.

For a simple MLP:

\[ W_{up}\in\mathbb R^{d_{ff}\times d_{model}} \]

and:

\[ W_{down}\in\mathbb R^{d_{model}\times d_{ff}}. \]

Ignoring biases, that is:

\[ 2d_{model}d_{ff} \]

parameters.

With:

\[ d_{model}=4096 \]

and:

\[ d_{ff}=14336, \]

that is approximately:

\[ 2(4096)(14336) \approx117\text{ million} \]

parameters per MLP block.

A gated MLP adds another large projection and can contain even more.

Those parameters create a huge learned repertoire of ways to recognize patterns in the residual stream and transform them.

Are facts literally stored in MLP neurons?¶

This is where we need an important caveat.

The key-value interpretation is useful, but it can become too literal.

It would be tempting to imagine:

neuron 7341 = "Paris is the capital of France"
neuron 9122 = "Einstein developed relativity"

Real networks are usually messier.

Knowledge can be distributed across:

multiple neurons,
multiple layers,
attention heads,
residual-stream directions,
combinations of features.

One neuron may also participate in many apparently unrelated contexts because of superposition.

So it is safer to say:

\[ \boxed{\text{MLPs implement many learned conditional feature transformations}} \]

rather than:

\[ \boxed{\text{each neuron stores one fact}} \]

But some neurons really are surprisingly interpretable¶

Despite that caveat, researchers often find intermediate units that activate strongly for recognizable patterns.

For example, a neuron might respond disproportionately to:

dates
Python code
legal language
a particular language
URLs
certain syntactic constructions

That does not necessarily mean the neuron is that concept.

But it means the learned detector:

\[ k_i^Tx \]

is sensitive to some direction strongly associated with that pattern.

This is one reason neuron activation inspection was an early mechanistic-interpretability technique.

Why sparse autoencoders became useful¶

Yesterday we discussed sparse autoencoders.

Now we can see why they help.

A raw MLP neuron may be polysemantic:

fires for concept A
and concept B
and strange subset of concept C

because the model has packed more features into its representation than there are clean coordinate axes.

An SAE attempts to recover a larger set of sparse feature directions:

\[ x \approx \sum_j a_jf_j. \]

Instead of asking:

What does neuron 4172 mean?

we can ask:

What sparse features are present in this activation, and which MLP components cause those features to appear or disappear?

That moves interpretability from neurons as concepts toward directions/features as concepts.

MLPs versus attention¶

We can now sharpen a distinction that has appeared throughout these lessons.

A useful approximation is:

\[ \boxed{\text{attention retrieves information from other token positions}} \]

while:

\[ \boxed{\text{MLPs transform information already present at one token position}} \]

Attention answers questions like:

Which earlier token matters?
What information should I copy here?
Which position contains the relevant subject?

MLPs answer questions more like:

Given the features currently present here,
what additional features should become active?

So:

attention = communication across positions
MLP       = computation within a position

Again, this is an approximation—attention can compute and MLPs can encode complex knowledge—but it is an excellent architectural mental model.

Why this matters for inference performance¶

There is also a systems consequence.

During autoregressive decode, attention must read the growing KV cache:

\[ K,V\propto T. \]

So attention's memory traffic grows with context length.

The MLP does not need a history cache.

For the current token it performs large matrix multiplications whose shapes are essentially independent of context length.

So decode has two very different workloads:

attention:
scan historical KV state
cost grows with context

MLP:
large dense projections
cost mostly fixed per generated token

This is one reason inference optimization cannot treat a Transformer layer as one homogeneous operation.

Attention and MLPs stress hardware differently.

Mixture-of-Experts takes this idea further¶

Recall our earlier lesson on MoE models.

A dense MLP has perhaps thousands of intermediate channels available for every token.

An MoE layer instead contains several entire MLP experts:

             router
          /    |    \
      expert 2 | expert 7
               |
            expert 11

and activates only a small subset.

You can think of this as adding another level of conditional computation.

Dense MLP:

\[ \text{choose which internal features activate} \]

MoE:

\[ \text{first choose which MLPs run} \]

then inside those experts:

\[ \text{choose which internal features activate}. \]

That lets the model increase its repertoire of learned transformations without paying the compute cost of running every expert for every token.

A deeper mental model: a wall of specialists¶

Return to yesterday's shared whiteboard.

The residual stream is the whiteboard.

Attention heads are researchers who can look around the room and copy useful information from other whiteboards.

The MLP is different.

Imagine a wall containing thousands of specialists.

Each specialist has a sign describing what catches their attention:

"If I see evidence of X, wake me up."

When the current whiteboard matches that specialist's trigger, they wake up and write their characteristic contribution onto the board.

Formally:

\[ \text{match strength} = k_i^Tx \]

then:

\[ \text{activation} = \sigma(k_i^Tx) \]

then:

\[ \text{written update} = \sigma(k_i^Tx)v_i. \]

All awakened specialists write simultaneously:

\[ MLP(x) = \sum_i \sigma(k_i^Tx)v_i. \]

And because the result is residual:

\[ x' = x+ \sum_i \sigma(k_i^Tx)v_i. \]

So a Transformer repeatedly alternates between two powerful operations:

ATTENTION
Who elsewhere in the sequence has useful information?
        ↓
write it to the shared state

MLP
Given what's now in the shared state,
which learned transformations should fire?
        ↓
write their results to the shared state

Then repeat:

\[ \boxed{ \text{communicate} \rightarrow \text{transform} \rightarrow \text{communicate} \rightarrow \text{transform} } \]

layer after layer.

That is a remarkably useful way to demystify the Transformer.

Attention moves information to where it is needed.

MLPs recognize combinations of that information and turn them into new features.

The residual stream lets every later component build on what earlier components discovered.

Together, those three ideas—attention, MLP feature transformations, and the residual stream—give us a fairly concrete picture of how a Transformer can assemble a computation one learned update at a time.