Multi-Head Attention & Chapter Summary · Build an LLM From Scratch

One rung remains. Causal attention computes one set of attention patterns per sequence — but a sentence carries many kinds of relationships at once (who did what, which words modify which, what refers back to what). Multi-head attention runs several causal-attention “heads” in parallel, each with its own $W_q$ , $W_k$ , $W_v$ , so each can learn to track a different kind of relationship. This is the final module — the one dropped into the GPT in chapter 4.

3.6.1 Stacking multiple single-head attention layers

The conceptually simplest version: instantiate num_heads complete CausalAttention modules and concatenate their outputs along the feature dimension (the book’s Figure on p. 45):

Two heads, two full sets of projections, two context-vector matrices — concatenated side by side. Choosing d_out = 2 per head gives a final embedding dimension of d_out × num_heads = 4.

</> Listing 3.4 — A wrapper class to implement multi-head attention

class MultiHeadAttentionWrapper(nn.Module):
  def __init__(self, d_in, d_out, context_length,
               dropout, num_heads, qkv_bias=False):
      super().__init__()
      self.heads = nn.ModuleList(
          [CausalAttention(d_in, d_out, context_length, dropout, qkv_bias)
           for _ in range(num_heads)]  #A
      )

  def forward(self, x):
      return torch.cat([head(x) for head in self.heads], dim=-1)  #B

#A num_heads complete, independent CausalAttention modules — each with its own W_q, W_k, W_v.
#B Run each head on the same input and glue the outputs along the last (feature) dimension.

</> Using the wrapper on the 2-sequence batch

torch.manual_seed(123)
context_length = batch.shape[1]  # number of tokens = 6
d_in, d_out = 3, 2
mha = MultiHeadAttentionWrapper(
  d_in, d_out, context_length, 0.0, num_heads=2
)
context_vecs = mha(batch)
print("context_vecs.shape:", context_vecs.shape)
# context_vecs.shape: torch.Size([2, 6, 4])  #A

#A First dim 2: two input texts (identical, so both batch elements print the same). Second dim 6: tokens. Third dim 4: num_heads × d_out = 2 × 2 — the concatenated per-head context vectors.

The first batch element of the printed output, with the two heads’ columns side by side — head 1 wrote the left two columns, head 2 the right two:

context_vecs[0] (6×4) — head 1 | head 2

Your	-0.45	0.22	0.48	0.11
journey	-0.59	0.01	0.59	0.33
starts	-0.63	-0.06	0.62	0.39
with	-0.57	-0.08	0.55	0.36
one	-0.55	-0.10	0.53	0.34
step	-0.53	-0.11	0.51	0.35
	h1·d0	h1·d1	h2·d0	h2·d1

Each row is one token's 4-D context vector: columns 0–1 come from head 1's value space, columns 2–3 from head 2's. The heads never mixed — yet. That's what the output projection in §3.6.2 adds.

3.6.2 Multi-head attention with weight splits

The wrapper works, but looping over separate modules is wasteful. The production version — the class chapter 4’s GPT will import — creates one set of full-width projections and splits them into heads by reshaping: each head works on a slice of size head dimension = d_out / num_heads, and every head’s attention runs inside a single batched matrix multiplication. A final output projection then mixes the concatenated head outputs:

</> The efficient MultiHeadAttention with weight splits

class MultiHeadAttention(nn.Module):
  def __init__(self, d_in, d_out, context_length,
               dropout, num_heads, qkv_bias=False):
      super().__init__()
      assert (d_out % num_heads == 0), \
          "d_out must be divisible by num_heads"  #A

      self.d_out = d_out
      self.num_heads = num_heads
      self.head_dim = d_out // num_heads

      self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)  #B
      self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
      self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
      self.out_proj = nn.Linear(d_out, d_out)  #C
      self.dropout = nn.Dropout(dropout)
      self.register_buffer(
          "mask",
          torch.triu(torch.ones(context_length, context_length),
                     diagonal=1)
      )

  def forward(self, x):
      b, num_tokens, d_in = x.shape

      keys    = self.W_key(x)      # (b, num_tokens, d_out)
      queries = self.W_query(x)
      values  = self.W_value(x)

      # Split the last dim: (b, T, d_out) -> (b, T, num_heads, head_dim)
      keys    = keys.view(b, num_tokens, self.num_heads, self.head_dim)  #D
      values  = values.view(b, num_tokens, self.num_heads, self.head_dim)
      queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

      # (b, T, num_heads, head_dim) -> (b, num_heads, T, head_dim)
      keys    = keys.transpose(1, 2)  #E
      queries = queries.transpose(1, 2)
      values  = values.transpose(1, 2)

      attn_scores = queries @ keys.transpose(2, 3)  #F
      mask_bool = self.mask.bool()[:num_tokens, :num_tokens]
      attn_scores.masked_fill_(mask_bool, -torch.inf)

      attn_weights = torch.softmax(
          attn_scores / keys.shape[-1]**0.5, dim=-1)
      attn_weights = self.dropout(attn_weights)

      # (b, num_heads, T, head_dim) -> (b, T, num_heads, head_dim)
      context_vec = (attn_weights @ values).transpose(1, 2)

      # Combine heads: d_out = num_heads × head_dim
      context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)  #G
      context_vec = self.out_proj(context_vec)  #H
      return context_vec

