Multi-Head Attention: Attending to Multiple Subspaces

Welcome to your third lesson in the "LLM and Transformer Fundamentals" module.

In the previous lesson, we dissected the scaled dot-product attention mechanism. We saw how a single "head" of attention allows a word to update its own representation by calculating a weighted sum of all other words' values, based on query-key similarity. This single head learns to identify one specific type of relationship in the text.

However, language is incredibly rich and layered. Consider the sentence: "The model didn't work because it was too complex." The word "it" has a syntactic relationship with "model" (anaphora) and a semantic relationship with "complex" (a property). A single attention head might struggle to capture both relationships simultaneously.

This lesson addresses that limitation. Our goal is to explain how multi-head attention allows the model to attend to different information subspaces simultaneously. We will see how the Transformer architecture runs multiple attention mechanisms in parallel to build a more nuanced and robust understanding of context.

1. The Core Idea: From One Generalist to a Team of Specialists

Instead of having one large attention mechanism, multi-head attention employs multiple, smaller attention "heads" that operate in parallel. Each head is a complete, independent scaled dot-product attention mechanism like the one we studied previously.

The key insight is that each head has its own set of learned weight matrices ($W^Q$, $W^K$, $W^V$). Because these are initialized randomly and trained independently, each head learns to specialize in detecting different kinds of relationships or patterns in the input text.

• One head might learn to track syntactic dependencies.
• Another might focus on semantic relationships, like identifying synonyms or antonyms.
• A third could learn to track positional information, like which words are nearby.

This is analogous to having a committee of experts analyze a problem. Each expert brings a different perspective, and their combined insights lead to a more comprehensive conclusion than any single expert could provide alone.

The following video provides an excellent high-level introduction to this idea.

2. The Mechanics of Multi-Head Attention

The process of multi-head attention can be broken down into three main steps: projection, parallel attention calculation, and combination.

Let's assume our model has a hidden dimension of $d_{\text{model}}$ (e.g., 768 for BERT-base) and we want to use $h$ attention heads (e.g., 12).

Step 1: Projection into Subspaces

First, we take our input embedding matrix $X$. For each head $i\in \{1,...,h\}$, we generate its own query, key, and value matrices using a unique set of weight matrices $W_i^Q$, $W_i^K$, and $W_i^V$.

A crucial detail is that these weight matrices project the input embeddings into a lower-dimensional "subspace". Typically, the dimension for each head is $d_k=d_v=d_{\text{model}}/h$. In our example, this would be $768/12=64$.

• $Q_i=X{W}_i^Q$
• $K_i=X{W}_i^K$
• $V_i=X{W}_i^V$

This splitting ensures that the total computational cost of multi-head attention remains roughly the same as a single-head attention using the full $d_{\text{model}}$ dimension. It's an efficient design that trades a single large matrix multiplication for several smaller ones that can be executed in parallel.

Step 2: Parallel Attention Calculations

Now that we have $h$ sets of (Q, K, V) matrices, we perform the scaled dot-product attention calculation for each head independently and in parallel.

Each ${\text{head}}_i$ is now a matrix containing the contextually updated representations as seen from that head's specialized perspective.

Step 3: Concatenation and Final Projection

We now have $h$ separate output matrices, but the next layer in the Transformer (the feed-forward network) expects a single input matrix. To reconcile this, we first concatenate the head outputs along the feature dimension:

The resulting matrix has dimensions (sequence_length, h * d_v), which is (sequence_length, d_model). We're back to our original embedding dimension.

Finally, this concatenated matrix is passed through one last linear projection, governed by a weight matrix $W^O$. This final projection mixes the information from all the heads, allowing the model to learn how to best combine their specialized outputs.

This video provides a very clear, step-by-step walkthrough of this exact process.

3. Visualizing Different "Information Subspaces"

The phrase "different information subspaces" refers to the specialized perspectives that each head learns. Because each head has its own trainable projection matrices ($W_i^Q,W_i^K,W_i^V$), it learns to project the input embeddings into a representation space where a specific kind of information is highlighted.

The following article provides one of the best visualizations of this concept.

This ability to simultaneously process different types of relationships is what gives multi-head attention its power.

4. Trade-offs: Are More Heads Always Better?

From a systems engineering perspective, it's crucial to question design choices and understand their trade-offs. While multi-head attention is powerful, is it optimally efficient? Is there redundancy? Your experience building high-performance services like Rudi will give you an appreciation for this kind of critical analysis.

A fascinating research paper explored exactly this question by systematically "pruning" (removing) attention heads from a trained BERT model to see what would happen.

This research provides a nuanced answer: multi-head attention introduces a beneficial "excess of capacity." This redundancy may be vital for stable and effective training, even if some heads are not critical for inference on a specific task. This is a classic trade-off between training robustness and inference efficiency.

Conclusion

In this lesson, we've seen how multi-head attention enhances the expressive power of the Transformer. By running multiple attention heads in parallel, the model can simultaneously analyze an input sequence from various perspectives, capturing a rich and diverse set of contextual relationships.

Key Takeaways:

• Multi-head attention consists of several attention heads running in parallel.
• Each head has its own learnable weight matrices ($W_i^Q,W_i^K,W_i^V$), allowing it to specialize and focus on a different information subspace.
• The process involves:
  1. Projecting the input embeddings into lower-dimensional Q, K, V subspaces for each head.
  2. Calculating attention in parallel for all heads.
  3. Concatenating the outputs and passing them through a final projection layer.
• This design allows the model to capture complex relationships, such as resolving pronoun ambiguity by considering both syntactic and semantic cues simultaneously.
• While highly effective, there is significant redundancy among heads, which has practical implications for model compression and efficiency.

Next Lesson Preview

So far, our attention mechanism treats the input as an unordered set of vectors. It has no inherent sense of sequence order. "The cat chased the mouse" and "The mouse chased the cat" would look identical to the attention mechanism alone. To solve this, the Transformer injects information about the position of each token into its embedding.

In our next lesson, we will explore Positional Encodings, the clever technique used to give the model a sense of sequence order, a crucial piece of the puzzle.