The Feed-Forward Block as Key-Value Memory

Rows of W_1 act as keys that match input patterns; columns of W_2 are the values those keys retrieve.

So we've spent a few posts on attention now. How tokens look at each other, build up context, decide what's relevant. But attention is only half the transformer block. The other half is a two-layer neural network that sits right after attention at every level. It processes each token position independently. No cross-talk between positions at all.

When I first encountered the feed-forward block I kind of glossed over it. Attention was the exciting new thing. The FFN was two matrix multiplies with a ReLU in between. Standard neural network stuff.

The first time I looked at a parameter count breakdown of a transformer I realized the "simple" feed-forward block contains about two-thirds of the model's total parameters. Not attention. The FFN. The real twist came when I read Geva et al.'s 2021 paper and saw a way to think about those parameters as a learned key-value memory. Rows of the first weight matrix are "keys" that match input patterns. Columns of the second matrix are "values" that get retrieved when those keys fire.

That reframing changed how I think about what transformers actually do with all those billions of parameters. It also made a bunch of other things click: why knowledge editing works, why MoE architectures replace the FFN specifically, why scaling $d_{ff}$ matters.

Let's build up to it from the basic math.

Prerequisites

You'll want to be comfortable with the self-attention mechanism from Post 4.1, including the residual stream picture introduced there: every block reads from a shared per-token vector and adds its result back. The FFN is one of those blocks. I'll compare its key-value lookup to attention's key-value lookup throughout. Familiarity with matrix multiplication at the level of Post 1.2 is enough for the math here.

The basic FFN

The standard feed-forward block in a transformer is:

where $x$ is the residual stream vector at a single position, $W_1$ has shape $(d_{ff}, d_{model})$, and $W_2$ has shape $(d_{model}, d_{ff})$. In the original "Attention Is All You Need" paper, the activation function was ReLU.

The key dimension here is $d_{ff}$. It's typically 4 times $d_{model}$. So if your model has a hidden dimension of 4096, the intermediate representation balloons up to 16,384 dimensions, then gets projected right back down to 4096. Expand, activate, contract. That's the full computation.

Why expand? The idea is that projecting into a higher-dimensional space and then projecting back gives the network room to carve out more complex decision boundaries. Think of it this way: in the original $d_{model}$-dimensional space, the features are packed tightly together. By blowing them up into $4 \times d_{model}$ dimensions, the network can separate patterns that were tangled in the smaller space, apply a nonlinearity to pick out the useful ones, and compress back down. It's the same intuition behind kernel methods in classical ML, just implemented as a learned linear expansion.

There's no interaction between token positions here. Each position gets the exact same linear transformations applied independently. Attention handles the cross-position communication; the FFN handles per-position computation.

The parameter budget

This is the part that caught me off guard. When I first read the Vaswani paper, I assumed most of the model's capacity was in attention, since that's the novel contribution. Let's count parameters for a single transformer layer with $d_{model} = 4096$ and $d_{ff} = 16384$.

The FFN block:

• $W_1$: $4096 \times 16384 = 67M$ parameters
• $W_2$: $16384 \times 4096 = 67M$ parameters
• Total: 134M parameters (ignoring biases)

The attention block (with 4 projection matrices $W_Q, W_K, W_V, W_O$):

• Each is $4096 \times 4096 = 16.8M$ parameters
• Total: 67M parameters

So the FFN has roughly twice the parameters of attention, per layer. In a 32-layer model, the FFN blocks alone account for over 4 billion parameters. The feed-forward blocks are where the bulk of the model's learned knowledge physically lives, stored as weight matrix entries.

When people talk about how many parameters a model has, they're mostly talking about FFN weights. When people talk about "scaling up" a model, they're mostly talking about making the FFN wider ($d_{ff}$) or deeper (more layers). The attention mechanism is important for what it does. The FFN is where the sheer mass of learned parameters accumulates.

The key-value memory interpretation

Here's the reframing from Geva et al. that I think is genuinely illuminating. Take the FFN computation and break it into steps.

