LM3-T5 — Multi-Head Attention

Learning objectives

After this lesson, you should be able to:

• Explain why multiple heads help and how the model width is divided among them.
• Describe the steps: project, split into heads, attend per head, concatenate, project out.
• Compute the parameter count of a multi-head attention layer and show that it does not depend on the number of heads.
• Implement causal multi-head attention with the right reshapes in PyTorch.

The idea: several readers, each with a different question

Picture a sentence being read by a team. One reader tracks who did what (subject and verb), another tracks pronouns and what they refer to, a third watches punctuation and the end of clauses. Each reader forms a summary of the same sentence. The final representation combines all the summaries. A head is one such reader: its own queries, keys and values, its own attention weights, its own output.

Splitting rather than multiplying

If the model width is d = 768 and there are H = 12 heads, each head works on vectors of width 768 / 12 = 64. Twelve heads of width 64 together cost the same as one head of width 768, but they can attend to twelve different patterns. A single head can only produce one weight distribution per token and so would have to blend all needs into one; several heads can specialise. Trained models show heads that track previous tokens, heads that attend to the start of the text, and heads with no obvious pattern; the specialisation is learned, not designed.

Why a final projection?

After concatenation the head outputs sit side by side, each in its own block of d/H features. An output projection matrix W_O (shape d × d) mixes the blocks so that information found by one head can combine with that from another. Without it, the heads could never interact until the next layer.

Implementation trick

The heads are not computed in a Python loop. Instead, one big projection produces all the queries at once, and a reshape splits the last dimension into (H, d/H). A transpose moves the head dimension next to the batch dimension, so a single batched matrix product computes all heads in parallel.

Cost

The extra structure is nearly free in parameters, but each head has its own (n, n) score matrix, so the memory for scores grows with the number of heads: H · n² entries. That is why long-context work focuses on this term (Lesson LM3-T6).

The mathematics: heads, concatenation and parameters

Let the input be X of shape (n, d_in), the output width d_out, and H heads with d_h = d_out / H (we require that H divides d_out). Head h has its own projection matrices:

H

the number of heads.

d_h

the width of one head: d_out / H.

W_Q^h, W_K^h, W_V^h

projection matrices of head h, each of shape (d_in, d_h).

M

the causal mask (Lesson LM3-T4), the same for all heads.

head_h

the output of head h, shape (n, d_h).

The heads are concatenated along the feature dimension and mixed:

Parameters. Stacking the per-head matrices side by side gives the three big matrices of shape (d_in, d_out) used in the implementation. The count is

The number of heads H does not appear. With d_in = d_out = 768 and no biases on Q, K and V this is 3 × 589,824 + 589,824 + 768 = 2,360,064; GPT-2 small also uses biases on Q, K and V (3 × 768 more), giving 4 × 768 × 768 + 4 × 768 = 2,362,368 parameters per attention layer. Twelve such layers hold about 28.3 million parameters, roughly a quarter of the 124 million.

Score memory. The score tensor has shape (B, H, n, n). For H = 12 and n = 1,024 that is 12 · 1,024² = 12,582,912 entries per sequence, about 50 MB in float32 for each layer, which is why fused kernels that never store it matter.

A worked example: width 8, two heads

Take d_in = d_out = 8 and H = 2, so d_h = 4. Three tokens (n = 3) in a batch of 2 are an input of shape (2, 3, 8).

1. Project: W_q(x), W_k(x), W_v(x) each give shape (2, 3, 8).
2. Split: view(2, 3, 2, 4) separates the 8 features into 2 heads of 4: shape (B, n, H, d_h).
3. Move heads forward: transpose(1, 2) gives (2, 2, 3, 4) = (B, H, n, d_h), so the head index acts like a batch index.
4. Scores: q @ k.transpose(2, 3) has shape (2, 2, 3, 3): one 3 × 3 matrix per head per sequence. Scale, mask, softmax over the last axis.
5. Mix values: weights @ v has shape (2, 2, 3, 4).
6. Merge heads: transpose(1, 2) to (2, 3, 2, 4), then contiguous().view(2, 3, 8) concatenates the heads.
7. Output projection: out_proj maps (2, 3, 8) to (2, 3, 8).

