Multi-Head Attention

This closing note of the Seq2Seq & Attention category covers multi-head attention — running several independent attention computations in parallel, each potentially learning to focus on a different type of relationship — the final piece needed before the full Transformer architecture in the next category.

The Core Idea

A single attention computation, as covered so far, produces one specific weighted view of the sequence — but there could be several genuinely different, simultaneously useful types of relationships to capture (e.g. one relating to grammatical structure, another to semantic similarity, another to positional proximity). Multi-head attention runs \(h\) independent attention computations — "heads" — each with its own learned \(\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V\), then combines all their outputs together.

Formula

Each head has its own independently learned projection matrices, so each can specialize in detecting a different kind of relationship. The outputs of all heads are concatenated together, then passed through one final learned linear layer \(\mathbf{W}_O\) to combine them into a single output.

Keeping the Total Computation Comparable

In practice, each head's dimension is set to \(d_k = d_{\text{model}}/h\) — so \(h\) heads, each operating on a smaller \(d_k\)-dimensional space, together cost roughly the same total computation as one head operating on the full \(d_{\text{model}}\)-dimensional space. This means multi-head attention isn't simply "more expensive" than single-head attention with the same total dimensionality — it trades some per-head capacity for the ability to specialize into multiple distinct relationship types simultaneously.

Diagram

Code

import torch
import torch.nn as nn

multihead_attn = nn.MultiheadAttention(embed_dim=64, num_heads=8, batch_first=True)

x = torch.randn(1, 10, 64)   # 10 tokens, 64-dim embeddings
output, attn_weights = multihead_attn(query=x, key=x, value=x)   # self-attention, multi-head
print(output.shape)         # (1, 10, 64) -- same shape as input, ready to feed to the next layer
print(attn_weights.shape)   # (1, 10, 10) -- averaged across heads by default

# Each of the 8 heads operates on dimension 64/8 = 8
print(64 // 8)   # 8 -- this is d_k per head

Common Mistakes

• Assuming more heads is always strictly better — more heads means each individual head has less capacity (smaller \(d_k\)), so there's a genuine tradeoff between the number of distinct relationship types captured and how much each head can represent; the right number of heads is an empirically tuned architectural choice.
• Forgetting the final \(\mathbf{W}_O\) linear layer after concatenation — without it, the concatenated heads' outputs would just be passed through unchanged, missing the learned combination step that lets the model weight and mix the different heads' contributions.

Interview Relevance

Q: "Why does the Transformer use multi-head attention instead of a single attention computation with the full model dimension?" Multiple heads let the model attend to different types of relationships simultaneously — one head might specialize in short-range syntactic relationships, another in long-range semantic ones — each with its own independently learned Q/K/V projections. By setting each head's dimension to \(d_{\text{model}}/h\), the total computational cost stays comparable to a single full-dimension head, so this specialization is gained without a proportional increase in cost.

Key Takeaways — Seq2Seq & Attention

• Basic encoder-decoder Seq2Seq models compress an entire input sequence into one fixed-size context vector — a genuine bottleneck that degrades quality on long sequences.
• Teacher forcing speeds up Seq2Seq training but introduces exposure bias, the mismatch between training-time and inference-time conditioning.
• Attention resolves the bottleneck by letting the decoder access every encoder position directly, via a query/key/value framework borrowed conceptually from search/retrieval.
• Scaled dot-product attention (dividing by \(\sqrt{d_k}\)) keeps softmax's gradients healthy regardless of the query/key dimensionality.
• Self-attention relates positions within one sequence directly, without the distance penalty an RNN's sequential processing imposes; cross-attention lets one sequence pull information from a different one; multi-head attention runs several attention computations in parallel to capture multiple relationship types at once.

Next: Transformers assembles self-attention, multi-head attention, and several supporting components (positional encoding, residual connections, feed-forward blocks) into the complete Transformer architecture — the foundation of every modern large language model.

Practice Question

A Transformer layer has \(d_{\text{model}}=512\) and uses 8 attention heads. What is \(d_k\) for each individual head?

Related DL Notes

• Scaled Dot-Product Attention
• Self-Attention
• Transformer Architecture