Multi-Head Attention

Summary: Multi-head attention runs h parallel attention computations on linearly projected Q, K, V, then concatenates and projects the outputs. This allows the model to jointly attend to information from different representation subspaces at different positions.

Motivation

In plain English: Single-head attention averages over all positions, which can wash out distinct relational patterns. Multiple heads let the model simultaneously capture different types of relationships (e.g., syntactic, semantic, positional) in parallel subspaces.

With single-head attention of dimension d_{model}:

head = Attention(Q, K, V) ∈ R^{n × d_{model}}

Multi-head splits this into h heads, each with dimension d_k = d_v = d_{model}/h:

MultiHead(Q, K, V) = Concat(head_1, ..., head_h)W^O

where:

head_i = Attention(QW^Q_i, KW^K_i, VW^V_i)

Key property: Total computation ≈ single-head attention with full dimension.

Transformer Configuration (Base Model)

Parameter Value
d_{model} 512
h (heads) 8
d_k = d_v 64
d_{ff} 2048
N (layers) 6

Projection Matrices

  • W^Q_i ∈ R^{512 × 64}
  • W^K_i ∈ R^{512 × 64}
  • W^V_i ∈ R^{512 × 64}
  • W^O ∈ R^{512 × 512} (since h · d_v = 8 · 64 = 512)

Total attention parameters per layer: 4 × 512 × 512 ≈ 1.05M

Why Multiple Heads?

1. Representation Subspaces

Each head learns different projection matrices, effectively learning different "views" of the relationships:

  • Syntactic heads: Subject-verb agreement, dependency parsing
  • Semantic heads: Coreference, entity tracking
  • Positional heads: Relative distance, ordering

2. Ensemble Effect

Analogous to ensemble learning — multiple "experts" voting on attention distribution.

3. Gradient Flow

Each head receives independent gradients, potentially improving optimization.

Ablation Study Results (Table 3, Row A)

Heads (h) d_k d_v BLEU (EN-DE dev)
1 512 512 24.9
4 128 128 25.5
8 64 64 25.8
16 32 32 25.8
32 16 16 25.4

Findings:

  • Single head: 0.9 BLEU worse than best
  • Too many heads (32): quality degrades (insufficient capacity per head)
  • Optimal: 8 heads with d_k=64 (base), 16 heads with d_k=64 (big)

Computational Considerations

Parallel Implementation

In practice, heads are computed in a single batched matrix multiplication:

# Q, K, V: (batch, seq_len, d_model)
# Reshape: (batch, seq_len, h, d_k) -> (batch, h, seq_len, d_k)
Q = Q.view(batch, seq_len, h, d_k).transpose(1, 2)
K = K.view(batch, seq_len, h, d_k).transpose(1, 2)
V = V.view(batch, seq_len, h, d_v).transpose(1, 2)

# Batched attention: (batch, h, seq_len, seq_len)
attn = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
attn = F.softmax(attn, dim=-1)
out = torch.matmul(attn, V)  # (batch, h, seq_len, d_v)

# Concat: (batch, seq_len, h * d_v)
out = out.transpose(1, 2).contiguous().view(batch, seq_len, h * d_v)
out = self.W_O(out)

Memory Optimization

  • FlashAttention flash-attention: Fuses attention computation, avoids materializing N^2 matrix
  • Grouped-Query Attention (GQA): Share K,V projections across heads (fewer KV heads than Q heads)
  • Multi-Query Attention (MQA): Single K,V head shared across all Q heads

Head Pruning & Analysis

Redundancy

  • Many heads can be pruned post-training with minimal accuracy loss
  • Voita et al. (2019): Prune 38/48 heads in BERT-large with <1% drop
  • Michel et al. (2019): Single head often sufficient per layer at inference

Specialization

  • Early layers: More syntactic, local patterns
  • Middle layers: Semantic, long-range dependencies
  • Late layers: Task-specific, output-oriented

Variants

Variant Description Use Case
Standard MHA h Q,K,V projections Original Transformer
GQA h_q Q heads, h_{kv} < h_q KV heads LLaMA, PaLM (faster inference)
MQA h_{kv} = 1 Extreme KV cache reduction
DSA Learned sparse top-k per head DeepSeek V3.2

Related Concepts

Sources

  • arxiv-1706.03762: Section 3.2.2 — Multi-head attention formulation; Table 3(A) — Head count ablation