Multi-Head Attention

Single attention is limited: one set of Q/K/V projections captures one type of relationship.

Multi-head: parallel attention heads, each learning different patterns.

H heads, each with its own W_q, W_k, W_v projecting to d_k = D/H.

Each head:
  Q_i = X @ W_q^i  → (T, d_k)
  K_i = X @ W_k^i
  V_i = X @ W_v^i
  head_i = attention(Q_i, K_i, V_i)  → (T, d_v)

Concatenate: concat(head_1, head_2, ..., head_H)  → (T, H*d_v) = (T, D)

Final projection: output = concat @ W_o

Implementation:

In practice, vectorized:

Why multiple heads?

• Different heads attend to different things.
• Some learn syntactic relations.
• Some learn long-range dependencies.
• Some learn specific patterns (subjects, verbs, etc.).

Number of heads:

• GPT-2 small: 12.
• GPT-3: 96.
• Llama 70B: 64.

Per-head dimension d_k = d_model / H typical: 64-128.

Multi-head is mostly an empirical win. Theoretical ceiling is the same as single-head with same total dim, but multi-head trains better.

Recent work (Multi-Query Attention, Grouped-Query Attention):

• Share K, V across heads (only Q is per-head).
• Fewer parameters.
• Faster inference (less KV cache).
• Used by Llama 2, GPT-3.5/4 (rumored).