First, $W_1 \cdot x$ computes a dot product between each row of $W_1$ and the input $x$. A dot product is a similarity measure. So each row of $W_1$ is a pattern (a "key"), and the dot product measures how well the current input matches that pattern. The $i$-th element of the intermediate vector is just: how much does $x$ look like the $i$-th row of $W_1$?

Then ReLU zeros out all the negative values. This is a hard gate. Only the keys that match strongly enough (positive dot product after bias) survive.

Finally, $W_2$ multiplies this sparse activation vector. You can think of this column by column. Each column of $W_2$ is a value associated with one intermediate neuron. When neuron $i$ fires with activation $a_i$, the output gets $a_i \cdot W_2[:, i]$ added to it. The stronger the match, the more that value contributes.

Put it all together:

where $k_i$ is the $i$-th row of $W_1$ (the key) and $v_i$ is the $i$-th column of $W_2$ (the value).

That's a key-value memory. The input is the query, the first matrix stores keys, the activation function decides which keys matched, and the second matrix stores the values that get retrieved. With $d_{ff} = 16384$, each layer has 16,384 memory slots.

This is structurally similar to attention, which is also a key-value lookup. But the two kinds of key-value lookup serve totally different purposes:

• Attention keys and values come from the current sequence. They're dynamic, computed on the fly from other tokens in the context window. Different input, different keys and values. Attention is like a conversation: the participants change.
• FFN keys and values are baked into the weight matrices during training. They're fixed. The same 16,384 memory slots exist regardless of what the input is. What changes is which ones get activated.

Another way to say it: attention is working memory. The FFN is long-term storage. (Sukhbaatar et al. formalized this same intuition by treating the FFN as a persistent memory layer with fixed key-value pairs.)

Seeing it in action

Let's make this concrete with a toy example. Below is an interactive FFN with $d_{model} = 4$ and $d_{ff} = 8$. The input vector on the left is a single token's representation. It flows through $W_1$ to produce 8 intermediate activations (the middle column), then through $W_2$ to produce the output vector on the right.

Click on any intermediate neuron to see its key (what input pattern it matches) and its value (what it writes to the output). The key is the corresponding row of $W_1$. The value is the corresponding column of $W_2$. Toggle between ReLU and SwiGLU to see how the activation function changes which neurons fire.

Notice how ReLU is binary about it: a neuron either fires or it doesn't. With SwiGLU, activations are smoother and some neurons that were dead under ReLU contribute a little. The overall output vector shifts noticeably between the two modes.

Also notice the expand-contract shape. Four input dimensions blow up to eight intermediate neurons and then collapse back to four output dimensions. In a real model, that expansion ratio is 4x (or about 2.7x for SwiGLU), giving the network thousands of intermediate "memory slots" per layer.

One thing the visualization makes tangible: the output vector is a weighted sum of the active value columns. Neuron 0 might push toward "Paris," neuron 3 might push toward "city," and neuron 5 might push toward "European." The output blends all of these based on activation strength. It's not winner-take-all. Multiple memories contribute simultaneously, and the resulting vector is a mixture of all the patterns the FFN matched.

SwiGLU: the modern FFN

If you read model cards for LLaMA 3, Mistral, Gemma, or basically anything released after 2022, you'll see "SwiGLU" mentioned under the architecture. It's the modern replacement for the ReLU FFN. Here's why it took over.

The original transformer used ReLU. But since 2020, basically every production model has switched to a gated variant. The one that won is SwiGLU, from Shazeer's "GLU Variants Improve Transformer" paper.

Instead of two weight matrices, there are now three: $W_1$, $W_{gate}$, and $W_2$. The Swish activation ($x \cdot \sigma(x)$, where $\sigma$ is the sigmoid) replaces ReLU, and the $W_{gate}$ projection provides a multiplicative gate. The element-wise product $\odot$ means both the "activation path" and the "gate path" have to agree for a neuron to fire strongly.

