07 Module 2 · Building the Model
Chapter 7 of 17 Run it in your browser

Chapter 7: Multi-Head Attention

Figure 7.1: We run multiple attention heads at the same time to gather richer context.

In Chapter 6, we built one attention head, which acts as a single "reader" scanning the text. But one reader is rarely enough to catch everything. In this chapter you will:

• Run multiple attention heads in parallel.
• Understand how different heads learn to track different patterns.
• Concatenate their outputs into a single rich representation.

Words to Know

• Multi-Head Attention: Running several attention heads at the same time, so the same text is read in several different ways.
• Concatenation: Joining several lists of numbers end to end to make one longer list.
• Projection: A final step that blends the heads' joined outputs into one list of numbers for each token.

Theory

One Reader Isn't Enough

Figure 7.2: Every head sees the same input through its own filters, and their answers are joined rather than merged.

When you read a sentence, you track multiple things at once:

• Who is the subject?
• What action is happening?
• What is the tone, serious or sarcastic?

A single attention head can only "look" for one kind of pattern at a time. Multi-head attention fixes this by running several heads in parallel (Figure 7.2). Think of it as a team of readers, each highlighting different details, who then combine all their notes.

How It Works

Multi-head attention repeats the mechanism from Chapter 6:

1. Run 4 independent attention heads (each one a SingleHeadAttention from Chapter 6) on the same input.
2. Each head produces an output of 32 numbers (head_size).
3. Glue (concatenate) all 4 outputs together: 4 heads × 32 numbers = 128 numbers.
4. Apply a final projection, one more step that blends them.

Because each head starts with different random weights (its own Q, K, and V filters), no two heads ever see the text the same way, and training pushes them further apart.

Why Four Heads and Not One Big One

Four readers sound expensive. They are not, and this is the part that surprises people.

A head's filters are sized by its head_size, not by the full 128 channels. Split 128 channels across four heads and each head gets filters that are a quarter as wide. Multiply it out and the totals match exactly:

head_size = C // n_heads

# Every head holds three filters (query, key, value), each C by head_size
four_heads = n_heads * 3 * C * head_size
one_head = 3 * C * C

Terminal

$ python src/examples/ch07_head_budget.py
4 heads of 32: 49,152 numbers
1 head of 128  : 49,152 numbers
Same budget: True

Four points of view cost the same as one. That is why every model in this family uses many heads. The split is free, so there is no reason to take one perspective when you can have four.

Do the Heads Really Differ?

That is the claim. Here is the check, on the model you will train in Chapter 13.

The script below loads the trained weights, feeds in a real line of Shakespeare, and prints how much attention the last character pays to each earlier character, one row per head.

for h, head in enumerate(model.blocks[0].attn.heads):
    with torch.no_grad():
        scores = head.query(x) @ head.key(x).transpose(-2, -1) * head.head_size ** -0.5
        weights = F.softmax(scores, dim=-1)[0, -1]     # the last character's row

Terminal

$ python src/examples/ch07_heads_differ.py
Prompt: "JULIET: O Romeo"

Attention paid by the last character, one row per head:
              J    U    L    I    E    T    :    _    O    _    R    o    m    e    o
  head 0:  0.02 0.04 0.01 0.06 0.03 0.04 0.04 0.02 0.06 0.04 0.05 0.04 0.11 0.38 0.07
  head 1:  0.02 0.02 0.03 0.03 0.04 0.08 0.12 0.06 0.05 0.07 0.02 0.10 0.16 0.14 0.06
  head 2:  0.04 0.07 0.04 0.06 0.05 0.11 0.04 0.06 0.09 0.07 0.07 0.06 0.09 0.09 0.06
  head 3:  0.03 0.02 0.01 0.05 0.05 0.03 0.03 0.05 0.04 0.04 0.05 0.08 0.12 0.25 0.14

Sharpest focus per head: 0->'e' (0.38)  1->'m' (0.16)  2->'T' (0.11)  3->'e' (0.25)

Read the rows, not the labels. Head 0 commits: it puts 0.38 of its attention on the single character just before the end and largely ignores the rest. Head 3 does something similar but softer, splitting between e and the final o. Head 1 spreads itself over the colon, the m and the e, holding several places at once. Head 2 is nearly flat. Its largest weight is 0.11 and its smallest is 0.04, which is close to paying equal attention to everything.

So the heads do differ, and not in the tidy way the textbook story suggests. One is sharp, one is soft, one is diffuse, and one is barely committing at all. That last one is worth sitting with. In a trained model, some heads do very little. Nobody assigned these roles, and nobody can promise that head 1 is the "grammar head". They are four different filters that started from four different random draws and were shaped by the same pressure to predict the next character.

Watch Out

It is tempting to read a story into each head: this one tracks subjects, that one tracks punctuation. Sometimes a head really is that clean, and researchers have found interpretable ones in large models. Often it is not, as head 2 shows. Look at the numbers before you tell the story.

The Output Projection

After concatenating the 4 heads, we pass the 128 numbers through one final linear layer (the proj layer).

We add this layer because the heads might have found redundant or conflicting information. The projection layer learns to blend the insights: "if head 1 and head 3 agree on this pattern, emphasize it, and if head 2 is unsure, ignore it." It mixes the 4 separate perspectives into a single unified context vector of 128 numbers. With the projection in place, the "Attention" stage of our map is complete (Figure 7.1).

