Multi-Head Attention: Architecture and Purpose

Recommended resources

The freely available resources this lesson will direct you to.

Introduction

In our last lesson, we derived the self-attention mechanism, establishing it as a powerful method for capturing dependencies within a sequence. We saw how a single "head" of attention calculates a set of weights to produce a new, context-rich representation for each token. However, a single set of attention weights might force the model to average different types of relationships. For instance, in the sentence, "The tired runner drank water," one token might need to attend to "tired" for semantic context and "runner" for syntactic agreement. A single attention mechanism might struggle to capture both optimally.

This lesson introduces multi-head attention, a core innovation from the "Attention Is All You Need" paper. We will explore how this architecture enables the model to process information from different representational subspaces simultaneously. In essence, instead of one generalist attention mechanism, we deploy a team of specialist attention "heads," each focusing on a different aspect of the input sequence. This parallel processing design is not only more expressive but also computationally efficient—a detail that should be of particular interest given your background in performance systems.

The Purpose of Multiple Heads

The fundamental idea behind multi-head attention is to allow the model to jointly attend to information from different perspectives. A single attention head learns one specific function for mapping queries to keys. By using multiple heads, we empower the model to learn several such functions in parallel.

Imagine you are analyzing a complex system. You might have different analytical tools: one to check for performance bottlenecks, one for memory leaks, and another for logical errors. Each tool gives you a different "view" of the system. Multi-head attention works analogously. Each head can be thought of as a separate tool that learns to focus on a distinct type of relationship within the data, such as:

Syntactic dependencies (e.g., subject-verb agreement).
Semantic relationships (e.g., synonymy or antonymy).
Positional information (e.g., attending to the previous word).

To get a high-level intuition for why this is so powerful, the following video provides an excellent conceptual overview.

Attention in transformers, step-by-step | Deep Learning Chapter 6

This video from 3Blue1Brown explains the motivation for using multiple attention heads to capture diverse contextual relationships.

Please watch the segment from this section on multi-head attention. Focus on the idea that each head can learn a different kind of contextual update, allowing the model to become much more expressive.

The Architecture of Multi-Head Attention

Let's now formalize the architecture. We begin with our input sequence representation $X \in R^{n \times d_{model}}$ , where $n$ is the sequence length and $d_{model}$ is the embedding dimension.

Instead of learning a single set of large projection matrices $W_{Q}, W_{K}, W_{V} \in R^{d_{model} \times d_{model}}$ , we learn $h$ sets of smaller projection matrices for each head $i \in {1, \dots, h}$ :

Query matrices: $W_{i}^{Q} \in R^{d_{model} \times d_{k}}$
Key matrices: $W_{i}^{K} \in R^{d_{model} \times d_{k}}$
Value matrices: $W_{i}^{V} \in R^{d_{model} \times d_{v}}$

A critical design choice is to set the head dimensions $d_{k}$ and $d_{v}$ such that $d_{k} = d_{v} = d_{model} / h$ . This ensures that the total computational cost of multi-head attention is roughly the same as that of a single-head attention with the full $d_{model}$ dimension.

The process then unfolds in three stages:

1. Parallel Projections and Attention Calculation

For each head $i$ , we project the input $X$ into its own query, key, and value representations and compute the attention output independently:

head_{i} = Attention (X W_{i}^{Q}, X W_{i}^{K}, X W_{i}^{V}) = softmax (\frac{( X W _{i}^{Q} ) ( X W _{i}^{K} ) ^{T}}{d _{k}}) (X W_{i}^{V})

Each $head_{i}$ is a matrix in $R^{n \times d_{v}}$ , representing one specialized view of the sequence. This is the "massively parallel" part of the design, as all $h$ heads can be computed simultaneously.

This diagram from the original "Attention Is All You Need" paper illustrates the multi-head architecture. The input Q, K, and V are projected into h different subspaces, where scaled dot-product attention is applied in parallel. The resulting outputs are then concatenated and passed through a final linear layer.

2. Concatenation

The outputs of all $h$ heads are concatenated along the feature dimension to form a single, larger matrix:

Concat (head_{1}, head_{2}, \dots, head_{h}) \in R^{n \times (h \cdot d_{v})}

This concatenated matrix now contains the rich, multi-faceted information captured by all the individual attention heads.

3. Final Projection

This large matrix is then projected back to the original model dimension, $d_{model}$ , using a final learnable weight matrix, $W^{O} \in R^{(h \cdot d_{v}) \times d_{model}}$ . This step allows the model to mix and combine the information from all the specialist heads.

