Skip to content
ginsberg.si

Transformers from the Weights Up/Part 5 of 20

Transformers from the Weights Up, Part 5: MLPs and what the FFN stores

Each FFN hidden unit is a memory slot: gate and up rows are keys, and a down_proj column is the value written back.

4 Nov 2026

Matrix-vector multiplication is a weighted sum of columns: W_down h = Σ hᵢ · W_down[:, i]. So call vᵢ = W_down[:, i] (a [4096] vector in Llama-3-8B) the value of slot i, and aᵢ(x) = hᵢ = SiLU(gᵢ·x) · (uᵢ·x) its coefficient, where gᵢ and uᵢ are row i of gate_proj and up_proj. Nothing has been approximated: this is the same computation as Lesson 3's SwiGLU, regrouped. Image 5's MLP is exactly this picture, 14,336 times over, in each of 32 layers.

IMAGE 05input xwidth dup-projectionactivation φwidth d_ffdown-projectionoutput ywidth dy = W_down · φ(W_up x)
Image 5. Transformer FFN shape: up-projection, a nonlinearity, then down-projection.
(v) An FFN as a bank of key–value memories, one slot at a time
slotg·xu·xa = SiLU(g·x)·u·x1. landmark + "is in"1.701.201.722. British context-0.120.14-0.013. function word next-0.100.31-0.014. tower-name continuation0.260.690.105. weak noise slot0.200.140.026. query, not a name1.17-0.84-0.75FFN output: nothing written yetdim 00.000dim 10.000dim 20.000dim 30.000dim 40.000

Toy model (weights and slot names made up for illustration): d = 5, d_ff = 6 slots. Each slot i has two keys, a gate row gᵢ and an up row uᵢ, and one value, column vᵢ of W_down. The arithmetic is real and computed live: coefficient aᵢ = SiLU(gᵢ·x)·(uᵢ·x), output = Σ aᵢ vᵢ. Watch the first slot dominate, four slots barely fire, and the last slot fire with a negative coefficient, which subtracts its value.

Keys, values, and the tensors they live in

RoleTensor in the checkpointSlice for slot iShape
Gate key (does the slot fire?)model.layers.ℓ.mlp.gate_proj.weight[i, :][4096]
Up key (signed strength)model.layers.ℓ.mlp.up_proj.weight[i, :][4096]
Value (what gets written)model.layers.ℓ.mlp.down_proj.weight[:, i][4096]

Count the memory. Llama-3-8B has 14,336 slots per layer × 32 layers = 458,752 slots. Each holds 3 × 4096 = 12,288 weights, or 24,576 bytes in bf16. All of them together are 5.64 billion parameters, and at batch-1 decode they are 11.27 GB of the 15.01 GB that Lesson 4 showed must stream through the chip for every token: about 75% of the bandwidth bill is spent reading memories.

(w) Where the parameters live
26.9%26.9%26.9%19.2%FFN gate: 58.72MFFN up: 58.72MFFN down: 58.72Mattention q/k/v/o: 41.94M2 RMSNorms: 8,192
One Llama-3-8B block
218,112,000 parameters
FFN share
176.16M = 80.8%

Exact counts from the published configs. Llama-3-8B: d = 4096, d_ff = 14336, 32 layers, 8 KV heads × 128, vocab 128,256, untied embeddings, no biases; total 8,030,261,248. GPT-2 small: d = 768, d_ff = 3072, with biases. Geva et al. (2021) open with the GPT-2-era fact that FFNs are two-thirds of the parameters; SwiGLU at d_ff = 3.5d pushes it past four-fifths of a Llama-3 block.

What the evidence says the slots hold

