1.3 The Full Block & the Residual Stream

Learn

One block, two writes

Module 1.2 gave you attention. A transformer block is attention plus one more sublayer, wired together in a specific way that turns out to matter more than either piece:

\begin{aligned} x &\leftarrow x + \mathrm{Attn}(\mathrm{LN}(x)) \\ x &\leftarrow x + \mathrm{MLP}(\mathrm{LN}(x)) \end{aligned}

Look at what those two lines are not. Neither says $x \leftarrow f(x)$ . Both say $x \leftarrow x + f(\ldots)$ . The sublayer never replaces the running vector; it computes a correction and adds it. That single design choice — the residual connection, borrowed from ResNets — is what this whole module is about.

One block. The straight line through the middle is the residual stream; it is never transformed, only added to. Each sublayer branches off, reads a normalised copy, computes, and merges its result back with a plus.

LayerNorm

LayerNorm standardises a vector across its own features — not across the batch, not across the sequence — then applies a learned gain and bias:

\mathrm{LN}(x) = \gamma \odot \frac{x - \mu(x)}{\sigma(x)} + \beta, \qquad \mu = \tfrac{1}{d}\textstyle\sum_t x_t,\; \sigma = \sqrt{\tfrac{1}{d}\sum_t (x_t-\mu)^2 + \epsilon}

It exists to keep the stream's scale in a range where gradients behave. But it has two consequences interpretability cares about enormously, and both come from what it throws away: subtracting $\mu$ deletes the component along the all-ones direction, and dividing by $\sigma$ deletes the overall magnitude. Whatever a layer writes along $(1,1,\ldots,1)$ is invisible to every later reader; and “how loudly” the stream is speaking is normalised away before anyone listens.

RMSNorm, and why pre-LN won

Modern models (Llama, most of the open-weight world) use RMSNorm, which drops the mean subtraction and the bias and just divides by the root-mean-square. It is cheaper and works just as well, which is decent evidence the centring was never doing much. Separately: the original 2017 paper put LayerNorm after the addition (post-LN). Essentially everything since GPT-2 puts it before the sublayer (pre-LN), because post-LN sits directly on the residual path and breaks the clean gradient highway — post-LN models need learning-rate warmup to train at all.

Where position comes from

Attention is permutation-equivariant: shuffle the tokens and the outputs shuffle with them. Something has to break that symmetry. GPT-2 uses learned absolute positional embeddings — a second lookup table with one vector per position, added to the token embedding at the very bottom. Simple, and it hard-caps the context at however many rows the table has.

Nearly every model since 2022 uses RoPE (rotary position embeddings). Instead of adding anything to the stream, RoPE rotates the query and key vectors inside each head by an angle proportional to position, in 2-D slices. The consequence is the whole point: after rotating $q_m$ by $m\theta$ and $k_n$ by $n\theta$ , their dot product depends only on $m - n$ — the attention score sees relative position for free, and there is no table to run off the end of.

Key idea

Positional information enters the model in exactly one place, and it is a place you can point at. In GPT-2 it is a row of a table added at the input; with RoPE it is a rotation applied inside the QK circuit and nowhere else — RoPE never touches the values, so it changes where heads look without changing what they move.

Learn

The residual stream is the object

Unroll the recursion. A 12-layer model's final vector at some position is not the output of layer 12. It is a sum:

x_{\text{final}} = \underbrace{x_{\text{emb}}}_{\text{token + position}} + \sum_{\ell=1}^{L}\Big( \underbrace{\textstyle\sum_{h} \mathrm{head}^{(\ell,h)}}_{\text{each head, separately}} + \underbrace{\mathrm{MLP}^{(\ell)}}_{\text{}} \Big)

For GPT-2 small that is 1 embedding + 144 head contributions + 12 MLP contributions, all added into the same 768-dimensional vector. Nothing was ever overwritten. Every term is still, in principle, recoverable.

Key idea

Stop thinking of a transformer as a pipeline where layer

\ell

transforms the output of layer

\ell-1

