Multi-Head Attention: Understanding Subspace Representation

Introduction

In our last lesson, we focused on causal masking, a critical adaptation of self-attention that makes autoregressive generation possible. We learned how to prevent a model from "cheating" by looking at future tokens during training, ensuring that predictions at any step are only based on past information.

So far, we've treated self-attention as a single, monolithic calculation. This is powerful, but it has a limitation: it learns only one way to relate words to each other. Consider the word "bank". It could refer to a financial institution or a river's edge. A single attention mechanism might struggle to capture these different contextual meanings simultaneously.

This brings us to today's learning objective: to describe multi-head attention and how it enables attending to different representation subspaces simultaneously. We will see how running multiple attention mechanisms in parallel allows the model to capture a richer and more diverse set of relationships within the data.

1. The Rationale: Why More Than One Head?

The core idea behind multi-head attention is that a single attention mechanism might not be enough to capture the complex relationships present in a sequence. By using multiple "heads," the model can learn to focus on different aspects of the input.

Think of it like a team of specialists analyzing a sentence:

• One specialist might focus on syntactic relationships (e.g., which noun is the subject of which verb?).
• Another might track semantic relationships (e.g., which words are related in meaning?).
• A third could identify co-references (e.g., which pronoun refers to which noun?).

Multi-head attention gives the model this ability by creating multiple "representation subspaces." Each head has its own set of learned weight matrices, allowing it to project the input embeddings into a different space and focus on different patterns.

The famous "Illustrated Transformer" article provides a great example.

As the article shows, when encoding the word "it" in the sentence "The animal didn’t cross the street because it was too tired", different attention heads learn to connect "it" to different words like "animal" and "tired", capturing different facets of its meaning in context.

2. The Multi-Head Attention Mechanism

Now, let's deconstruct the mechanism itself. It consists of four main steps:

1. Linear Projections: Instead of just one set of weight matrices ($W_Q,W_K,W_V$), we now have $h$ sets, where $h$ is the number of heads. For each head $i$, we create unique Query, Key, and Value matrices by projecting the input embeddings $X$ with that head's specific weights:
   $Q_i=X{W}_Q^{(i)},K_i=X{W}_K^{(i)},V_i=X{W}_V^{(i)}$

2. Parallel Attention Calculation: We then perform scaled dot-product attention for each head in parallel, producing $h$ separate output matrices:
   ${\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$

3. Concatenation: The outputs of all the heads are concatenated together into a single large matrix.

4. Final Projection: This concatenated matrix is then projected back to the model's original embedding dimension using one final learned weight matrix, $W_O$.

This entire process is beautifully visualized in the following diagram.

To see this process explained step-by-step with clear animations, the following video is an excellent resource.

The key insight from the video is that the final linear projection layer ($W_O$) is not just for reshaping the output. It's a trainable layer that learns how to best combine the information from all the different attention heads. It can learn to amplify the outputs of heads that provide useful signals and down-weight those that are less relevant for a given context.

3. A Look at the Dimensions

Understanding the matrix dimensions is crucial for grasping how multi-head attention is implemented efficiently. Let's assume an input sequence of length seq_len and an embedding dimension d_model.

In the original Transformer paper, the total d_model is split across the h heads. This means the dimension for each head's Query, Key, and Value vectors is $d_{head}=d_{model}/h$.

Let's trace the dimensions through the process with a common example where d_model = 512 and h = 8 heads.

1. Input Embeddings ($X$): The shape is (seq_len, 512).
2. Projection: For each head $i$, the projection matrices $W_Q^{(i)},W_K^{(i)},W_V^{(i)}$ project the 512-dimensional embeddings down to 64-dimensional vectors ($d_{head}=512/8=64$). So, $Q_i,K_i,V_i$ each have the shape (seq_len, 64).
3. Attention Output: The output of each attention head, head_i, also has the shape (seq_len, 64).
4. Concatenation: When we concatenate the 8 heads, we combine their 64-dimensional outputs along the feature dimension. The resulting matrix has the shape (seq_len, 8 * 64), which is (seq_len, 512). We are back to our original embedding dimension.
5. Final Projection: The final weight matrix, $W_O$, has a shape of (512, 512). It takes the concatenated matrix and performs a final linear transformation, producing the output of the multi-head attention block, which also has the shape (seq_len, 512).

This design ensures that the output of a multi-head attention layer has the same dimensions as its input, allowing us to stack these layers on top of each other to build a deep model.

The following resource provides another excellent step-by-step walkthrough, focusing on the matrix dimensions and including code snippets that make the process very concrete.

Conclusion

Today, we've explored multi-head attention, a key innovation of the Transformer architecture that significantly enhances the model's representative power.

Key Takeaways:

• Purpose: Multi-head attention allows a model to jointly attend to information from different "representation subspaces" at different positions. This is like having a team of specialists analyzing the input from various perspectives.
• Mechanism: It works by running multiple self-attention mechanisms (heads) in parallel. Each head has its own learned set of projection matrices ($W_Q,W_K,W_V$).
• Combining Heads: The outputs from the parallel heads are concatenated and then transformed by a final learned linear projection layer ($W_O$), which learns to synthesize the information from all heads.
• Dimensionality: A key design choice is to split the model's total embedding dimension (d_model) across the heads, which allows for this richer representation without a massive increase in computational cost compared to a single large head.

Preview of the Next Lesson:
We have now covered the core attention mechanisms: scaled dot-product attention, causal (masked) attention, and multi-head attention. In our next lesson, we will put all the pieces together and describe the full structure of a Transformer block, including the feed-forward network, residual connections, and layer normalization. This will give us a complete picture of the fundamental building block of modern LLMs.