Geva, Schuster, Berant and Levy (2021) took this regrouping literally. For each key they collected the training prefixes that activated it most, and for each value they projected it through the output embedding to see which next tokens it pushes up. Their findings, from the paper's abstract:

  1. 1

    Keys correlate with human-readable input patterns

    Top-activating prefixes for a key share something nameable. Lower layers tend to capture shallow patterns, such as a particular word or n-gram ending; upper layers capture more semantic ones, such as a topic.

  2. 2

    Values induce distributions over the vocabulary

    Projecting vᵢ through the output embedding concentrates probability on tokens likely to come right after the key's pattern, especially in the upper layers.

  3. 3

    The output is a composition of memories

    No single slot decides the prediction. Many fire together, their weighted sum is added to the stream, and later layers refine it through the residual connection.

  4. 4

    Updates promote concepts (Geva et al., 2022)

    Viewing each FFN update in vocabulary space, the sub-update from a single value vector often promotes a human-interpretable concept. Scaling chosen sub-updates reduced GPT-2's toxicity by almost 50%, and a simple early-exit rule saved 20% of compute on average.

(x) Reading a value vector in vocabulary space
W_U · v1: what this value says about the next tokenP(next token): before FFN (outline) vs after (filled)Paris+1.32Paris42% → 79%France+0.80France22% → 16%London-0.98London6% → 0%tower+0.36tower11% → 2%the+0.02the12% → 2%is-0.10is7% → 1%
Slot 1 coefficient
1.725 × 1.00 = 1.725
Top token after FFN
Paris (79%)

Same toy model as (v); the 6-token unembedding W_U is made up for illustration, and every number is computed live. The left bars are W_U · vᵢ, which tokens slot i's value promotes or suppresses on its own (Geva et al. 2022 do this with GPT-2's real 50,257-row embedding matrix). The right bars are next-token probabilities from the stream before and after the whole FFN writes into it. Scaling one slot's coefficient is a crude version of the interventions those papers run.

Facts in the middle layers, and editing them

Meng, Bau, Andonian and Belinkov (2022) asked where a model recalls a fact such as "The Space Needle is in Seattle." Their causal tracing corrupts the subject tokens, then restores one hidden state at a time from a clean run and measures how much of the right answer comes back. The decisive states sit in middle-layer feed-forward modules while the model processes the subject tokens. That fits the key–value reading: the subject, at its last token, is the key; the attribute is written as the value.

If a fact is a key-to-value mapping inside W_down, you can insert a new one with a single rank-one update. That is ROME (Rank-One Model Editing):

k* is the key the subject produces, v* is a value chosen by optimization so that the model outputs the new attribute, and C = KKᵀ is an uncentered covariance of the keys this layer normally sees, estimated from sample text. In ROME's framing the "key" is the whole [14336] activation vector h at the subject's last token, and W_down is a linear associative memory mapping it to a [4096] value: the same matrix, viewed as one memory rather than 14,336 slots. The edit maps k* exactly to v* (check: Ŵk* = Wk* + Λ · (C⁻¹k*)ᵀk* = Wk* + (v* − Wk*) = v*). Every other key k moves by Λ · (C⁻¹k*)ᵀk, which is small when k is unrelated to k* in the metric C defines.

(y) Writing a new memory with a rank-one edit
key space (input to W)value space (output of W)k*k1k2k3v*
Edit applied
0%
k1: cos with k*, output drift
0.76, 0.000
k2: cos with k*, output drift
-0.45, 0.000
k3: cos with k*, output drift
-1.00, 0.000

Toy 2-D W_down-like matrix (made up for illustration; arithmetic live). Left: three existing keys and the new key k* (angle set by the slider). Right: each key's output W·k before the edit (dashed) and after (solid), and the target v*. The update Ŵ = W + (v* − Wk*) k*ᵀ / (k*ᵀk*) is ROME's formula with the key covariance C set to the identity; it maps k* exactly to v*, and moves every other output by (k*·kⱼ / |k*|²)(v* − Wk*). In 2-D, keys overlap and the edit bleeds; in 4096-D, unrelated keys are nearly orthogonal (Lesson 2), and ROME's C⁻¹ shrinks the bleed further.

Where the clean picture breaks

