Skip to content
The loss curve

Chapter 9 · 14 min

Multi-head and residuals

From one head to many. Add residual connections and LayerNorm — the wiring that makes Transformers trainable at depth.

In chapter 8 you built one attention head. It computed one attention pattern — one matrix of who-listens-to-whom — and reshuffled the token representations accordingly. That's one routing pattern per layer.

That's not enough. Real text has many simultaneous relationships: subject/verb agreement, pronoun resolution, modifier attachment, semantic similarity, position bias. Any one of these can be the right thing to attend to depending on the token. A model with one per layer has to pick.

The fix is the obvious one: run several heads in parallel. Each head has its own W_Q, W_K, W_V, so each ends up with a different attention pattern. Concatenate their outputs, project the result, move on. That's multi-head attention.

This chapter also introduces residual connections and layer normalization — the two pieces of plumbing that make it possible to actually stack attention layers without the network collapsing during training. They look like afterthoughts; they're load-bearing. Locally, you will add those shape-preserving operations to the pieces you already wrote. The and are essential for deep networks.

Same toy sentence as chapter 8 (the cat sat on mat), but with d_model = 8, H = 4 heads, so each head's d_head = 2.

1. Combine multiple heads

You already have a working single attention head. The chapter has pre-run 4 different heads (different random seeds for W_Q, W_K, W_V) and hands you each head's output as a [seq_len × d_head] matrix. The cell combines them.

The standard recipe:

  1. Concatenate every head's output along the feature axis. Each token's row becomes [head₀, head₁, head₂, head₃], length H × d_head = d_model.
  2. Project the concatenated matrix through a learned W_O of shape [d_model × d_model].

The output projection lets the model decide how to mix the heads. Without it, the heads' outputs would just be glued together with no chance to interact. This is the mechanism.

Code · JavaScript

The result has the same shape as the input — [seq_len × d_model] — but each token's new representation now reflects four different attention patterns blended together.

2. Inspect what the heads actually learn

We've claimed the heads see different things. Let's verify. The four pre-computed attention matrices are shown below as heatmaps. Some are sharp (a few cells dominate each row); some are diffuse (mass spread evenly across the row).

A standard way to quantify "concentration" is : H = -Σ p log p. Low entropy means the head focuses on a few tokens. The maximum entropy for a row of 5 tokens is log(5) ≈ 1.61.

Compute the average entropy per head.

Code · JavaScript

Attention patterns across heads · log(n) ≈ 1.61 = maximum entropy for 5 tokens

Head 0

thecatsatonmatthecatsatonmatthe, the: 0.240the, cat: 0.172the, sat: 0.211the, on: 0.193the, mat: 0.183cat, the: 0.181cat, cat: 0.218cat, sat: 0.191cat, on: 0.206cat, mat: 0.203sat, the: 0.210sat, cat: 0.198sat, sat: 0.198sat, on: 0.206sat, mat: 0.187on, the: 0.193on, cat: 0.201on, sat: 0.201on, on: 0.195on, mat: 0.210mat, the: 0.193mat, cat: 0.213mat, sat: 0.192mat, on: 0.211mat, mat: 0.192

Head 1

thecatsatonmatthecatsatonmatthe, the: 0.196the, cat: 0.205the, sat: 0.198the, on: 0.193the, mat: 0.207cat, the: 0.197cat, cat: 0.194cat, sat: 0.200cat, on: 0.210cat, mat: 0.198sat, the: 0.197sat, cat: 0.205sat, sat: 0.199sat, on: 0.193sat, mat: 0.206on, the: 0.188on, cat: 0.196on, sat: 0.198on, on: 0.210on, mat: 0.208mat, the: 0.204mat, cat: 0.198mat, sat: 0.201mat, on: 0.202mat, mat: 0.195

Head 2

thecatsatonmatthecatsatonmatthe, the: 0.234the, cat: 0.180the, sat: 0.205the, on: 0.201the, mat: 0.181cat, the: 0.169cat, cat: 0.223cat, sat: 0.192cat, on: 0.197cat, mat: 0.220sat, the: 0.206sat, cat: 0.195sat, sat: 0.202sat, on: 0.200sat, mat: 0.196on, the: 0.217on, cat: 0.192on, sat: 0.201on, on: 0.201on, mat: 0.190mat, the: 0.156mat, cat: 0.229mat, sat: 0.191mat, on: 0.195mat, mat: 0.229

Head 3

