Multi-Head Attention Deep Dive

🎯 Learning Objectives

• Explain why multiple attention heads are beneficial over a single head
• Understand how different heads specialize in different linguistic patterns
• Compute multi-head attention outputs through the full pipeline
• Analyze head importance and redundancy

📋 Prerequisites

• Self-Attention & QKV (previous topic)
• Matrix concatenation and linear algebra

1. 📖 Core Content

1.1 Intuition: Why Multiple Heads?

• Head 1: Might focus on syntactic dependencies (subject-verb relationships)
• Head 2: Might focus on semantic similarity (related concepts)
• Head 3: Might focus on positional proximity (nearby words)
• Head 4: Might focus on coreference (pronouns and their antecedents) Each head sees a different projected version of the input and computes attention independently.

1.2 Formal Definition

• Each head has 3 projection matrices: $W_i^Q, W_i^K, W_i^V \in \mathbb{R}^{d_{model} \times d_k}$
• Output projection: $W^O \in \mathbb{R}^{h d_v \times d_{model}}$
• Total attention parameters: $3 \cdot h \cdot d_{model} \cdot d_k + d_{model} \cdot h \cdot d_v$ With $d_k = d_v = d_{model}/h$:
• Per head: $3 \cdot d_{model} \cdot (d_{model}/h) = 3d_{model}^2/h$
• All heads: $3d_{model}^2$
• Output projection: $d_{model}^2$
• Total: $4d_{model}^2$ — same as a single head with d_model dimension! Key insight: Multi-head attention doesn't increase the total parameter count compared to a single large head. It just distributes the capacity across multiple subspaces.

1.3 Multiplication by 4 Rule

1.4 Complete Multi-Head Worked Example

• 2 tokens, d_model = 6, h = 3 heads, d_k = d_v = 2 Step 1: Input

1.5 Head Specialization

• Lower layers: Heads focus on positional relationships and surface features
• Middle layers: Heads capture syntactic relationships
• Higher layers: Heads capture semantic relationships

1.6 Head Redundancy and Pruning

1. Many heads are redundant: Pruning up to 40% of heads can maintain performance
2. Head importance varies: Some heads contribute critically to certain tasks
3. Task-specific heads: Different tasks rely on different subsets of heads

python
Run Inline
# runnable
import numpy as np
def compute_head_importance(scores_before, scores_after, baseline_performance):
    """
    Estimate head importance by measuring performance drop when head is masked
    Args:
        scores_before: Performance with all heads
        scores_after: Performance with one head ablated
        baseline_performance: Initial performance metric
    Returns:
        importance: Relative importance of each head
    """
    num_heads = len(scores_after)
    importance = np.zeros(num_heads)
    for i in range(num_heads):
        # Larger drop = more important
        importance[i] = (baseline_performance - scores_after[i]) / baseline_performance
    return importance
# Hypothetical example: 8 heads
baseline = 0.92  # F1 score
scores_with_ablation = [0.91, 0.88, 0.90, 0.85, 0.91, 0.89, 0.92, 0.91]
importance = compute_head_importance([], scores_with_ablation, baseline)
print("Head importance scores:")
for i, imp in enumerate(importance):
    print(f"Head {i+1}: {imp:.3f} (ablated score: {scores_with_ablation[i]:.3f})")
# Heads with low importance could potentially be pruned
threshold = 0.02
prunable = [i for i, imp in enumerate(importance) if imp < threshold]
print(f"\nPrunable heads (importance < {threshold}): {[h+1 for h in prunable]}")

1.7 Why This Matters

1. Parallel pattern capture: Attend to syntax, semantics, and position simultaneously
2. Robustness: Redundancy provides graceful degradation
3. Interpretability: Head analysis reveals what the model learns
4. Efficiency: Same parameter budget as single head, but richer representations

4. 📐 Key Formulas / Concepts

5. ⚠️ Common Pitfalls

Pitfall 1: Assuming all heads learn unique functions

Pitfall 2: Forgetting the output projection W^O

Pitfall 3: Thinking more heads is always better

6. 📝 Practice Questions

Q1: For GPT-3 (d_model=12288, h=96), what is d_k?

d_k = 12288 / 96 = 128. Each head operates in 128-dimensional space.