There is a simpler thing we could have done here, and it is worth seeing why we did not. We could average the four heads instead of gluing them end to end. Averaging would give us 32 numbers rather than 128, which sounds tidy, but it throws away exactly what we paid for. Head 0's sharp focus on one character and head 2's flat spread would cancel each other into a lukewarm middle. Concatenation keeps every head's answer intact and lets the projection layer decide what each one is worth. Averaging decides in advance that they are all worth the same.

In Business

Figure 7.3: Multi-head attention gathers different perspectives.

When building an assistant to draft emails in your company's house style, you look for more than one pattern. You track tone, structure, and vocabulary (Figure 7.3). In the same way, multi-head attention asks 4 different "departments" to evaluate the text, then compiles a final executive summary.

Code

Figure 7.4: The heads run in parallel, get concatenated, and projected.

We use src/ch06_multihead_attention.py to build the MultiHeadAttention class (Figure 7.4).

        # Run each head in parallel
        head_outputs = [h(x) for h in self.heads]

        # Concatenate outputs along the last dimension
        out = torch.cat(head_outputs, dim=-1)

        # Apply projection and dropout
        out = self.dropout(self.proj(out))

Running the script compares one head with four:

Terminal

$ python src/ch06_multihead_attention.py
--- Comparing one head vs multi-head ---
Single head output  shape: torch.Size([2, 10, 32])  (head_size=32)
Multi-head output   shape: torch.Size([2, 10, 128])  (C=128)

Multi-head output has 4x more channels - it sees 4 perspectives at once.

Multi-head attention done! Ready for Chapter 8.

What just happened:

• Line 2 ran 4 independent attention heads at the same time.
• Line 2's result is that each head produced an output with 32 channels.
• Line 5 concatenated them into 128 channels (the original embedding dimension).
• Line 8 applied the projection, which blends the 128 channels and keeps the shape of the input. The same line also applied dropout, which Chapter 14 explains.

Shape Check: Table 7.1 lists the shapes as the data passes through the heads.

Table 7.1: Tensor shapes during the multi-head attention step.

Variable	Shape	Meaning
head_outputs	4 × [B, T, 32]	A list of outputs from the 4 heads.
out (after cat)	[B, T, 128]	The 4 outputs joined end-to-end.
out (final)	[B, T, 128]	The blended representation, ready for the next step.

Try It

Try It

Change the number of heads in src/utils/config.py by setting n_heads = 8 instead of 4. Run the script again and notice how the head_size automatically drops to 16, so the final concatenated size is still 128 (8 × 16 = 128). The model can have more perspectives, but each one has less detail.

Watch Out

For multi-head attention to work cleanly, your embedding dimension (n_embd) must be divisible by your number of heads (n_heads). If you try n_embd = 128 and n_heads = 5, the program will crash because it cannot divide the channels equally.

Key Takeaways

• Multi-head attention runs several single-head attention modules in parallel.
• Splitting the channels across heads costs nothing: four heads of 32 hold the same 49,152 numbers as one head of 128.
• The heads do end up different, but not in a tidy way. In our trained model one head is sharp, one is soft and one is close to flat.
• The outputs are concatenated, not averaged, so that no head's answer is diluted by the others before the projection can weigh it.
• A final linear projection blends the separate perspectives.
• The output shape [B, T, C] is exactly the same as the input shape, making it easy to stack layers.

Check Your Understanding

1. Why is one attention head not enough to understand complex text?
2. If n_embd = 256 and n_heads = 8, what is the head_size?
3. Four heads of 32 hold the same number of weights as one head of 128. Why does splitting cost nothing?
4. Why do we concatenate the heads rather than average them?
5. In the run above, head 2's attention was almost flat. What does that tell you about the claim that each head learns its own linguistic role?
6. What is the purpose of the final projection layer?

Assignment

Assignment 7: Same heads, new word

About 15 minutes. You change one word and write no new code. Work in this chapter's Colab notebook or on your own laptop.

1. Run it. Run src/examples/ch07_heads_differ.py as it is. Copy the last line it prints, the one that starts Sharpest focus per head.
2. Change one word. Near the top of the file, find PROMPT = "JULIET: O Romeo". The prompt is the line of text the model reads. Change Romeo to love, keep the quotation marks, and run the script again.
3. Compare. Which character does each head focus on most now? Which heads pick the character just before the last one in both runs? Which two heads have the smallest numbers in both runs?
4. Explain it. In two or three sentences a manager could follow, say how the four heads differ, and why it is worth testing a second prompt before you say what job a head does.

Hand in: both Sharpest focus per head lines, your answers to step 3, and your sentences.

Further Reading

The architecture this book builds. Reading a sequence one step at a time is slow, because step 500 cannot start until step 499 has finished, and distant words stay hard to connect. This paper removed the step-by-step reading entirely and kept only attention, plus a note of each token's position. Every token can then be processed at once, which is what made training on very large amounts of text practical. The model you build in Chapters 6 to 10 is this design, made small.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., & Polosukhin, I. (2017). Attention is all you need (arXiv:1706.03762). arXiv. https://doi.org/10.48550/arXiv.1706.03762

Next up: Chapter 8: Feed-Forward and Norms