The residual stream

Key idea

x ← x + attention(x), then x ← x + MLP(x), repeated in every layer

Real world

A shared whiteboard in a meeting room. Each person who walks in reads what is there, adds a note, and leaves. Nobody rubs anything out, so the last person sees everyone’s notes at once.

In an LLM

At the last “the” of “The cat sat on the”, the embedding writes “this is the word the”, an attention head adds “it follows on”, another adds “the story is about a cat”, and an MLP adds “floor, mat or bed comes next”. The output layer reads the sum.

Two lines of code

Strip away the details and a transformer layer is two lines:

Here $x$ is one token’s vector (768 numbers in GPT-2 small). Attention (the previous section) moves information between tokens; the MLP processes each token on its own. The important symbol is the plus. Neither layer replaces $x$. Each one computes a change and adds it on. The vector that flows from the embedding at the bottom to the output at the top, collecting all these additions, is the Residual stream.

“Residual” is the name for the thing each layer adds: the change, not the whole new vector. The idea comes from image recognition: He et al., 2015 found that a network where each layer adds a residual can be trained 100 layers deep or more, while a plain stack of the same depth trains worse than a shallow one.

A shared bus

Think of the residual stream as a bus, or a shared whiteboard, that runs past every component. Each component:

1. reads the current vector (in practice it reads a normalised copy, see normalization),
2. computes something from it,
3. writes by adding its result back.

Because they all write to the same place, a later layer can read what any earlier layer wrote. An attention head in layer 5 can use a fact an MLP found in layer 2. Step through it:

The dimensions in this cartoon have names so you can read them. In a real model they don’t: the 768 numbers are not labelled, and a single fact is usually spread over many of them as a direction in the space. But the bookkeeping is exactly this: start from the embedding, add, add, add.

Written out for a model with $N$ layers, the final vector is literally a sum:

where $a_l$ and $m_l$ are what attention and the MLP in layer $l$ wrote. This is why interpretability researchers can ask “which layer contributed this prediction?”: you can take the sum apart.

Measured in GPT-2

Here is the real thing. The top chart is the length (the norm) of one token’s residual vector after every attention and MLP step in GPT-2 small. The bars below are the lengths of what each layer wrote.

Three things stand out:

• It grows. Every token starts as a vector of length about 5 and ends at several hundred. Since layers only add, the stream accumulates.
• Late layers write the most. The early and middle layers make small, careful edits. The last two layers write much bigger vectors: they are shaping the final prediction.
• The first token is strange. Pick “The” in any sentence. GPT-2 writes a vector of length about 3,000 there in its first three layers (mostly the MLPs of layers 1 and 2), leaves it nearly alone, and the attention of the last layer then cancels most of it. Many attention heads use the first token as a resting place when they have nothing better to attend to (an “attention sink”, see attention patterns), and this huge vector is part of how that works.

Because the stream keeps growing, every layer reads a normalised copy of it rather than the raw vector. Otherwise the late layers would see vectors several times longer than the early layers did, and every layer would have to cope with a different scale. That is the job of the next page, normalization.

Why the plus makes deep models trainable

Training adjusts each weight using its gradient: how much the loss would change if that weight changed. The gradient is computed backwards from the loss through every layer (backpropagation).

In a plain stack, $x_{l+1} = f(x_l)$, the gradient reaching layer 1 is a product of one factor per layer. Multiply 48 numbers that are each a bit below 1 and you get almost nothing; a bit above 1 and it explodes. Early layers either stop learning or blow up.

With a residual, $x_{l+1} = x_l + f(x_l)$. The derivative of $x_{l+1}$ with respect to $x_l$ is $1 + f'(x_l)$: there is always a “1” in it, a path where the gradient passes straight through unchanged. Stack $N$ layers and there is a clean path from the loss all the way down to the embedding. Here is a simulated 16-wide network with and without the plus. Press play and watch the gradient travel down from the loss:

The same widget appears again in backpropagation, where the chain rule behind it is worked through.

It also means a new layer starts life doing almost nothing harmful. With small initial weights, $f(x)$ is close to zero and the layer is close to the identity. Training then teaches it to add something useful. A deep model is never worse at the start than a shallow one, which is the problem residuals were invented to fix.

class Block(nn.Module):
    def __init__(self, d, n_heads):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = CausalSelfAttention(d, n_heads)
        self.mlp = MLP(d)
 
    def forward(self, x):                 # x: (batch, T, d)
        x = x + self.attn(self.ln1(x))    # read a normalised copy, add back
        x = x + self.mlp(self.ln2(x))
        return x