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.

Attention lets token positions exchange information. The feed-forward network that follows it processes each position independently, using the same learned weights at every position.

In the original transformer, the FFN is two matrix multiplications with a ReLU between them.

At the original 4× expansion ratio, the FFN contains two-thirds of an encoder layer’s main projection weights. Geva et al. (2021) interpret those weights as keys and values: match an input pattern, then add its associated output direction.

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 expands to 16,384 dimensions, then gets projected right back down to 4096. The nonlinearity acts between the expansion and contraction.

A wider intermediate layer supplies more learned projections of the input and more activation coefficients for the output to combine. The nonlinearity lets these coefficients change differently with the input; without it, the two matrices would compose into one linear map.

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

Count a 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

The FFN has twice the projection parameters of standard multi-head attention in this example. Across 32 layers that is over 4 billion FFN parameters. Parameter share is not a measurement of where a corresponding share of knowledge is stored.

The key-value memory interpretation

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

Attention also has persistent learned projection weights. The comparison here is specifically between its input-dependent K/V activations and the FFN’s fixed key/value matrices.

Seeing it in action

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.

ReLU has an exact zero region, but its positive outputs are continuous, not binary. SwiGLU multiplies two input-dependent projections and can produce signed values. The toy view keeps the same width to compare the operations; it therefore gives SwiGLU an extra matrix.

Also notice the expand-contract shape. Four input dimensions blow up to eight intermediate neurons and then collapse back to four output dimensions. Common matched-budget examples use 4x (or about 2.7x for SwiGLU), giving the network thousands of intermediate "memory slots" per layer.

SwiGLU: the modern FFN

Llama and Mistral use SwiGLU, a gated alternative to the original ReLU FFN. Other models use different gates; Gemma, for example, uses a GELU-based gated FFN.

Shazeer's 2020 GLU study compared gated variants under a matched parameter and compute budget and found improvements on its tested language tasks.

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$ lets one input-dependent projection modulate another. Both factors can be negative, so this is not a Boolean agreement test.

Each intermediate feature now depends on two learned projections. The memory analogy still helps, but the coefficient is a signed product rather than a positive match score.

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.

SiLU and Swish name the same activation used in SwiGLU. A model card mentioning SiLU alone does not establish the full gated architecture; check the projections too.

Attention vs FFN: the division of labor

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 both blocks as reading from and writing to the residual stream. Both contain learned weights and can participate in factual predictions.

For an illustrative prediction of “Paris” after “The capital of France is,” attention can bring information about “capital” and “France” to the final position. An FFN update there can then increase the output score for “Paris.” This describes one possible contribution, rather than a complete account of a trained model’s prediction.

Each layer receives representations changed by the layers before it. Identically shaped FFNs can therefore learn different functions at different depths.

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

The residual path carries the previous vector forward while the FFN adds its update. That update can reinforce or cancel existing features. In the memory interpretation, it is a sum of retrieved value vectors.

Knowledge neurons: finding facts in the weights

Does the key-value interpretation help explain factual predictions in a trained model?

Dai et al. (2022) tested this interpretation with interventions in pretrained BERT models.

They used an integrated-gradients-based attribution method, then refined neurons across prompts expressing the same fact. Suppressing or amplifying selected activations changed the probability of the associated completion. This is not simply selecting whichever neurons have the largest activations.

The selected coordinates participate in those predictions, but need not store independent, readable facts or work for every paraphrase.

ROME and MEMIT modify MLP output weight matrices to change selected factual associations. ROME uses a rank-one matrix update, which can change many entries rather than one neuron or column. Such edits can also have unwanted effects on other prompts.

A model with 32 FFN layers and $d_{ff}=16384$ has 524,288 intermediate coordinates. Those coordinates are not a count of stored facts: features can be distributed across neurons, and one neuron can participate in multiple features. The memory interpretation is a decomposition of the computation, not a literal database schema.

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 a positive value in dimension 0 and a negative value in dimension 2, 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. Their contribution is zero for this input.

ReLU can produce exact zeros, so only a subset of value vectors contributes for a given input. The active fraction depends on the model, layer, and input; the toy example below deliberately uses sparse activations.

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.

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. Training adjusts the key and value matrices through 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 above to compare the coefficients applied to the value columns. SwiGLU multiplies two projections element by element, after applying SiLU to one. Both magnitude and sign affect each column's contribution.

Common misconceptions

"The FFN is just a nonlinearity between attention layers." In the standard $d_{ff}=4d_{model}$ block, its two projections contain twice as many weights as the four attention projections. Interventions show that FFN computation can affect factual predictions, though it also serves many other functions.

"Only FFNs store knowledge." FFNs have a useful memory interpretation, and interventions have localized some factual associations there. Attention's projection weights are learned parameters too. Dynamic attention weights should not be confused with the persistent parameters that compute them; the division into routing and memory is a working model, not an exclusive allocation of knowledge.

"Bigger $d_{ff}$ always helps." Widening the FFN adds parameters and arithmetic. Whether that improves quality at a fixed training budget depends on the rest of the architecture, data, and optimization; the memory analogy cannot answer it.

"The FFN mixes information across the sequence." A batched FFN applies the same weights independently to each position. Processing positions in parallel does not mix their information.

What's next

The layer equation above also includes Norm() calls. Their placement and choice of normalization affect 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: applies a rank-one update to an MLP output matrix to change a factual association.
• 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.