The complete formula for Multi-Head Attention is:

MultiHead (X) = Concat (head_{1}, \dots, head_{h}) W^{O}

The diagram below provides a more detailed view of the data flow within a single attention head, which is the fundamental building block replicated in parallel.

A detailed view of a single attention head's computation. The input hidden states are projected by weight matrices $W_q$, $W_k$, and $W_v$ to produce queries, keys, and values. The queries and keys are combined to form attention scores, which are then used to create a weighted sum of the values, yielding the head's output.

Socratic Question: Given our design choice that $d_{v} = d_{model} / h$ , what is the dimension of the concatenated matrix $Concat (head_{1}, \dots, head_{h})$ ? What does this imply about the dimensionality of the final projection matrix $W^{O}$ , and why is this design computationally efficient?

Answer

If $d_v = d_{\text{model}} / h$, the concatenated matrix will have dimensions $\mathbb{R}^{n \times (h \cdot d_v)} = \mathbb{R}^{n \times (h \cdot (d_{\text{model}}/h))} = \mathbb{R}^{n \times d_{\text{model}}}$.

This means the final projection matrix $W^{O}$ will have dimensions $R^{d_{model} \times d_{model}}$ .

This design is efficient because the total computation for all heads is comparable to a single large attention head. The main computational cost is in the matrix multiplications. For a single head of dimension $d_{model}$ , the $Q K^{T}$ multiplication is roughly $O (n^{2} \cdot d_{model})$ . For $h$ heads of dimension $d_{k} = d_{model} / h$ , the cost for each head is $O (n^{2} \cdot d_{k})$ . The total cost across all heads is $h \cdot O (n^{2} \cdot d_{k}) = h \cdot O (n^{2} \cdot d_{model} / h) = O (n^{2} \cdot d_{model})$ . The computation remains on the same order of magnitude while providing a much richer representational capacity.

To see how these mathematical steps translate into code, let's turn to Andrej Karpathy's excellent from-scratch implementation.

Let's build GPT: from scratch, in code, spelled out.

In this segment, Andrej Karpathy transitions from implementing a single self-attention head to implementing the full multi-head attention mechanism.

Watch from this part of the video. Pay close attention to how he uses nn.ModuleList to create multiple heads that run in parallel and torch.cat to concatenate their outputs. This directly mirrors the mathematical formulation we just discussed.

A Deeper Dive into the Math and Implementation

To solidify your understanding, the following article provides a clear, step-by-step mathematical breakdown and a NumPy implementation of multi-head attention. This resource is perfect for bridging the gap between theory and from-scratch coding.

The Math Behind Multi-Head Attention in Transformers

This article explains the mathematical steps for constructing multi-head attention and provides a from-scratch implementation.

First, read section "Multi-Head Attention" for a concise mathematical description. Then, carefully work through the section "Building from Scratch". Focus on how the transpose_qkv and transpose_output methods handle the reshaping of data to manage the multiple heads efficiently.

This implementation detail of reshaping tensors is fundamental to achieving high performance in practice. By packing the heads into an additional dimension of the tensor, all heads can be processed with single, highly optimized matrix multiplication calls, fully leveraging the power of modern GPUs.

Conclusion

In this lesson, we have deconstructed the multi-head attention mechanism. We've established its purpose as a method for allowing a model to learn multiple types of relationships within a sequence simultaneously.

Key Takeaways:

Purpose: Multi-head attention enhances model expressiveness by running multiple self-attention computations in parallel, each focusing on a different "subspace" of the input data.
Architecture: The mechanism involves three key stages: (1) projecting the input into $h$ sets of lower-dimensional Q, K, V representations, (2) computing scaled dot-product attention for each head in parallel, and (3) concatenating the results and passing them through a final linear projection layer.
Computational Efficiency: By reducing the dimension of each head ( $d_{k} = d_{model} / h$ ), the total computation remains comparable to a single, full-dimension attention head, making it a highly efficient way to increase model capacity.

Preview of the Next Lesson:

Multi-head attention is the engine of the Transformer, but it doesn't operate in a vacuum. To build a robust and deep network, we need to integrate it into a larger structure. In the next lesson, we will assemble a complete Transformer encoder block. This involves wrapping our multi-head attention layer with two other crucial components: residual connections and layer normalization. We will explore why these are indispensable for stabilizing training and enabling the construction of very deep Transformer models.