Multi-Head Attention: Unpacking Representation Subspaces

The Limitation of a Single Perspective

Imagine you're analyzing the sentence: "The mole on his arm was not a mole but a rare insect."

A single attention mechanism might learn to focus on the biological context, connecting "mole" to "arm" and "insect". But it would struggle to also capture the grammatical structure (e.g., connecting "was" to "mole") or other nuances at the same time. It's forced to average its attention across all these different types of relationships.

This is the core limitation of what we've learned so far, which is often called "single-head attention."

The Solution: Multiple Heads Working in Parallel

The solution proposed in the original "Attention Is All You Need" paper is elegant: instead of one attention mechanism, let's use several in parallel. This is called Multi-Head Attention.

Each "head" is essentially an independent scaled dot-product attention mechanism. The key idea is that each head will be encouraged to learn a different type of relationship from the data.

This diagram provides a clear visual comparison between the single-head architecture we know and the multi-head architecture we're about to explore.

As the diagram shows, the process involves three main stages:

1. Projecting the Q, K, and V inputs for each head.
2. Running scaled dot-product attention in parallel for each head.
3. Concatenating the results and projecting them back to the expected dimension.

How Multi-Head Attention Works

Let's watch a video that provides a fantastic high-level intuition for why this parallel structure is so powerful.

The video brilliantly illustrates the goal: to give the model the capacity to learn many distinct ways that context can change meaning. Now, let's look at the mechanics of how this is achieved.

1. Creating Different Representation Subspaces

This is the most critical step. If we simply copied the same Q, K, and V matrices for each head, they would all perform the same calculation and learn the same thing. To prevent this, we first project the original Q, K, and V matrices into different, lower-dimensional spaces for each head.

For a model with $h$ heads, we create $h$ sets of unique, learnable weight matrices:

• $W_1^Q,W_2^Q,...,W_h^Q$
• $W_1^K,W_2^K,...,W_h^K$
• $W_1^V,W_2^V,...,W_h^V$

For each head $i$, we calculate its specific query, key, and value:
$Q_i=Q{W}_i^Q$
$K_i=K{W}_i^K$
$V_i=V{W}_i^V$

These projection matrices ($W_i$) are the key. They act like different lenses, each transforming the original embedding into a different representation subspace. In each subspace, the model can focus on a different aspect of the word's meaning or its relationship with other words.

2. Calculating Attention in Parallel

Now, inside each subspace, we perform the same scaled dot-product attention calculation that we learned in the last lesson:
${\text{Head}}_i=\text{Attention}(Q_i,K_i,V_i)=\text{softmax}\left(\frac{Q_i{K}_i^T}{\sqrt{d_k}}\right)V_i$
Since each head is working with its own projected version of Q, K, and V, each one will produce a different attention score matrix and, consequently, a different output.

3. Concatenating and Projecting

Once we have the output from all $h$ heads, we concatenate them back together into one large matrix.
$\text{MultiHeadOutput}=\text{Concat}({\text{Head}}_1,{\text{Head}}_2,...,{\text{Head}}_h)$
This concatenated matrix now contains the information learned from all the different subspaces. However, its dimension is too large. So, we use one final, learnable weight matrix, $W^O$, to project it back down to the model's original embedding size.
$\text{Final Output}=\text{MultiHeadOutput}\cdot W^O$

This final output vector is a rich representation that has incorporated context from multiple perspectives simultaneously.

Let's watch a more detailed walkthrough that shows the matrix operations and dimensions, which should resonate with your engineering background.

Capturing Richer Interpretations

So, what do these different subspaces actually learn? It's hard to say definitively, as the model learns whatever is optimal for its task. But we can build some intuition.

As the article explains, by splitting the embedding into different subspaces, each head can learn to focus on different aspects of a word's meaning and relationships. One head might focus on syntactic structure, another on semantic categories, and another on long-distance dependencies.

We can even visualize this. The different attention patterns produced by each head show what they are "paying attention" to.

This visualization from the video by Umar Jamil (based on the original paper) shows how different heads (represented by colors) create different connections between words. For the word 'making', some heads (blue, green) strongly connect it to 'difficult', while others (pink) connect it to '2009', demonstrating they've learned to focus on different contextual relationships.

Conclusion

In this lesson, we've unpacked the "multi-head" part of the attention mechanism, which is central to the power of Transformers.

Key Takeaways:

• Single-Head Limitation: A single attention head struggles to capture multiple types of relationships (e.g., syntactic, semantic) simultaneously.
• Parallel Processing: Multi-head attention addresses this by running multiple attention mechanisms in parallel.
• Representation Subspaces: The key to multi-head attention is the use of unique, learnable projection matrices ($W_i$) for each head. These matrices project the input embeddings into different subspaces, allowing each head to learn a unique perspective.
• Richer Context: By combining the outputs from all heads, the model produces a final representation that is enriched with contextual information from many different viewpoints.

Preview of the Next Lesson:

So far, our attention mechanism is a "bag of words" model—it considers all words in the context but has no inherent sense of their order. If we shuffled a sentence, the attention output would be the same! For language, sequence order is crucial. In our next lesson, "Describe the role of positional encodings," we will explore the clever technique Transformers use to inject this vital sequential information into the model.