Why does this help? Swish is smooth everywhere (unlike ReLU, which has a hard kink at zero), so gradients flow more cleanly during training. The bigger win is the multiplicative gate. With plain ReLU, a neuron's fate is decided by a single dot product: positive means on, negative means off. With SwiGLU, the network computes two separate projections of the input and multiplies them. A neuron only fires strongly when both paths agree, which lets the network express more nuanced activation patterns than a hard threshold allows.

In the key-value memory framing, SwiGLU makes the "matching" step more sophisticated. Instead of a single key per neuron, you effectively have two keys that both need to match for the value to be retrieved. Think of it as a security system with two-factor authentication: both the "do I match?" check and the "should I fire?" check need to agree.

To keep the total parameter count comparable, models using SwiGLU reduce $d_{ff}$. Three matrices instead of two means you need to shrink each one. The standard ratio drops from $4 \times d_{model}$ to about $\frac{8}{3} \times d_{model}$ (often rounded to a multiple of 256 for hardware alignment). For LLaMA 2 with $d_{model} = 4096$, the actual $d_{ff}$ is 11008. Not a round number, but hardware-aligned.

LLaMA, Mistral, Gemma, and most models released after 2022 use SwiGLU. It's one of those changes that happened quietly but universally. If you're reading a model card and see "SiLU activation" or "SwiGLU FFN," now you know exactly what they're talking about.

Attention vs FFN: the division of labor

There's a clean conceptual split in the transformer block.

Attention moves information between positions. It's how the model at position 10 can "see" what's at position 3. It reads from the residual stream at every position and writes a weighted combination back. It's a communication channel.

The FFN transforms information at each position independently. It reads the residual stream at position 10 and writes back to position 10. No other position is involved. It's a computation step, or if you buy the key-value memory framing, a memory recall step.

Elhage et al. describe this nicely as a kind of assembly line. Attention layers move information to where it needs to be. FFN layers process whatever has arrived. The residual stream is the shared workspace that both read from and write to.

I like to think of it with a library analogy. Attention is the librarian who takes your question and walks through the shelves, pulling relevant books and stacking them on your desk. The FFN is the reference book itself. Once the right information has been gathered to your desk, you look things up in the reference material. The librarian knows where things are; the reference material contains the actual answers. Both are necessary, but they do fundamentally different jobs.

So when the model processes "The capital of France is" and needs to predict "Paris," the attention layers gather the relevant context ("capital," "France") into the residual stream at the final position. Then the FFN at that position does a lookup in its learned key-value memory: the key for "capital of France" matches strongly, and the corresponding value pushes the representation toward the token "Paris."

At least, that's the idealized picture. Real models are messier. Attention heads sometimes do computation, FFN layers sometimes do routing. But as a first-order mental model, "attention moves, FFN computes" is solid.

And this alternation happens at every layer. A 32-layer transformer does 32 rounds of "communicate then compute." Early layers tend to handle low-level patterns (syntax, token identity). Later layers handle higher-level abstractions (semantics, factual recall, planning). The FFN at layer 2 is doing different work than the FFN at layer 28, even though they have identical architecture. What differs is what the attention layers have already routed into the residual stream by the time each FFN reads it.

One important detail: the FFN output is added to the residual stream, not substituted into it. The full computation at each layer is:

This residual connection means the FFN doesn't need to preserve information that's already in the stream. It just writes its delta. Retrieved a fact about Paris? Write the "Paris" direction into the stream. Everything else the stream already carries passes through untouched. This additive structure is what makes the "memory write" interpretation concrete: the FFN literally adds retrieved value vectors to the running representation.