Parameter count for this layer: 3 × 8 × 8 = 192 for Q, K and V, plus 8 × 8 + 8 = 72 for the output projection: 264. With one head of width 8 the count is identical, because the heads divide the width.

import torch
import torch.nn as nn

class CausalMHA(nn.Module):
    def __init__(self, d, nheads, context_length):
        super().__init__()
        assert d % nheads == 0
        self.h, self.dh = nheads, d // nheads
        self.q, self.k, self.v = (nn.Linear(d, d, bias=False) for _ in range(3))
        self.out = nn.Linear(d, d)
        self.register_buffer("mask", torch.triu(torch.ones(context_length, context_length), 1).bool())

    def forward(self, x):
        B, n, d = x.shape
        split = lambda t: t.view(B, n, self.h, self.dh).transpose(1, 2)       # (B, H, n, dh)
        q, k, v = split(self.q(x)), split(self.k(x)), split(self.v(x))
        scores = q @ k.transpose(2, 3) / self.dh ** 0.5                       # (B, H, n, n)
        scores = scores.masked_fill(self.mask[:n, :n], float("-inf"))
        w = torch.softmax(scores, dim=-1)
        ctx = (w @ v).transpose(1, 2).contiguous().view(B, n, d)             # merge the heads
        return self.out(ctx)

torch.manual_seed(0)
x = torch.randn(2, 3, 8)
for H in (1, 2, 4):
    m = CausalMHA(8, H, context_length=16)
    print(H, tuple(m(x).shape), sum(p.numel() for p in m.parameters()))   # same shape and 264 parameters every time

The output shape and the parameter count are the same for 1, 2 and 4 heads. What changes is the way the 8 features are partitioned for attention.

Engineering implications and limitations

Choose d_out divisible by H. Typical head widths are 64 or 128. If the width is not divisible, the reshape fails. Check this when building the config.

contiguous() before view. After transpose the tensor is no longer laid out contiguously, and view then raises an error. Either call .contiguous() first or use reshape, which copies if needed.

Mixing up the order of the reshape. Splitting with view(B, n, H, d_h) is correct; view(B, H, n, d_h) on the same tensor silently mixes tokens and features. The shapes agree, the results are wrong, and only a comparison with a reference (the notebook does this) or a leakage test catches it.

Heads are not independent experts. Studies of trained models show that some heads can be removed with little loss, and some carry specific functions. Do not assume one head per concept; do not use head patterns as an explanation without testing it.

Variants change the K/V heads. Multi-query and grouped-query attention share keys and values among several query heads. They reduce the size of the KV cache and speed up inference; Module 10 covers them. The code above is the standard multi-head form.

Bias terms. Whether Q, K and V projections have biases is a configuration choice (GPT-2 does; many modern models do not). Count parameters with the setting you actually use.

Dropout. Apply dropout to the attention weights during training (and set it to 0 for evaluation) only if your recipe does; large-scale pretraining often uses none.

Common misconceptions

Engineering takeaways

1. Multi-head attention splits the width d into H heads of width d/H; each head runs scaled dot-product attention with its own projections.
2. The head outputs are concatenated back to width d and mixed by an output projection.
3. Parameters: 3 · d_in · d_out + d_out² + d_out, independent of H; 2,362,368 per layer for GPT-2 small with Q, K and V biases.
4. Implementation: project once, view as (B, n, H, d_h), transpose to (B, H, n, d_h), attend, transpose back, contiguous().view, project.
5. Score memory grows with H · n², which motivates fused attention kernels.

References and further reading

See Chapter 3 of the book (extending single-head attention to multi-head attention) and Chapter 4 for the GPT-2 sized configuration.

• Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762 (Section 3.2.2, multi-head attention).
• Build a Large Language Model (From Scratch), Chapter 3: Coding Attention Mechanisms (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: a systematic look at the tensor shapes and matrix operations, with the costs and the common shape bugs.

Related content

Nothing below is required to finish this lesson.

• Before
  LM3-T4: Causal Masking

  Mask future positions so each token only sees the past, and verify no information leaks.

• Lab
  LM3-L3: Implement Causal Multi-Head Attention in PyTorch

  Implement causal self-attention and multi-head attention in PyTorch and check them against a reference.