. Think of it as a shared bus that every component reads from and writes to. Layers are not stages; they are subroutines making incremental additive edits to a single running state. This reframing is due to Elhage et al. (2021) and it is the foundation of every circuit result in Part 3.

The stream as a communication channel, not a pipeline. A head in layer 2 can write a direction that nothing reads until layer 9 — the intervening layers simply add alongside it. The two components are then in communication, however far apart they are.

Virtual weights

Because the stream is linear and additive, there is an implied weight matrix between any two components, even though no such matrix appears in the checkpoint. If head $A$ in layer 1 writes $W_{OV}^{A} x$ and head $B$ in layer 4 reads through $W_Q^{B}$ , then the composed map

W_Q^{B} \, W_{OV}^{A}

is a real, computable matrix describing exactly how much $A$ 's output steers $B$ 's queries. Elhage et al. call these virtual weights, and the phenomenon composition (Q-, K-, or V-composition depending on which of $B$ 's inputs is affected). You can compute the virtual weight between any pair of components in a trained model without running it on a single token.

bandwidth

The residual stream has a fixed width — 768 floats in GPT-2 small — shared by 157 writers and as many readers. Two components can communicate privately only by using directions nothing else uses, and there are not enough orthogonal directions to go around. So the stream is crowded, and components must tolerate reading each other's interference. That pressure is exactly what Module 3.3 calls superposition.

A privileged basis, or not?

Whether individual residual-stream coordinates mean anything is a live question. Nothing in the architecture prefers the standard basis — attention and the linear layers are rotation-equivariant, so in principle only directions matter, not axes. Yet in practice the stream has outlier coordinates with enormous magnitude, and Anthropic attributes this to the optimiser (Adam is per-coordinate) rather than to anything meaningful in the representation. Treat basis-aligned claims about the residual stream with suspicion; treat direction-based claims as the default.

Safety tie-in

Almost every practical interpretability and steering technique is downstream of additivity. Linear probes work because a concept can be a direction in the stream. Steering vectors (Module 5.1) work because adding a vector to the stream is exactly the operation the architecture is built from — you are not hacking the model, you are doing what layer 7 does. Direct logit attribution works because you can push each of the 157 terms through the unembedding separately and ask which ones voted for the answer. If the architecture composed its layers multiplicatively, auditing a model would be dramatically harder than it already is.

Learn

MLPs: two thirds of the parameters

The other sublayer is almost embarrassingly plain — one hidden layer, no bells:

\mathrm{MLP}(x) = W_{\text{out}} \, \mathrm{GELU}(W_{\text{in}} x + b_{\text{in}}) + b_{\text{out}}

with $W_{\text{in}}$ of shape $d_{\text{mlp}} \times d_{\text{model}}$ and $d_{\text{mlp}} = 4 d_{\text{model}}$ by convention. It acts on one position at a time, sees nothing else, and contains roughly two thirds of every parameter in the layer stack. Whatever a transformer knows, most of it is stored here.

The useful reframing — from Geva et al. — is to read the two matrices as a lookup table. Write $W_{\text{in}}$ as a stack of rows $k_i$ and $W_{\text{out}}$ as a stack of columns $v_i$ . Then

\mathrm{MLP}(x) = \sum_{i=1}^{d_{\text{mlp}}} \underbrace{\mathrm{GELU}(k_i \cdot x + b_i)}_{\text{how much key } i \text{ matched}} \; \underbrace{v_i}_{\text{what neuron } i \text{ writes}}

Key idea

An MLP is a key-value memory with

d_{\text{mlp}}

entries. Each neuron holds a key direction that pattern-matches against the residual stream, and a value direction it adds to the stream in proportion to how well the key matched. Attention decides which facts are available; MLPs are where the facts are.

Geva et al. found the keys are often human-legible — neurons that fire on a topic, a template, a language — and the values frequently promote a coherent set of next tokens. Module 5.2 (ROME) uses exactly this picture to locate and rewrite a specific fact.

Neurons are not features

The tempting next step — “so we just read off what each of the 3,072 neurons means” — does not work. Most neurons are polysemantic: they fire on several unrelated things. The reason is superposition (Module 3.3): a model with 3,072 neurons represents far more than 3,072 features, so features get packed into overlapping combinations rather than one per neuron. Sparse autoencoders (Module 3.4) are the field's attempt to undo that packing. The key-value picture is the right mechanism; the individual neuron is the wrong unit.

What happened to 4×

Modern models mostly use a gated MLP — SwiGLU — with three matrices instead of two:

W_{\text{down}}\big(\mathrm{SiLU}(W_{\text{gate}} x) \odot (W_{\text{up}} x)\big)

. To hold the parameter count fixed against the old two-matrix design,

d_{\text{mlp}}

drops from

4 d_{\text{model}}

to about

\tfrac{8}{3} d_{\text{model}}

— which is why Llama-family configs have hidden sizes like 11008 rather than a round multiple of four. The key-value reading survives the change; there are just two keys per entry now, one of which gates the other.

Learn

The logit lens: reading the stream mid-flight

Here is the payoff of additivity. The model's final step is LayerNorm and then multiplication by the unembedding matrix $W_U$ . But the residual stream at layer 5 lives in exactly the same space as the residual stream at layer 12 — same basis, same width, same units. So nothing stops you applying that final step early:

\text{logit-lens}_\ell = \mathrm{softmax}\big(W_U \, \mathrm{LN}_{\text{final}}(x_\ell)\big)

You get a distribution over the vocabulary for every layer: the model's prediction if you cut the remaining layers off. This is the logit lens (nostalgebraist, 2020), and it is the cheapest useful interpretability tool that exists — three lines of code, no training.

Illustrative logit-lens trace for a factual-recall prompt. Hand-drawn schematic of the pattern such runs typically show, not output from a live model — you will run the real thing in Module 3.1. Note the shape: a long flat stretch, then a jump at one MLP, then sharpening.

Read that shape. For six layers the model is not “gradually becoming more confident” — it has no idea, and its top guess is a generic function word. Then one MLP fires and the answer appears. The remaining layers sharpen and clean up. Prediction in transformers tends to be lumpy and event-like, not smooth.

The step-through widget below walks exactly this narrative, one sublayer at a time.

The logit lens is not a neutral instrument

It works far better on some models than others — on many it produces nonsense in the early layers, and on some it never works at all. The reason is that

W_U

was trained to read the final layer's stream, and intermediate layers may use a different effective basis. The tuned lens (Belrose et al., 2023) fixes this by training a small affine probe per layer, and it is both more faithful and less honest-looking: you are now reading the stream through something you fit yourself, so “the model believed X at layer 5” becomes a claim about your probe as much as about the model. Both lenses are observational. Neither tells you that the layer's output was used.

Explore

Play: watch the stream fill, then count the cost

The first widget follows one token position through a 4-layer model, one sublayer at a time: what reads, what writes, how the stream accumulates, and what the logit lens says at each step. The second answers the question every architecture diagram dodges — given a config, where do the parameters actually go?

The residual stream, one write at a time

Following a single token position — the final “of” in “The Eiffel Tower is in the city of” — through a 4-layer model. Every sublayer reads the stream, computes, and adds its result back. Nothing is ever overwritten.

embed + position — The stream starts as the token embedding of “of” plus its positional information. It knows what word it is and where it sits. Nothing else.

The stream as a sum of writes

Each segment is one sublayer's contribution. The final stream is literally their sum — which is why you can subtract any one of them out and ask what the model would have predicted without it.

Logit lens: decode the stream here

“ the”

12%

“ a”

6%

“ his”

4%

Apply the final LayerNorm and the unembedding to the partial stream, as if the remaining layers did not exist. The answer does not fade in smoothly — it arrives.

Step 1 of 10: embed + position. Total written into the stream so far: 1.00 (arbitrary units) across 1 contributions. The logit-lens readout is hand-authored and illustrative — it shows the shape real logit-lens runs have on factual recall, not a live model.

Where the parameters live

GPT-2-family arithmetic: learned positional embeddings, pre-LayerNorm blocks with biases, and a tied unembedding (so the output matrix is free). Move the sliders and watch the stacked bar redistribute.

d_model768

n_layers12

n_heads12 (d_head 64)

d_mlp3072 (4.0× d_model)

vocab50,257

n_ctx (positions)1,024

124.4M

124,439,808 parameters

token embedding38,597,37631.0%

positional embedding786,4320.6%

attention28,348,41622.8%

MLP56,669,18445.5%

LayerNorm38,4000.0%

Two things worth noticing. Moving n_heads does not change the total — heads split a fixed budget of 4·d_model², they do not add to it. And inside the layer stack, MLP beats attention 2.00:1 — at the standard 4× expansion it is always exactly 2:1, before biases.

Total: 124,439,808 parameters = 124.4M. Head width d_head = 64. Per layer: 2,362,368 attention + 4,722,432 MLP + 3,072 LayerNorm. Published count for GPT-2 small: 124,439,808 ✓ exact match.

Things to try: (1) Step the flow widget through once and watch the “sum of writes” bar — notice that the embedding is a small fraction of the final stream by the end, which is why late-layer representations barely resemble the token that produced them. (2) In the calculator, drag n_heads from 1 to 32 and confirm the total does not move: heads split a fixed budget of $4 d_{\text{model}}^2$ rather than adding to it. (3) Set d_model to 256 and push n_layers to 48, then do the opposite — d_model 2048, n_layers 4. Same rough total, wildly different models. Watch what happens to the embedding's share in each case: for small models the vocabulary table dominates everything, which is why a 50k-token vocabulary is a real design constraint at small scale and an afterthought at frontier scale.

Practice

Problem set

The parameter count is the one to do properly, by hand, before touching the widget. Getting it exactly right — to the last of the eight digits — means you understand the architecture with no gaps, and there is no other exercise in Part 1 that checks that as ruthlessly.

1.Count GPT-2 small's parameters by handpencil & paper

GPT-2 small: $d_{\text{model}} = 768$ , $n_{\text{layers}} = 12$ , $n_{\text{heads}} = 12$ , $d_{\text{head}} = 64$ , $d_{\text{mlp}} = 3072$ , $n_{\text{vocab}} = 50257$ , $n_{\text{ctx}} = 1024$ . Every linear layer has a bias; there are two LayerNorms per block plus one at the end, each with a gain and a bias; the unembedding is tied to the token embedding.

Produce the exact integer, broken down by category. Then answer: what fraction sits in MLPs, and what fraction of the non-embedding parameters sit in MLPs? Check your answer against the published count for the released checkpoint.

2.What LayerNorm destroyspencil & paper

Let $\mathrm{LN}(x) = (x - \mu(x))/\sigma(x)$ (ignore $\gamma, \beta, \epsilon$ ). Prove: (a) $\mathrm{LN}(x + c\mathbf{1}) = \mathrm{LN}(x)$ for any scalar $c$ ; (b) $\mathrm{LN}(ax) = \mathrm{LN}(x)$ for any $a > 0$ .

Then: what does (a) imply about a component that writes along the all-ones direction? What does (b) imply about trying to interpret the magnitude of a head's contribution? And why do interpretability libraries offer a “fold LayerNorm” option?

3.Virtual weights and their rankpencil & paper

Head $A$ in layer 1 writes $W_O^A W_V^A x$ into the stream, where $W_V^A$ is $d_{\text{head}} \times d_{\text{model}}$ and $W_O^A$ is $d_{\text{model}} \times d_{\text{head}}$ . Head $B$ in layer 4 forms its queries with $W_Q^B$ .

(a) Write the matrix describing the total effect of $A$ 's output on $B$ 's queries. (b) What is its maximum rank, for GPT-2 small numbers? (c) Ignoring LayerNorm, why is this map exact rather than approximate, even though layers 2 and 3 sit in between? (d) What is the name for a composition that changes $B$ 's keys instead?

4.A full transformer block in NumPycode

Reusing your multi-head attention from Module 1.2, write block(x, params) implementing pre-LN GPT-2: LayerNorm → attention → add; LayerNorm → MLP (with GELU) → add. Then stack $L$ of them, add embeddings at the bottom, and finish with a final LayerNorm and a tied unembedding.

Success checks: (1) load the real gpt2 weights (HuggingFace transformers or transformer_lens) and reproduce the reference logits for a short prompt to within 1e-3; (2) assert your parameter count is exactly 124_439_808; (3) verify empirically that adding a constant vector $c\mathbf{1}$ to the residual stream before a block changes nothing downstream.

5.Run the logit lens for realcode

In TransformerLens, run gpt2-small on The Eiffel Tower is in the city of with run_with_cache. Pull resid_post for every layer, apply model.ln_final and then model.unembed, and plot $P(\text{“ Paris”})$ against layer. Print the top-3 tokens at each layer.

Then two extensions. (1) Do the same for a prompt requiring syntax rather than a fact — e.g. subject-verb agreement across a clause — and compare the shapes of the two curves. (2) Use cache.decompose_resid() or accumulated_resid to check that the per-component contributions really do sum to the full stream.

Success check: your layer-12 distribution matches the model's actual output distribution exactly, and the sum of decomposed components matches resid_post[-1] to floating-point tolerance.

6.Meet a real MLP neuronexplore

Open Neuronpedia and browse GPT-2 small MLP neurons (not SAE features — those come in Module 3.4). Pick three from different layers.

For each, record: the top activating text snippets, whether you can state a one-sentence hypothesis for what the key direction detects, and what the neuron's top positive logit contributions are. Then classify each neuron as monosemantic-looking, clearly polysemantic, or illegible, and count how many of each you found.

0 of 6 problems marked done

Check

Check yourself

1.

What is the best way to think about what layer 7 does to the residual stream?

2.

In GPT-2 small, roughly what fraction of the non-embedding parameters live in MLPs?

3.

A component writes a vector proportional to

\mathbf{1} = (1,1,\ldots,1)

into the residual stream. What do later sublayers see?

4.

What is a “virtual weight”?

5.

The key-value memory view says an MLP computes

\sum_i \mathrm{GELU}(k_i \cdot x)\, v_i

. What are

k_i

and

v_i

?

6.

You run the logit lens and find the correct answer's probability jumps from 0.02 to 0.31 at layer 6. What may you conclude?

7.

What does RoPE change about how position enters the model, compared with GPT-2's learned positional embeddings?

Answer all 7 questions to submit.

Submit your answers to complete this check.

Go deeper

The Elhage framework is the one that changes how you see the architecture — read only the residual-stream sections now, and let Module 3.2 handle the rest. Karpathy is the do-along that makes all of Part 1 concrete.

EssentialA Mathematical Framework for Transformer Circuitspaper

Elhage, Nanda, Olsson et al. (Anthropic) · 2021 · 1.5h

Read three sections and stop: "Transformer Overview", "The Residual Stream as a Communication Channel" (including the subsections on virtual weights and subspaces/bandwidth), and "Attention Heads are Independent and Additive". This is where the reframing in this module comes from, stated better than anyone has restated it since. Skip the two-layer induction analysis — that is Module 3.2's job.

EssentialLet's build GPT: from scratch, in code, spelled outvideo

Andrej Karpathy · 2023 · 2h (do-along)

Finish it here if you started it in Module 1.2. The last third — where he adds residual connections and LayerNorm to a network that was not training well, and it suddenly does — is the best available demonstration of why the residual stream exists at all. Type every line.

Transformer Feed-Forward Layers Are Key-Value Memoriespaper

Geva, Schuster, Berant & Levy · 2021 · 50 min

The source of this module's MLP framing. Read §2 (the formal key-value equivalence) and §3 (human evaluation of what the keys detect). Skim §4–5. The critical caveat to carry forward: their analysis treats one neuron as one memory slot, which superposition (Module 3.3) shows is not generally true.

interpreting GPT: the logit lensblog

nostalgebraist · 2020 · 30 min

The original post, still the clearest statement of the idea. Read for the plots rather than the code. Pay attention to where the author flags that it does not work well — that honesty is the reason the tuned lens exists.