Multi-head and residuals

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 head 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 residual connection and LayerNorm 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 multi-head attention mechanism.

// 'headOutputs' is an array of 4 head outputs, each [seq_len × d_head].
// 'W_O' is the output projection matrix, shape [d_model × d_model].
// The script defines matmul(A, B).
//
// 1. Concatenate every head's output along the feature axis.
//    Each token's row becomes [head0_out..., head1_out..., ...].
// 2. Project the concatenated matrix through W_O.
// 3. Return the [seq_len × d_model] result.
return [];

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 entropy: 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.

// 'headAttentions' is an array of 4 attention matrices, each [seq_len × seq_len].
// Compute the average Shannon entropy of each head's attention rows:
//   row_entropy = -Σ p_i * log(p_i)
//   head_entropy = mean(row_entropy over rows)
//
// Low entropy → the head focuses on a few tokens (concentrated pattern).
// High entropy → it spreads attention evenly (no clear preference).
// Return an array of 4 entropies, one per head.
return [];

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 residual connection 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.)

LayerNorm 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:

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.

// 'input' is the original X (the input to the attention sublayer).
// 'sublayerOutput' is what came out of the attention sublayer.
// 'eps' is a small constant (1e-5) for numerical stability.
//
// 1. Residual: add input + sublayerOutput element-wise.
// 2. LayerNorm: for each row, subtract its mean and divide by its
//    standard deviation (with the eps added inside the sqrt).
//
// Return the normalized [seq_len × d_model] matrix. Each row should
// have mean ≈ 0 and std ≈ 1.
return sublayerOutput;

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 residual connection. 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. LayerNorm 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.

Going further

• Vaswani et al., "Attention Is All You Need" (2017). The original transformer paper. The architecture is still the dominant one nine years later.
• Anthropic's transformer-circuits.pub — the mechanistic-interpretability lab's writings on what heads actually do.
• He et al., "Deep Residual Learning" (2015) — the paper that put residual connections on the map.

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