#A The split only works if d_out divides evenly into num_heads slices of head_dim each.
#B ONE full-width projection per role — not one per head. The heads are implicit slices of these matrices.
#C The output projection: mixes the concatenated head outputs. Technically optional, but standard in LLM architectures.
#D view() reinterprets the d_out axis as (num_heads × head_dim) without copying data.
#E transpose(1, 2) moves heads next to the batch dim: PyTorch matmul now treats (b, num_heads) as batch axes and runs EVERY head in parallel.
#F keys.transpose(2, 3) flips (T, head_dim) so the per-head matmul yields (b, num_heads, T, T) attention scores — one T×T matrix per head per batch element.
#G transpose back, then contiguous().view() glues the head outputs side by side — the tensor equivalent of the wrapper’s torch.cat.
#H Finally the out_proj lets the heads exchange information.

The whole trick is a shape choreography. Follow one tensor through the forward pass (the book’s running numbers: b=2, T=6, d_out=2, num_heads=2, so head_dim=1):

Project once, full width

W_query, W_key, W_value each map (2, 6, 3) → (2, 6, 2). One matmul per role — the "heads" don't exist yet; they are hiding inside the d_out=2 axis.

1 / 8

The weight-split dataflow, one reshape at a time — with the running example's concrete shapes.

</> Using MultiHeadAttention on the batch

torch.manual_seed(123)
batch_size, context_length, d_in = batch.shape
d_out = 2
mha = MultiHeadAttention(d_in, d_out, context_length, 0.0, num_heads=2)
context_vecs = mha(batch)
print("context_vecs.shape:", context_vecs.shape)
# context_vecs.shape: torch.Size([2, 6, 2])  #A
print(context_vecs)
# tensor([[[0.3190, 0.4858],
#          [0.2943, 0.3897],
#          [0.2856, 0.3593],
#          [0.2693, 0.3873],
#          [0.2639, 0.3928],
#          [0.2575, 0.4028]],
#         [[0.3190, 0.4858],
#          ... (identical — same input twice) ...
#          [0.2575, 0.4028]]], grad_fn=<ViewBackward0>)

#A Here d_out=2 is the TOTAL width (head_dim=1 each) — unlike the wrapper, where each head added its own d_out. Same interface chapter 4 will use with d_out=768, num_heads=12.

Wrapper vs. weight splits, side by side

Aspect⇅	MultiHeadAttentionWrapper (3.6.1)⇅	MultiHeadAttention (3.6.2)⇅
Projection matrices	num_heads separate W_q, W_k, W_v sets	ONE full-width W_q, W_k, W_v	+
Head computation	Python loop over modules	one batched 4-D matmul	+
Output dimension	num_heads × d_out (concat grows it)	d_out total (head_dim = d_out / num_heads)	+
Head mixing	none — pure concatenation	out_proj linear layer	+
Used in	teaching / exercises	the GPT model of ch4 (d_out=768, num_heads=12)	+

Two ways to get multi-head attention — click a row for the rationale.

3.7 Summary — the four rungs, climbed

The book’s Figure 3.23 mental model, completed: every module built, each extending the previous one.

What this chapter established, in the book’s own summary points — each now backed by working code:

Attention transforms inputs into enriched context-vector representations that incorporate information about all (visible) inputs, computed as a weighted sum over the inputs.
In the simplified mechanism, attention weights come from dot products — element-wise multiply, then sum — normalized by softmax; matrix multiplications aren’t strictly required but make the computation compact and efficient, replacing nested loops.
The LLM version, scaled dot-product attention, adds trainable matrices computing intermediate transformations of the inputs: queries, keys, and values, with $\sqrt{d_k}$ scaling for healthy gradients.
Left-to-right LLMs add a causal mask so attention never sees future tokens, plus a dropout mask to reduce overfitting.
Stacking causal-attention modules gives multi-head attention; the efficient implementation replaces the stack with batched matrix multiplications via view/transpose weight splits.

Where this goes next

MultiHeadAttention(d_in, d_out, context_length, dropout, num_heads, qkv_bias) is a finished, importable component. In chapter 4 it slots into the transformer block between layer normalization and the feed-forward network — instantiated at GPT-2 scale: d_out=768, num_heads=12, context_length=1024.

📝 Check yourself: multi-head attention

0 / 5

1.Predict: MultiHeadAttentionWrapper(d_in=3, d_out=2, num_heads=2) processes a batch of shape (2, 6, 3). What is the output shape?
2.Why run multiple attention heads instead of one bigger head?
3.In the efficient MultiHeadAttention class, why must d_out be divisible by num_heads?
4.What is the key implementation difference between MultiHeadAttentionWrapper and the efficient MultiHeadAttention?
5.What does the final out_proj linear layer in MultiHeadAttention do?