If you remove the residual connection (don't do this), the FFN would have to simultaneously preserve all existing information and inject new information through the same bottleneck. The skip connection frees it to focus purely on what it's adding.

Knowledge neurons: finding facts in the weights

OK, so the FFN can be interpreted as a key-value memory. But is it actually used that way? Or is this just a nice mathematical analogy that doesn't correspond to anything real inside a trained model?

Dai et al. showed that it's real, and more literally than I expected.

They introduced a method to identify "knowledge neurons" that correspond to factual associations. For a prompt like "The Eiffel Tower is located in ___," you can trace which intermediate FFN neurons activate most strongly for the correct completion "Paris" and suppress them. When you zero out those specific neurons, the model can no longer produce the right answer. When you amplify them, the model becomes more confident.

The striking part is how few neurons are involved. A single factual association often localizes to a handful of neurons across a few layers. This is not some diffuse representation spread across every parameter. It's relatively sparse.

This connects directly to the key-value memory interpretation. The neurons that fire for "Eiffel Tower is in ___" have W1 keys that match the pattern "famous landmark + location query," and their W2 values push the output distribution toward "Paris." The fact is stored in specific rows and columns of specific weight matrices in specific layers.

This has practical implications. If facts are localized in specific FFN neurons, you can edit them. Model editing techniques like ROME and MEMIT directly modify $W_2$ column entries to update or insert factual knowledge without retraining. You want the model to think the Eiffel Tower is in London? Change a few hundred weight values. (Whether you should do this is another conversation.)

It also raises interesting questions about how knowledge scales. A model with 32 layers and $d_{ff} = 16384$ has about 500,000 memory slots total. That's a lot, but it's not infinite. And a large language model knows millions of facts. Are some neurons polysemantic, storing multiple facts in superposition? Almost certainly yes. A single neuron might activate for "capital cities" generally, not just for one specific capital. This is an active area of interpretability research, and the FFN-as-memory lens is what makes the question even legible.

There's also a layer-depth story here. Dai et al. found that knowledge neurons for factual associations tend to concentrate in the upper-middle layers of the network, not the first or last layers. This makes sense if you think about what each layer has to work with. Early layers haven't aggregated enough context yet (the attention layers haven't had enough rounds to route the relevant information). Late layers are more focused on formatting the output distribution. The middle layers are where context has been assembled and factual recall is most useful.

I find this satisfying as a mental model. The transformer isn't one undifferentiated blob of matrix multiplies. Different layers develop different specializations, and those specializations are discoverable. The "FFN as memory" lens is what makes this kind of analysis possible in the first place.

Worked example

Let's trace a concrete input through a small FFN with $d_{model} = 4$ and $d_{ff} = 8$. I'm using the same numbers as the visualization above, so you can follow along by clicking the neurons.

Input vector: $x = [0.80, 0.30, -0.50, 0.60]$

Step 1: Key matching ($W_1 \cdot x + b_1$)

Each of the 8 rows of $W_1$ computes a dot product with $x$. Take neuron 0 as an example. Its key (row 0 of $W_1$) is $[0.60, 0.20, -0.30, 0.10]$. The dot product is:

After adding the bias ($b_1[0] = 0.10$), neuron 0 has a pre-activation of $0.85$. That's a strong match. This key pattern wanted positive values in dimensions 0 and 2 (via the negative-negative product), and the input delivered.

Now do the same for all 8 neurons. Some will get large positive pre-activations (strong match), some near zero, some negative (no match).

Step 2: Activation (ReLU)

ReLU zeros out all negative values. This is where the "gating" happens. Only the neurons whose keys actually matched the input get to contribute. In the visualization, those are the warm-colored circles. The gray ones had negative pre-activations and got clamped to zero. They're dead for this input.

This sparsity is important. With ReLU, typically only 10-30% of intermediate neurons fire for any given input. The FFN isn't retrieving all 16,384 memories at once. It's doing a sparse lookup, and which subset gets retrieved depends entirely on the input.

Step 3: Value retrieval ($W_2 \cdot \text{activated}$)

Each surviving neuron contributes its corresponding column from $W_2$ to the output, scaled by how strongly it activated. If neuron 0 activated at $0.85$ and its value vector (column 0 of $W_2$) is $[0.30, -0.10, 0.40, 0.20]$, it adds $0.85 \times [0.30, -0.10, 0.40, 0.20] = [0.26, -0.09, 0.34, 0.17]$ to the output. The final output is the sum of all these weighted value vectors plus the output bias.

The crucial thing: the output is determined by which keys matched and how strongly. Change the input even slightly and different neurons fire, retrieving different values. The FFN has learned, through training, which patterns to respond to and what to write when they're detected.

This is where the "memory" metaphor really pays off. You can think of each intermediate neuron as a conditional rule: "If the input looks like this key pattern, add this value vector to the output." Training doesn't explicitly program these rules. Backpropagation carves them into the weight matrices over billions of gradient updates. But the structure of the computation is a pattern-matching lookup. Geva et al. showed that you can even decode the value vectors through the unembedding matrix to see which output tokens each neuron "wants to promote." Many neurons have interpretable specializations: some activate for number words, others for geographic terms, others for specific syntactic patterns.

Toggle between ReLU and SwiGLU in the visualization above to see how the gating mechanism changes which "memories" get retrieved. Under SwiGLU, the gating path adds a second opinion: even if the key matches (high $W_1 \cdot x$), the gate ($W_{gate} \cdot x$) can suppress the neuron, and vice versa.

Common misconceptions

"The FFN is just a nonlinearity between attention layers." This understates what it's doing. The FFN contains roughly two-thirds of the model's parameters and does the heavy lifting of storing and retrieving learned knowledge. The activation function is the gate that decides which of those memories fire for the current input.

"Attention is where the knowledge lives." Attention computes relevance patterns dynamically. It doesn't store facts in any persistent way. The attention weights change completely with every new input sequence. Factual knowledge, world knowledge, linguistic patterns are predominantly encoded in FFN weights, which are fixed after training. Attention figures out which knowledge to retrieve by routing information to the right positions; the FFN actually retrieves it from its static key-value store.

"Bigger $d_{ff}$ always helps." More intermediate neurons means more memory slots, yes. But there are diminishing returns. Doubling $d_{ff}$ doubles the FFN parameters but doesn't double performance. At some point you're adding memory capacity the model can't effectively use given the signal it receives from the attention layers. This is part of why Mixture of Experts (Post 4.9) became attractive: instead of making one huge dense FFN, use many smaller ones and route to the relevant subset. Same total parameter count, but only a fraction of it activates for any given token.

"The FFN processes the whole sequence at once." It processes each token position independently. The same weights are applied to position 0 and position 1023, but there is zero information flow between positions inside the FFN. All cross-position communication happens in the attention layer. The FFN is purely a per-position operation.

What's next

We've now covered both of the major computational blocks inside a transformer layer: attention and the FFN. But I skipped over an important detail in the full layer equation above, the Norm() calls. Where you place normalization, and which flavor you use, turns out to matter a lot for training stability. The next post, Layer Normalization: Pre-Norm, Post-Norm, and RMSNorm, digs into that.

Additional reading (and watching)

• Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017. The original transformer paper; introduces the FFN as a position-wise two-layer network with ReLU.
• Geva, M., et al. (2021). Transformer Feed-Forward Layers Are Key-Value Memories. EMNLP 2021. The paper that crystallized the key-value memory interpretation of FFN weights.
• Shazeer, N. (2020). GLU Variants Improve Transformer. The compact note that introduced SwiGLU and the gated FFN family that took over modern transformers.
• Dai, D., et al. (2022). Knowledge Neurons in Pretrained Transformers. ACL 2022. Localizes specific factual associations to small sets of FFN neurons.
• Elhage, N., et al. (2021). A Mathematical Framework for Transformer Circuits. Anthropic. Frames attention and FFN as cooperating reads/writes against a shared residual stream.
• Meng, K., et al. (2022). Locating and Editing Factual Associations in GPT. NeurIPS 2022. ROME: edits a single $W_2$ row to update one fact at a time.
• Meng, K., et al. (2023). Mass-Editing Memory in a Transformer. ICLR 2023. MEMIT: scales the ROME idea to thousands of edits in a single pass.
• Sukhbaatar, S., et al. (2019). Augmenting Self-attention with Persistent Memory. NeurIPS 2019. Independently formalizes the FFN as a layer of persistent key-value pairs.