Multi-Head Attention

Multi-head attention runs Self-Attention multiple times in parallel, each head learning different relationships.

Key design:

• Total dimension $d_\text{model}$ split across $h$ heads: each head has dimension $d_k = d_\text{model}/h$
• Typical: $d_\text{model} = 512$, $h = 8$ → each head operates in $d_k = 64$ dimensions
• Final projection $W_O$ mixes the heads' outputs

Why multiple heads:

• Different heads can attend to different types of relationships:
  • Head 1: syntactic (subject-verb agreement)
  • Head 2: semantic (coreference)
  • Head 3: positional (nearby tokens)
• One head per relationship type is more expressive than a single large attention
• Empirically, some heads become specialized, others become redundant (prunable)

Implementation — reshape to expose the head dimension ($D \to N \times H$), transpose so heads and sequence swap, run attention, then reshape back:

batch, seq_len, _ = x.shape
qkv = self.qkv_proj(self.norm(x))                 # (batch, seq_len, 3 * d_model)
qkv = qkv.reshape(batch, seq_len, 3, self.num_heads, self.head_dim)
q, k, v = qkv.unbind(dim=2)                        # (batch, seq_len, heads, head_dim)
q, k, v = (t.transpose(1, 2) for t in (q, k, v))  # (batch, heads, seq_len, head_dim)
 
mask = torch.tril(torch.ones(seq_len, seq_len)).bool()
out = scaled_dot_product_attention(q, k, v, mask)
out = out.transpose(1, 2).reshape(batch, seq_len, -1)  # recover (batch, seq, d_model)
return x + self.out_proj(out)                          # residual

Computation cost: same as single-head attention with full $d_\text{model}$ (the split makes it equivalent), so multi-head is "free" in compute but adds expressiveness.

Grouped-Query Attention (GQA): modern LLMs share K/V across groups of heads to reduce KV cache memory at inference. Used in LLaMA 2+.

See also: Self-Attention, Residual Connections, Causal Masking