Multi-Head Attention › Head Specialization

Michael BrenndoerferUpdated May 30, 20255 min read
Language AI HandbookAll chapters
Part 13 of 67Self-Attention
Chapter 94Multi-Head Attention › Head Specialization
Reading section 8 of 14 · Head Specialization · 5 minSection 8 of 14 · 5 min

Book contents

Section 8 of 14 · 5 min read

Multi-Head Attention › Head SpecializationOne section · Sections in this chapter8 / 14 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Research on trained transformers reveals that attention heads develop distinct specializations. The emergence of interpretable, consistent patterns across very different models is one of the most striking properties of multi-head attention. It suggests that these specializations are not accidents of initialization or quirks of a particular training run; they reflect real structure in language that gradient descent reliably discovers.

Some commonly observed patterns include:

  • Positional heads: Attend to fixed relative positions (e.g., always attend to the previous token, or always attend to the first token)
  • Syntactic heads: Track grammatical relationships like subject-verb or modifier-noun connections
  • Coreference heads: Link pronouns to their antecedents across the sequence
  • Rare word heads: Focus on infrequent tokens that carry high information content
  • Separator heads: Attend to punctuation and sentence boundaries

This specialization emerges purely from training on language modeling objectives. No explicit supervision tells heads what to focus on; they discover useful patterns through gradient descent.

Historical Context: Discovering What Heads Learn

The systematic study of attention head interpretability began in earnest after the original "Attention is All You Need" paper (Vaswani et al., 2017). Early analyses by Vig (2019) and Clark et al. (2019) visualized BERT's attention patterns and found that individual heads reliably tracked specific linguistic phenomena: one head consistently attended to the direct object of a verb, another to the head of a noun phrase, and so on. These findings were surprising because BERT was trained only on masked language modeling and next-sentence prediction, with no explicit linguistic supervision whatsoever. The implication was that language modeling itself, applied at scale, was sufficient to induce representations that recover classical linguistic categories. Michel et al. (2019) later showed that a majority of BERT's 144 attention heads could be pruned at test time with minimal accuracy loss, suggesting that while many heads are useful, the critical patterns are concentrated in a subset.

What does this mean in practice? When you deploy a transformer model, its attention heads are not interchangeable. Some heads carry far more weight for specific task types: syntactic heads matter most for tasks requiring grammatical understanding, while semantic heads matter most for tasks requiring world-knowledge reasoning. This is one reason why fine-tuning a pre-trained transformer often works so well: the heads that are relevant for your task already exist in the pre-trained model, and fine-tuning adjusts their relative weights rather than learning entirely new patterns from scratch.

In[15]:
Code
# Simulate different head specialization patterns
np.random.seed(42)
seq_len = 8
tokens = ["[CLS]", "The", "quick", "fox", "jumps", "over", "dogs", "[SEP]"]


# Create simulated attention patterns for different head types
def create_positional_head(n, offset=-1):
    """Head that attends to a fixed relative position."""
    weights = np.zeros((n, n))
    for i in range(n):
        target = max(0, min(n - 1, i + offset))
        weights[i, target] = 0.8
        # Some attention to self
        weights[i, i] = 0.1
        # Distribute remaining mass
        remaining = 0.1 / (n - 2) if n > 2 else 0
        for j in range(n):
            if j != i and j != target:
                weights[i, j] = remaining
    return weights


def create_first_token_head(n):
    """Head that always attends to [CLS] token."""
    weights = np.zeros((n, n))
    weights[:, 0] = 0.7  # Strong attention to first token
    # Distribute remaining mass
    for i in range(n):
        weights[i, i] += 0.2  # Some self-attention
        remaining = 0.1 / (n - 2) if n > 2 else 0
        for j in range(1, n):
            if j != i:
                weights[i, j] += remaining
    return weights


def create_local_context_head(n, window=2):
    """Head that attends to nearby tokens."""
    weights = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            distance = abs(i - j)
            if distance <= window:
                weights[i, j] = 1.0 / (distance + 1)
        weights[i] = weights[i] / weights[i].sum()
    return weights


# Generate example patterns
prev_token_head = create_positional_head(seq_len, offset=-1)
cls_head = create_first_token_head(seq_len)
local_head = create_local_context_head(seq_len, window=2)
Out[16]:
Visualization
Heatmap showing diagonal attention pattern for previous token head.
Previous Token Head: attends strongly to position i-1, capturing local sequential dependencies.
Heatmap showing column attention pattern for CLS aggregation head.
[CLS] Aggregation Head: all tokens attend to the first token, useful for sequence-level classification.
Heatmap showing banded attention pattern for local context head.
Local Context Head: attention concentrated within a sliding window, capturing phrase-level context.

The three patterns illustrate different specializations:

  • Previous token head: The diagonal pattern shows each token attending primarily to its predecessor. This captures local sequential dependencies, useful for language modeling where the previous word strongly predicts the next.

  • [CLS] aggregation head: The first column lights up, showing all tokens attending to the special classification token. This pattern aggregates information from the entire sequence into one position, commonly used for classification tasks.

  • Local context head: The banded pattern shows attention concentrated within a local window. This captures phrase-level relationships without being distracted by distant tokens.