Multi-Head Attention: Capturing Diverse Context

Hello! Welcome to the sixth lesson in our journey through the Transformer architecture.

In our last lesson, we explored the self-attention mechanism, understanding how it uses Query, Key, and Value vectors to create context-aware representations for each token. We ended with a crucial question: What if a single attention calculation isn't enough to capture all the complex relationships in a sentence?

This lesson directly addresses that question. We'll explore multi-head attention, which is the Transformer's solution for capturing diverse contextual patterns simultaneously. This is not a new type of attention but rather a clever way of running multiple self-attention "heads" in parallel.

Our learning outcome is to: Describe the role of multi-head attention in capturing diverse contextual patterns.

By the end of this lesson, you'll understand why modern models like GPT use dozens of attention heads and how they work together to build a rich, multi-faceted understanding of text.

1. Why One Head Isn't Enough

In the last lesson, we saw how a single attention head can learn to connect a pronoun like "it" to the noun it refers to. But a sentence contains many different kinds of relationships happening at once:

• Syntactic dependencies: Subject-verb-object relationships (e.g., "creature ate berry").
• Modification: Adjectives modifying nouns ("fluffy creature"), adverbs modifying verbs ("happily ate").
• Semantic relationships: Words that are part of a broader concept (e.g., "Eiffel," "tower," and "Paris").

A single attention head, with its one set of learned weight matrices ($W_Q,W_K,W_V$), might become a specialist in one of these areas. Forcing it to handle all of them at once can be inefficient and limit its effectiveness.

The solution is to use multiple attention heads, each with its own set of learned weights. This allows each head to specialize in detecting a different kind of relationship.

This video provides a great analogy for why multiple "perspectives" or "filters" are useful.

Just as you might want to analyze different parts of an image to make a decision, an LLM uses multiple attention heads to analyze different linguistic patterns in the text. One head might track subject-verb agreement, another might link related nouns, and a third could focus on dependencies between adjacent words.

2. The Architecture of Multi-Head Attention

So, how does this work in practice? The core idea is to "split" the work of attention across multiple heads, let them process the input in parallel, and then combine their results.

Here is a high-level view of the architecture:

Let's break down the process shown in panel (b) of the diagram:

1. Initial Projections: We start with the input embeddings for our sequence of tokens. For each of the $N$ attention heads, we have a unique set of weight matrices: $W_Q^{(i)},W_K^{(i)},W_V^{(i)}$, where $i$ is the head index from 1 to $N$.

2. Dimensionality Split: This is a key point. We don't just run $N$ full-sized attention calculations. Instead, we split the embedding dimension (d_model) across the heads. If our main embedding vector has a size of d_model = 768 and we have n_heads = 12, then each head will work with smaller Query, Key, and Value vectors of dimension d_head = 768 / 12 = 64. This keeps the total computation roughly the same as a single head operating on the full dimension, but allows for parallel specialization.

3. Parallel Attention Calculation: Each head $i$ now performs the same scaled dot-product attention calculation we learned in the previous lesson, but on its smaller, dedicated Q, K, and V vectors. This produces an output vector for each token, for each head.
   ${\text{head}}_i=\text{Attention}(Q{W}_Q^{(i)},K{W}_K^{(i)},V{W}_V^{(i)})$

4. Concatenation & Final Projection: The output vectors from all $N$ heads are concatenated together. For each token, we now have a large vector that combines the "perspectives" of all the heads. This concatenated vector is then passed through a final linear projection layer with its own weight matrix, $W_O$, to produce the final output of the multi-head attention layer. This projection mixes the information from the different heads.

This diagram provides another excellent view of the entire data flow.

3. A Software Engineer's Perspective

Given your background, it can be very helpful to think about this structure in terms of code. The following resource explains the Transformer architecture from a software engineer's point of view, using Rust for its examples. The type system and structure make the relationships between components very clear.

As the article shows, the AttnLayer effectively orchestrates the work. It distributes the task to individual AttnHeads, each of which operates in its own lower-dimensional subspace. It then gathers the results and produces a single, unified update vector for each token in the sequence. This combination of parallel processing and final integration is what gives multi-head attention its power.

To reinforce these steps, the "Transformer Explainer" resource provides another great step-by-step walkthrough.

Conclusion

In this lesson, we scaled up our understanding from a single self-attention mechanism to the full multi-head attention architecture used in all modern Transformers. By running multiple attention heads in parallel, each specializing in different patterns, the model can build a far more nuanced and robust understanding of the input text.

Key Takeaways:

• Diverse Patterns: Multi-head attention allows a model to capture different types of relationships (e.g., syntactic, semantic) in parallel, as each head can learn to specialize.
• Architecture: The process involves splitting the input into lower-dimensional chunks for each head, performing scaled dot-product attention independently in each head, concatenating the results, and applying a final linear projection.
• Computational Efficiency: Splitting the d_model dimension across multiple heads means the total computation is comparable to a single, full-dimensional attention calculation, but with the benefit of parallelization and specialization.
• Mixing Information: The final projection layer ($W_O$) is crucial as it allows the model to learn how to best combine the outputs from the different specialist heads.

In the previous lessons, we've assembled the key components: token embeddings, positional encodings, and multi-head attention. In our next lesson, we will see how these are put together inside a decoder-only model like GPT to generate text one token at a time.