Multi-Head Attention › Visualizing Attention Heads

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

Book contents

Section 7 of 14 · 5 min read

Multi-Head Attention › Visualizing Attention HeadsOne section · Sections in this chapter7 / 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

The true power of multi-head attention becomes apparent when we visualize what different heads are doing. Even with random initialization (before any training), the different projection matrices cause each head to compute different attention patterns. After training on language data, these patterns become meaningful: heads specialize to detect specific types of relationships.

Let's visualize the attention patterns for each head on a simple sentence. We'll create random embeddings and see how the four heads distribute their attention differently:

In[10]:
Code
# Create a simple sentence for visualization
sentence = ["The", "cat", "sat", "on", "the", "mat"]
seq_len = len(sentence)

np.random.seed(7)
d_model = 64
num_heads = 4

# Simulate embeddings with structured token relationships:
# "The" (0) and "the" (4) are the same word → similar vectors
# "cat" (1) and "mat" (5) share phonetic similarity → partial overlap
# "sat" (2) and "on" (3) are verb/preposition → distinct from nouns
base = np.random.randn(seq_len, d_model)
# "The"/"the" are highly correlated
base[4] = 0.85 * base[0] + 0.15 * np.random.randn(d_model)
# "cat"/"mat" share partial structure
base[5] = 0.55 * base[1] + 0.45 * np.random.randn(d_model)
embeddings = base[np.newaxis, :, :]  # (1, 6, d_model)

# Create MHA and compute attention
mha = MultiHeadAttention(d_model, num_heads, seed=17)
output, attention_weights = mha.forward(embeddings)
Out[11]:
Visualization
Heatmap showing attention weights for head 1.
Head 1 attention pattern showing how each token distributes attention across the sequence.
Heatmap showing attention weights for head 2.
Head 2 attention pattern with a different learned focus than head 1.
Heatmap showing attention weights for head 3.
Head 3 attention pattern demonstrating further specialization.
Heatmap showing attention weights for head 4.
Head 4 attention pattern completing the four-head perspective.

The heatmaps reveal that even with random initialization, each head produces a distinct attention pattern. Some heads distribute attention more uniformly, while others concentrate it on specific positions. In a trained model, these differences become meaningful and interpretable: one head might consistently attend to the previous word (useful for language modeling), another to syntactically related words (useful for parsing), and another to semantically similar words (useful for coreference resolution).

The key observation is that the same input, "The cat sat on the mat," produces four different perspectives on token relationships. This diversity is precisely what we wanted when we motivated multi-head attention: the ability to capture multiple types of dependencies simultaneously.

Quantifying Head Focus with Entropy

We can quantify how "focused" or "distributed" each head's attention is using entropy. A head with low entropy concentrates attention on a few positions, while a head with high entropy distributes attention more uniformly. This provides a numerical measure of the diversity we observed visually.

In[12]:
Code
def attention_entropy(weights):
    """
    Compute the average entropy of attention distributions.

    Higher entropy means more uniform attention (less focused).
    Lower entropy means attention is concentrated on fewer positions.
    """
    # Avoid log(0) by adding small epsilon
    eps = 1e-10
    # Compute entropy for each query position, average over positions
    entropy = -np.sum(weights * np.log(weights + eps), axis=-1)
    return entropy.mean()


# Compute entropy for each head
head_entropies = []
for h in range(num_heads):
    entropy = attention_entropy(attention_weights[0, h])
    head_entropies.append(entropy)

# Maximum possible entropy (uniform distribution over 6 positions)
max_entropy = np.log(seq_len)
Out[13]:
Visualization
Bar chart comparing entropy values for four attention heads.
Attention entropy by head. Lower entropy indicates more focused attention patterns, while higher entropy indicates more distributed attention. The dashed line shows maximum possible entropy (uniform attention across all 6 positions).
Out[14]:
Console
Entropy analysis by head:
  Maximum possible entropy (uniform over 6 positions): 1.792

  Head 1: entropy = 1.516 (15.4% focused)
  Head 2: entropy = 1.445 (19.4% focused)
  Head 3: entropy = 1.474 (17.7% focused)
  Head 4: entropy = 1.369 (23.6% focused)

Heads with lower entropy relative to the maximum are more "focused," concentrating their attention on specific positions. Heads closer to the maximum entropy distribute attention more uniformly. This variation in focus is another dimension of head diversity: some heads act as sharp selectors, while others aggregate information more broadly.