Multi-Head Attention: Beyond Single Focus

Hello! Welcome back to our journey through the Transformer architecture.

In our last lesson, we explored scaled dot-product attention, the mechanism that allows a model to weigh the importance of different words in a sequence to create context-aware embeddings. We saw how Query, Key, and Value vectors interact to produce an output that is highly parallelizable, overcoming a major limitation of RNNs.

However, a single attention mechanism can be limiting. Consider the sentence: "The animal didn't cross the street because it was too tired." When processing the word "it," the model needs to understand that "it" refers to "the animal" (its antecedent) but also that its state is "tired." A single attention head might be forced to average its focus between these two different kinds of relationships.

Today, we will address this limitation by exploring multi-head attention. You will learn how the Transformer runs multiple attention calculations in parallel to capture a richer variety of contextual relationships. This lesson directly addresses the learning outcome: Describe multi-head attention and its advantage over single-head attention.

Why One Head Isn't Always Enough

The core idea behind multi-head attention is to give the model more capacity to understand complex relationships. It improves upon single-head attention in two primary ways:

1. It expands the model’s ability to focus on different positions. As in our example above, it allows the model to jointly attend to information from different parts of the sentence. One "attention head" could learn to connect "it" to "animal," while another connects "it" to "tired."
2. It gives the attention layer multiple “representation subspaces.” This is a more abstract but powerful concept. Instead of having just one set of Query, Key, and Value projection matrices ($W^Q,W^K,W^V$), we will have multiple sets. Each set projects the input embeddings into a different learned subspace. Think of this as looking at the sentence from multiple perspectives simultaneously. One head might focus on syntactic structure, another on semantic meaning, and a third on co-reference.

To use a simple analogy from the resource "The Math Behind Multi-Head Attention," you can think of single-head attention as one spotlight on a stage. It can only illuminate one area at a time. Multi-head attention is like having multiple spotlights, each able to shine on a different performer, allowing the model to see the whole scene more clearly.

The Mechanism of Multi-Head Attention

Now, let's break down how multi-head attention is calculated. The process can be summarized in four main steps, which elegantly combine the single-head attention mechanism we've already learned with the concept of parallel processing.

Take a look at this diagram, which contrasts single-head (Scaled Dot-Product) attention with multi-head attention. Keep it in mind as we go through the steps.

Step 1: Create Per-Head Projections

Instead of just one set of weight matrices for Q, K, and V, we now have h sets, where h is the number of attention heads (a hyperparameter, typically 8 or 12). For each head $i\in \{1,...,h\}$, we have its own weight matrices: $W_i^Q,W_i^K,W_i^V$.

The original input embeddings are used h times, each time being projected by a different set of matrices to create the queries, keys, and values for each head:

Step 2: Calculate Attention in Parallel

Each head now performs the exact same scaled dot-product attention calculation we learned in the last lesson, completely independently of the others.

Because these calculations are independent, they can be executed in parallel, making the mechanism highly efficient on modern hardware like GPUs.

Step 3: Concatenate the Head Outputs

The outputs from all h heads are concatenated into a single, large matrix.

Step 4: Final Linear Projection

This large concatenated matrix is passed through a final linear layer, with its own weight matrix $W^O$, to project it back down to the model's expected dimension ($d_{\text{model}}$). This final projection allows the information learned by all the different heads to be integrated and combined.

The full formula for multi-head attention is:

Let's watch a short video segment that visualizes this process beautifully.

The Mathematics and Implementation

Given your math background, it's valuable to see the formal mathematical representation and how it maps to a practical implementation.

In practice, frameworks like PyTorch and TensorFlow don't literally create h separate matrix multiplications. For efficiency, they perform a single large matrix multiplication to get the Q, K, and V projections for all heads at once, and then reshape the resulting tensor to add a "number of heads" dimension. This achieves the same outcome with fewer, larger operations, which is much faster on a GPU.

Let's see how this is done in code.

Conclusion

In this lesson, we expanded on our knowledge of self-attention by introducing the multi-head mechanism. By running multiple attention heads in parallel, the Transformer can capture a much richer set of relationships and features from the input data.

Key Takeaways:

• Single-head attention can be limiting, as it may need to average its focus across different types of contextual relationships.
• Multi-head attention addresses this by running multiple self-attention mechanisms in parallel, each with its own set of learnable weight matrices ($W_i^Q,W_i^K,W_i^V$).
• The primary advantages are:
  1. It allows the model to focus on different parts of the input sequence simultaneously.
  2. It gives the model multiple "representation subspaces" to analyze the input from different perspectives.
• The process involves projecting inputs for each head, calculating attention in parallel, concatenating the results, and a final linear projection to integrate the information.

Preview of the Next Lesson

So far, our attention mechanism is "permutation invariant"—if we shuffled the words in a sentence, it would produce nearly the same output because it lacks any sense of order or position. This is a problem, as word order is crucial for meaning. In our next lesson, we will explore positional encoding, the clever technique the Transformer uses to inject information about the position of each word into the model.