thecatsatonmatthecatsatonmatthe, the: 0.215the, cat: 0.194the, sat: 0.200the, on: 0.189the, mat: 0.201cat, the: 0.189cat, cat: 0.205cat, sat: 0.200cat, on: 0.209cat, mat: 0.197sat, the: 0.190sat, cat: 0.203sat, sat: 0.196sat, on: 0.203sat, mat: 0.209on, the: 0.199on, cat: 0.201on, sat: 0.202on, on: 0.203on, mat: 0.196mat, the: 0.185mat, cat: 0.204mat, sat: 0.195mat, on: 0.206mat, mat: 0.210

You should see meaningful differences between heads. Some end up nearly uniform (high entropy — the head doesn't strongly prefer anything). Others have sharp peaks (low entropy — the head has decided to focus). In a real trained transformer, you'd find specialized heads: copy heads, induction heads, name-following heads, syntax-tracking heads. We named patterns we found, not patterns we asked for.

3. Residual + LayerNorm

If we wired the attention sublayer straight into the next attention sublayer, training would collapse. Two reasons, both gnarly:

  • Gradient vanishing. Each sublayer compresses the gradient signal a bit. Stack 12 of them and the gradient at the bottom is microscopic.
  • Representation drift. Each sublayer transforms the activations into a different geometry. Stack many and the activation magnitudes blow up or shrink to zero.

The fixes the first: instead of output = sublayer(input), we use output = input + sublayer(input). The gradient now has a clean path back through the addition, regardless of what the sublayer does. (This is the trick that made ResNets practical in 2015 and is now in every deep model.)

fixes the second: after the residual addition, normalize every token's row to have mean 0 and standard deviation 1. The downstream layer sees activations of a known scale regardless of what came before.

Combined, the per-sublayer recipe is:

output=LayerNorm(input+sublayer(input))\text{output} = \text{LayerNorm}(\text{input} + \text{sublayer}(\text{input}))

Run it. The chapter feeds the cell the original input (X) and sublayerOutput (the multi-head output from cell 1), plus a small eps for numerical stability inside the standard-deviation division.

Code · JavaScript

Look at the row statistics. Each row's mean should be 0 (to floating-point precision) and each row's std should be very close to 1. That's the invariant LayerNorm gives the next sublayer: every token, every layer, comes in with the same scale.

Why this chapter matters

We've now got every piece of a transformer block:

  • Attention (chapter 8): pull information from other positions.
  • Multi-head (this chapter): run several attention routes in parallel.
  • Residual + LayerNorm (this chapter): the connectivity and normalization that let us stack many blocks.

The next chapter assembles them into a complete transformer block and stacks several of them into the actual architecture.

4. Add residual and normalization helpers

Append these helpers to llm/nn.py:

import math
 
 
def add(a: Matrix, b: Matrix) -> Matrix:
    return [
        [x + y for x, y in zip(row_a, row_b)]
        for row_a, row_b in zip(a, b)
    ]
 
 
def layer_norm(x: Matrix, eps: float = 1e-5) -> Matrix:
    out: Matrix = []
    for row in x:
        mean = sum(row) / len(row)
        var = sum((value - mean) ** 2 for value in row) / len(row)
        denom = math.sqrt(var + eps)
        out.append([(value - mean) / denom for value in row])
    return out

These helpers are small because their job is structural:

  • add is the . It keeps the old representation and adds the sublayer's proposed change.
  • layer_norm works row by row, so every token is normalized independently.
  • mean recenters a token's features around zero.
  • var measures how spread out that token's features are.
  • Dividing by sqrt(var + eps) gives the next layer a predictable scale. eps prevents division by zero. stabilizes training by maintaining consistent activation scales.

You now have the invariant a transformer depends on: every sublayer accepts a matrix and returns a matrix of the same shape, so the residual stream can keep flowing.

Recap

  • Multi-head attention runs H attention computations in parallel with separate Q/K/V projections per head, then concatenates the outputs and projects with W_O. - Different heads learn different patterns by virtue of starting with different random weights and being trained on the same loss. Interpretability research catalogs the recurring patterns. - Residual connection = output = input + sublayer(input). Lets gradients flow cleanly through deep stacks. - LayerNorm normalizes each token's row to mean 0 std 1. Stabilizes activation scales across layers. - Your local project now has residual and LayerNorm helpers, the glue that lets attention stack. - The transformer block is just these three pieces glued together, twice (once for attention, once for feed-forward). Next chapter assembles the whole thing. The and are essential for deep networks.

Going further

Next up: the full transformer block — combine everything we've built into the actual unit you stack.