Multi Head Attention — Module 10.1

How This Lesson Fits the Module & Engineering Practice

Single-head self-attention or cross-attention forces all relationships through one subspace. Multi-head attention (MHA) runs \(h\) attention operations in parallel so different heads can specialize (syntax, long-range links, local patterns), then concatenates and projects with \(W_O\).

This is the attention block you meet in every modern transformer API (nn.MultiheadAttention). Mastering heads, concat, and the output projection is required before Module 10.2’s full stack.

Learning Objectives

By the end of this lesson, students should be able to:

Explain why multiple heads improve representational capacity.
Write the multi-head formula including concat and \(W_O\).
Relate \(d_model\), \(h\), and \(d_k = d_model / h\).
Implement or call multi-head attention in PyTorch.
Describe how MHA wraps both self- and cross-attention.
Spot mistakes such as forgetting \(W_O\) or mismatched head dimensions.

Definition

Multi-head attention projects inputs into \(h\) sets of Q/K/V, runs scaled dot-product attention per head, concatenates the head outputs, and applies a learned linear map \(W_O\) back to model dimension.

The Formula

For head \(i\):

\[\mathrm{head}_i = \mathrm{Attention}(Q W_Q^{(i)},\, K W_K^{(i)},\, V W_V^{(i)})\]

Then:

\[\mathrm{MultiHead}(Q,K,V) = \mathrm{Concat}(\mathrm{head}_1,\ldots,\mathrm{head}_h)\, W_O\]

If each head has width \(d_k\) (and often \(d_v = d_k\)), concat has width \(h \cdot d_k\). Usually \(d_model = h \cdot d_k\), so \(W_O \in \mathbb{R}^{(h d_k) \times d_model}\) restores the channel size expected by residual paths.

1. Project

Map to \(h\) Q/K/V slices (or one big linear then split).

2. Attend

Independent scaled attention per head.

3. Concat

Join head outputs along the feature axis.

4. \(W_O\)

Mix head features into a single \(d_model\) vector.

Why Multiple Heads?

Single head

One averaging pattern per layer.
Cheaper, but less diverse.
Can miss complementary relations.

Multi-head

Parallel subspaces / viewpoints.
Same asymptotic form, richer mix.
Standard in transformers.

Symbol	Meaning	Typical relation
\(d_model\)	Model / residual width	e.g. 512, 768
\(h\)	Number of heads	e.g. 8, 12
\(d_k\)	Per-head key/query dim	\(d_model / h\)
\(W_O\)	Output projection	\((h d_k) \to d_model\)

Self vs Cross Under Multi-Head

MHA is a wrapper, not a new attention type. Pass the same tensor for Q/K/V sources to get multi-head self-attention; pass target as Q-source and encoder memory as K/V-source for multi-head cross-attention. See Self Attention and Cross Attention.

PyTorch: Heads, Concat, \(W_O\)

import math
import torch
from torch import nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int):
        super().__init__()
        assert d_model % num_heads == 0
        self.h = num_heads
        self.d_k = d_model // num_heads
        self.W_qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        B, T, _ = x.shape
        qkv = self.W_qkv(x).reshape(B, T, 3, self.h, self.d_k)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, B, h, T, d_k)
        Q, K, V = qkv[0], qkv[1], qkv[2]
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)
        weights = torch.softmax(scores, dim=-1)
        heads = weights @ V                      # (B, h, T, d_k)
        concat = heads.transpose(1, 2).reshape(B, T, self.h * self.d_k)
        return self.W_o(concat)                  # (B, T, d_model)

x = torch.randn(2, 10, 64)
y = MultiHeadAttention(64, num_heads=8)(x)
print(y.shape)  # torch.Size([2, 10, 64])

# Library equivalent for self-attention:
mha = nn.MultiheadAttention(embed_dim=64, num_heads=8, batch_first=True)
y2, attn = mha(x, x, x, need_weights=True)

Strengths

Diverse attention patterns in one layer.
\(W_O\) learns how to combine heads.
Drop-in for self or cross wiring.

Tradeoffs

More parameters and matmuls than one head.
Still quadratic in sequence length per head.
Heads are not guaranteed to be “interpretable.”

Common Mistake

Concatenating heads and returning that tensor without \(W_O\). The output projection is part of the definition: it mixes head channels and matches residual width for residual connections.

Misconception

“Eight heads means eight times the sequence length cost of one full attention.” Each head usually uses a thinner \(d_k = d_model/h\), so total FLOPs are comparable to one full-width attention—plus the cost of \(W_O\)—not naïvely \(h\times\) a full \(d_model\) head.

Sibling links: Attention, Scaled Dot-Product Attention, Positional Encoding, Layer Normalization, Feed-Forward Network.

Knowledge Check

Short Answer: What does multi-head attention do after computing each head? Answer: Concatenates head outputs and multiplies by \(W_O\).
True/False: \(W_O\) is optional cosmetic fluff. Answer: False; it is a learned mixing projection back to \(d_model\).
Multiple Choice: If \(d_model=64\) and \(h=8\), typical \(d_k\) is: (a) 8, (b) 64, (c) 512. Answer: (a).
Short Answer: Why use multiple heads? Answer: To attend in multiple subspaces / capture diverse relationships in parallel.
True/False: Multi-head attention can implement both self- and cross-attention. Answer: True.
Multiple Choice: Concat of heads has width: (a) \(d_k\), (b) \(h\cdot d_k\), (c) \(h\) only. Answer: (b).
Short Answer: Name the PyTorch module for MHA. Answer: nn.MultiheadAttention.
True/False: \(d_model\) must be divisible by the number of heads in the usual split design. Answer: True.
Multiple Choice: Which projection restores residual width after concat? (a) \(W_O\), (b) softmax, (c) LayerNorm alone. Answer: (a).
Short Answer: What lecture follows MHA in this track? Answer: Positional Encoding.

Key Takeaways

MHA = \(h\) attentions in parallel → concat → \(W_O\).
Usually \(d_k = d_model / h\); total width stays \(d_model\).
Same block supports self- and cross-attention wiring.
Do not drop the output projection when implementing from scratch.
Next: Positional Encoding.

Trainer’s Guide

Hands-on idea: Visualize attention maps from two different heads on the same sentence and discuss specialization (even if imperfect).

Discussion prompt: If you double heads but keep \(d_model\) fixed, what happens to per-head capacity?

Recap: Multi-head attention runs parallel attentions, concatenates them, and projects with \(W_O\). Continue with Positional Encoding.