The residual stream — Linear algebra for LLMs

One transformer block

\displaystyle x \leftarrow x + \operatorname{Attn}(\operatorname{Norm}(x))

\displaystyle x \leftarrow x + \operatorname{MLP}(\operatorname{Norm}(x))

Each sublayer reads a normalized copy of $x$ and adds its result back. This “pre-norm” layout is what GPT-2 and Llama use. $x$ keeps its shape $(B, T, d)$ through every block.

Write

Block:

x \mathrel{+}= \operatorname{Attn}(\operatorname{Norm}(x))

, then

x \mathrel{+}= \operatorname{MLP}(\operatorname{Norm}(x))

.

Unrolled, it's one big sum

\displaystyle x_{\text{final}} = x_0 + \sum_{\ell=1}^{L} \big(a_\ell + m_\ell\big)

+

=

The final vector is the embedding plus what every attention ( $a_\ell$ ) and MLP ( $m_\ell$ ) sublayer wrote. Nothing is overwritten; information is only added. (The numbers are Part 1’s toy example.)

A shared space that everyone reads and writes

In a transformer

Every matrix that reads the stream ( $W_Q, W_K, W_V$ , the MLP’s up and gate projections) projects it; every matrix that writes ( $W_O$ , the MLP’s down projection) maps back into the same $d$ dimensions. Layers communicate by writing into directions that later layers read: the projections of Part 3. Interpretability work studies the stream exactly this way.

Why add instead of replace

\displaystyle \frac{\partial\,(x + f(x))}{\partial x} = I + \frac{\partial f}{\partial x}

The identity $I$ gives the gradient a straight path back to every earlier layer, whatever $f$ does. Without it, gradients would have to pass through every layer’s Jacobian (Part 8) in turn and could shrink or explode over dozens of layers.

Write

\partial(x + f(x))/\partial x = I + \partial f/\partial x

: a direct gradient path.

✏️ Your turn

Count the writers

A model has 32 blocks, each with one attention and one MLP sublayer. How many vectors are added together in its final residual stream, counting the embedding?

Work it out on paper, then go to the next card to check.

Answer: 65

\displaystyle 1 + 2 \times 32 = 65

The embedding plus 64 sublayer outputs. A final norm and the unembedding then read the result.

📌 Recap

The residual stream in one card

$x \mathrel{+}= \operatorname{Attn}(\operatorname{Norm}(x))$ ; $x \mathrel{+}= \operatorname{MLP}(\operatorname{Norm}(x))$ .
Final $x$ = embedding + every sublayer’s output.
All layers read from and write into the same $d$ -dimensional space.
The identity path keeps gradients flowing.

Write

The four lines above.