Total attention parameters = 4 × 12288² ≈ 604 million parameters (out of 175B total). Q2: If we have d_model=512 and h=16, what's d_k? Is this reasonable?

d_k = 512/16 = 32. This is at the lower end of typical values. With only 32 dimensions per head, each head has limited representational capacity. This might still work but is less common than h=8 (d_k=64) or h=12 (d_k≈43). Q3: Show that multi-head attention with h=1 is equivalent to single-head attention.

With h=1:

• d_k = d_model (since d_model/1 = d_model)
• $W^Q, W^K, W^V \in \mathbb{R}^{d_{model} \times d_{model}}$
• Concat has only 1 head, so Concat(head₁) = head₁
• $W^O \in \mathbb{R}^{d_{model} \times d_{model}}$

MultiHead(Q, K, V) = head₁ · W^O = Attention(QW^Q, KW^K, VW^V)W^O

Since W^O is a learned square matrix, this is equivalent to single-head attention with an additional linear transform. In practice, the two W^O's can be absorbed into the projections. Q4: Compute the number of attention parameters for d_model=1024, h=16.

d_k = 1024/16 = 64

Per head: 3 × 1024 × 64 = 196,608 All heads: 16 × 196,608 = 3,145,728 W^O: 1024 × 1024 = 1,048,576

Total: 3,145,728 + 1,048,576 = 4,194,304 = 4 × 1024² ✓ Q5: Why is the FFN dimension typically 4× d_model?

This is an architectural choice from the original Transformer paper that has proven effective. The expansion factor 4 provides enough capacity for non-linear transformations. Common dimensions: 512→2048, 768→3072, 1024→4096. The FFN has 8d_model² parameters (two layers: d_model → 4d_model → d_model).

FFN params = d_model × 4d_model + 4d_model × d_model = 8d_model² Q6: If a model's attention patterns show that all heads have almost identical attention distributions, what might this indicate?

This could indicate:

1. Underfitting: The model hasn't learned to differentiate head roles (training may be insufficient)
2. Low model capacity: d_k might be too small for meaningful specialization
3. Data issue: The training data may lack the diversity needed for specialization
4. Over-regularization: Strong regularization might be pushing heads toward uniformity

In well-trained models, heads typically show distinct attention patterns. Q7: What happens to multi-head attention if we remove the output projection W^O?

Without W^O, each head's output is simply concatenated. This means:

• The output has h independent "blocks" of information
• No mixing occurs between heads
• The model can't weight heads differently
• Performance typically degrades because the model can't integrate information across heads Q8: In a Transformer with 6 encoder layers, each with 8 heads, how many unique attention mechanisms process the input?

6 layers × 1 self-attention sublayer × 8 heads = 48 total attention mechanisms. Each of these 48 heads learns a different attention pattern. In the decoder, there are 2 attention sublayers per layer (masked self-attention + cross-attention), so 6 × 2 × 8 = 96 additional mechanisms. Q9: A paper finds that Head 5 in layer 3 consistently attends to the [SEP] token. What might this head's role be?

The [SEP] token in BERT separates sentences. A head that always attends to [SEP] might be involved in:

1. Sentence boundary detection: Tracking where one sentence ends and another begins
2. Next sentence prediction: The NSP task requires understanding sentence relationships
3. Global aggregation: Accumulating information from both sentences before classification

This is an example of head specialization — certain heads develop consistent, interpretable functions. Q10: Derive why multi-head attention doesn't increase the parameter count relative to a single large head.

Single head (d_k = d_model):

• W^Q: d_model² params
• W^K: d_model² params
• W^V: d_model² params
• Total: 3d_model² params

Multi-head (h heads, d_k = d_model/h):

• Per head: 3 × d_model × (d_model/h) = 3d_model²/h params
• All heads: h × 3d_model²/h = 3d_model² params
• W^O: h × (d_model/h) × d_model = d_model² params
• Total: 4d_model² params

Multi-head adds the output projection W^O (d_model² params) but doesn't change the QKV projection total. So it adds d_model² parameters (33% increase), not a factor of h increase.

7. 🔗 Cross-References

• Previous topic: Self-Attention & QKV
• Next topic: Positional Encoding
• Related: Transformer Encoder & Decoder Layers (Week 2) Join Discord PreviousSelf-Attention & QKVNextPositional Encoding