Clean storyWhat is actually trueWhere it comes back
One slot, one patternMany hidden units are polysemantic: one unit fires for several unrelated patterns, because the model stores more features than it has units (superposition).Lesson 20: sparse autoencoders recover cleaner features from mixtures of units
Slots are independentCoefficients are computed from the same stream, and many slots fire at once; meaning is often in the combination, not one slot.Lesson 20: circuits
Facts live only in FFNsAttention moves information between positions (it carries the subject's attribute to the last token); FFNs transform it in place. Recall needs both.Lessons 8–11
The value is the answerValues are added to a stream that later layers keep rewriting; the logit lens shows a vote, not the final outcome.Lessons 11–12

For engineering, the key–value view still pays its way. It explains why ffn_down is one of the tensors quantization recipes treat most carefully (its 14336-wide input carries the rare large coefficients SwiGLU produces, Lesson 3), why mixture-of-experts models split exactly this block into many smaller banks of slots and pick a few per token (Lesson 19), and why pruning or skipping inactive slots is a recurring inference trick.

Lesson 5: key numbers

14,336
slots per Llama-3-8B layer= d_ff
458,752
slots in the whole model14,336 × 32 layers
12,288
weights per slotgate row + up row + down column; 24 KiB in bf16
≈ −0.278
SiLU minimumat z ≈ −1.28; the gate is nearly off below zero
80.8%
FFN share of a Llama-3-8B block176.2M of 218.1M
70.2%
FFN share of all 8.03B params5.64B; embeddings and lm_head take 13.1%
≈ 2/3
FFN share in GPT-2-style modelsGeva et al. 2021; d_ff = 4d, two matrices
11.27 GB
FFN bytes per bf16 decode step75% of Lesson 4's 15.01 GB

Lesson 5: self-check

1. In Llama-3-8B, which exact slices hold the keys and the value of FFN slot 5000 in layer 20?

Answer. Keys: model.layers.20.mlp.gate_proj.weight[5000, :] and model.layers.20.mlp.up_proj.weight[5000, :]. Value: model.layers.20.mlp.down_proj.weight[:, 5000]. All three are [4096]. Note the value is a column, because PyTorch stores weights as [out, in] (Lesson 2).

2. Why is writing FFN(x) = Σ aᵢ(x)·vᵢ exact rather than an interpretation?

Answer. Because W_down · h is by definition Σ hᵢ · W_down[:, i]. The key–value view regroups the same arithmetic; the interpretation lies in claiming that individual keys and values mean something, which is the empirical part.

3. A ReLU-FFN slot can never write −vᵢ. A SwiGLU slot can. Why?

Answer. ReLU(k·x) ≥ 0, so a ReLU slot's coefficient is non-negative. In SwiGLU the coefficient is SiLU(g·x)·(u·x); the up score u·x can have either sign (and SiLU itself dips to about −0.278), so the slot can subtract its value.

4. With C = I, which existing memories does a rank-one edit disturb, and by how much?

Answer. Ŵk = Wk + (k*·k / |k*|²)(v* − Wk*). Only keys with a component along k* move, in proportion to their overlap. In 4096-D unrelated keys have near-zero overlap, so they barely move; ROME's C⁻¹ further steers the edit away from directions the layer uses often.

5. Why do FFN weights dominate the cost of generating one token at batch 1?

Answer. Decode is bandwidth-bound (Lesson 4), so cost tracks bytes read. FFN weights are 5.64B of the 7.505B streamed parameters: 11.27 GB of 15.01 GB in bf16, about 75%. Shrinking or skipping FFN slots is the biggest lever on decode speed.

6. Causal tracing finds factual recall in middle-layer FFNs at the subject's last token. Why the subject token and not the final token?

Answer. The FFN at a position can only read the stream at that position. The subject's last token is where the full subject (e.g. 'Space Needle') has been assembled, so that is where the key can match and the attribute value can be written. Attention then moves it to the final position, where the answer is predicted.

Lesson 5: resources