Kiru Lab  /  Language Models  /  Anatomy of a Transformer

The Residual Stream and the Transformer Block

The most useful mental model in interpretability: a shared bus that every layer reads from and writes to.

Concept  ·  about 35 minutes

A transformer block does two things in sequence, each wrapped in a residual connection: attention moves information between positions, then an MLP transforms information within a position.

x = x + Attention(LayerNorm(x)) x = x + MLP(LayerNorm(x))

Because every sublayer adds to `x` rather than replacing it, `x` is best understood not as "the activations at layer n" but as a running sum — a communication channel that every component reads from and writes into. This is the residual stream, and it is the single most productive frame for interpretability work.

  • Each component reads a subspace of the stream, computes, and writes its result back additively.
  • Because the operations are additive, a final output can be decomposed into per-component contributions — which is what makes techniques like logit attribution possible at all.
  • Components in different layers communicate by writing to and reading from the same directions, forming circuits across depth.

What the MLP is for

The MLP is two linear layers with a nonlinearity, and it holds roughly two-thirds of a transformer's parameters. Evidence suggests it functions substantially as key-value memory: the first layer detects patterns in the residual stream and the second writes associated content back. A great deal of a model's factual knowledge appears to live here rather than in the attention layers.

Structure dictates function

Attention moves information across positions; MLPs transform it in place. That division is architectural, and it predicts where to look for different kinds of behavior — routing failures in attention, factual and associative failures in MLPs.

Hold on to

  • The residual stream is an additive bus shared by every component
  • Additivity is what makes per-component attribution possible
  • Attention routes between positions; MLPs store and transform within one

Work through

Try each one before opening the solution. Getting it wrong first is most of where the learning happens.

  1. Hook a small open model and record residual stream norms at every layer. Plot growth across depth.
    Hint

    Register a forward hook on each block and record the norm of its output.

    Solution

    The residual stream norm grows steadily with depth, often roughly linearly or faster, because every sublayer adds to it and nothing subtracts. This is why raw activation magnitudes are not comparable across layers and why layer norm is applied before each sublayer reads the stream — a fact that has confused a lot of interpretability results by people who did not account for it.

    import torch

norms = []
handles = [block.register_forward_hook(
    lambda m, i, o, store=norms: store.append(float(o[0].norm(dim=-1).mean())))
    for block in model.blocks]

with torch.no_grad():
    model(input_ids)
for h in handles:
    h.remove()

for layer, n in enumerate(norms):
    print(layer, round(n, 2))
  2. Zero out one MLP layer and measure the change in output distribution. Compare against zeroing an attention layer.
    Hint

    Zero the sublayer output so only the residual path passes through, and measure KL divergence.

    Solution

    Ablating a single MLP layer usually shifts the output distribution more than ablating a single attention layer in the middle of the network, consistent with MLPs holding most of the parameters and much of the factual content. But early layers and the final layer behave differently from the middle, and single-layer ablation understates importance wherever backup behavior exists — so report this as a distribution across layers, not one number.

    import torch

def kl_after_ablating(model, ids, layer, component):
    with torch.no_grad():
        base = torch.log_softmax(model(ids)[0, -1], dim=-1)

    target = getattr(model.blocks[layer], component)
    handle = target.register_forward_hook(lambda m, i, o: torch.zeros_like(o))
    with torch.no_grad():
        ablated = torch.log_softmax(model(ids)[0, -1], dim=-1)
    handle.remove()

    return float(torch.nn.functional.kl_div(ablated, base, log_target=True,
                                            reduction="sum"))

Sign in to track your progress through the lab.

← Self-Attention, Built From Scratch Back to track →