Multi-Head Attention: Diverse Features, Diverse Subspaces

Introduction

In our last lesson, we focused on the mechanics and theoretical implications of self-attention. We established that by allowing every token in a sequence to attend to every other token, we create a direct computational path of constant length $O(1)$, overcoming the long-range dependency limitations of recurrent models. However, we concluded by posing a question: what if a single attention mechanism is insufficient to capture the rich and varied relationships within a sequence? A sentence's meaning is often constructed from multiple layers of structure—syntactic, semantic, and positional, among others. Forcing a single attention head to average its focus across all these different aspects could be a significant bottleneck.

This lesson directly addresses that challenge. Our goal is to derive multi-head attention and explain its role in capturing diverse feature representations from different subspaces. We will deconstruct this mechanism, which is a cornerstone of the Transformer architecture, to understand how it allows the model to process information from different representational subspaces in parallel. This enables the model to simultaneously attend to various types of relationships, leading to a much richer and more nuanced understanding of the input.

Motivation: Why a Single Attention Head is Not Enough

Language is inherently complex and often ambiguous. A single attention distribution per token may not have the capacity to capture multiple, equally valid interpretations or relationship types. Consider the sentence, "She saw the man with the telescope." This sentence has at least two valid interpretations:

1. She used a telescope to see the man. (A strong relationship between "she" and "telescope").
2. She saw a man who was holding a telescope. (A strong relationship between "man" and "telescope").

A single-headed attention mechanism would be forced to choose one dominant attention pattern or average them, potentially losing the nuance of the ambiguity. To truly understand language, a model needs to be able to identify different types of dependencies in parallel.

The following video provides an excellent intuitive explanation for this need, using linguistic examples to illustrate why a single perspective on a sentence is insufficient.

The Multi-Head Attention Architecture

The solution proposed in the original Transformer paper is not to design a more complex, single attention mechanism, but to run the simple scaled dot-product attention mechanism multiple times in parallel. This is the essence of Multi-Head Attention.

The core idea is to have multiple "heads," where each head is a separate scaled dot-product attention unit. Before being fed into its respective head, the input is projected into a lower-dimensional representation space. Each head, therefore, learns a different projection and can focus on a different aspect of the relationships between tokens.

Deriving Multi-Head Attention Mathematically

Let's formally derive the mechanism. We begin with our input Query ($Q$), Key ($K$), and Value ($V$) matrices, each of dimension $n\times d_{\text{model}}$, where $n$ is the sequence length and $d_{\text{model}}$ is the model's embedding dimension.

1. Linear Projections for Each Head: Instead of a single set of weight matrices, we now introduce $h$ different sets of learnable weight matrices, one for each head: $W_i^Q$, $W_i^K$, $W_i^V$ for $i=1,\dots ,h$. A key detail is that these projections map the $d_{\text{model}}$-dimensional embeddings into a lower-dimensional space of size $d_k=d_v=d_{\text{model}}/h$.

   For each head $i$, we compute a distinct set of queries, keys, and values:

   $$\begin{gathered} Q_i=Q{W}_i^Q \\ {K}_i=K{W}_i^K \\ {V}_i=V{W}_i^V \end{gathered}$$

   Here, $W_i^Q,W_i^K\in {\mathbb{R}}^{d_{\text{model}}\times d_k}$ and $W_i^V\in {\mathbb{R}}^{d_{\text{model}}\times d_v}$. This projects the original Q, K, and V matrices into $h$ different representational subspaces.

2. Parallel Attention Calculations: Next, we apply scaled dot-product attention to each of these projected sets of queries, keys, and values in parallel.

   $${\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$$

   Each ${\text{head}}_i$ is a matrix of size $n\times d_v$, containing the attention-weighted value vectors for that head's specific subspace.

3. Concatenation and Final Projection: The outputs from all $h$ heads are concatenated along the feature dimension to form a single large matrix of size $n\times (h\cdot d_v)$. Since $h\cdot d_v=h\cdot (d_{\text{model}}/h)=d_{\text{model}}$, this brings us back to our original embedding dimension.

   $$\text{Concat}({\text{head}}_1,{\text{head}}_2,\dots ,{\text{head}}_h)$$

   Finally, this concatenated matrix is passed through one last linear projection, using a learnable weight matrix $W^O\in {\mathbb{R}}^{d_{\text{model}}\times d_{\text{model}}}$, to produce the final output of the multi-head attention layer. This final projection allows the model to mix the information learned from the different heads.

The complete formula is:

where ${\text{head}}_i=\text{Attention}(Q{W}_i^Q,K{W}_i^K,V{W}_i^V)$.

For a detailed walkthrough of this derivation with concrete numbers and matrix dimensions, the following video is exceptionally clear.

The Role: Capturing Diverse Feature Representations

The mathematical derivation shows how multi-head attention works, but the crucial question is why it works. The mechanism's power lies in its ability to let the model jointly attend to information from different representational subspaces at different positions.

• Subspace Projections: Each set of projection matrices $(W_i^Q,W_i^K,W_i^V)$ for head $i$ learns to project the original embeddings into a different subspace. In one subspace, the notion of "similarity" (as computed by the $Q{K}^T$ dot product) might be based on syntactic roles. In another, it might be based on semantic meaning or even simple positional proximity.
• Parallel Specialization: Because the heads are independent during the attention calculation, each can specialize. This is analogous to convolutional filters in a CNN, where different filters learn to detect different features (e.g., vertical edges, horizontal edges, specific color patterns). Similarly, attention heads can learn to detect different types of linguistic patterns.

This is not just a theoretical claim. Empirical analysis of trained Transformer models reveals this specialization in practice. Different heads demonstrably learn to focus on distinct and interpretable linguistic phenomena.

This specialization allows the model to build a much richer, multi-faceted representation of each token. The final linear projection with $W^O$ then acts as a learned aggregator, deciding how to best combine these different views to produce the final output embedding for the layer.

Computational Considerations

From a computer science perspective, it is worth noting that despite the apparent increase in complexity, multi-head attention is computationally efficient. The total computation is similar to that of a single-head attention with the full $d_{\text{model}}$ dimensionality. However, because the dimension of each head is reduced ($d_k=d_{\text{model}}/h$), and since the computations for each head are independent, the entire process is highly parallelizable on modern hardware like GPUs.

The article below provides further insight into the computational aspects and the intuition behind the independent heads.

Conclusion

In this lesson, we have moved beyond the single-head attention model to understand the more powerful and widely used multi-head architecture.

Here are the key takeaways:

• Motivation: A single attention head can be a bottleneck, struggling to capture the multiple types of relationships (e.g., syntactic, semantic) present in language.
• Derivation: Multi-head attention addresses this by running $h$ scaled dot-product attention mechanisms in parallel. It does so by first projecting the Q, K, and V matrices into $h$ different, lower-dimensional subspaces. The outputs are then concatenated and passed through a final linear projection.
• Role: This architecture allows different heads to specialize in learning different types of features or relationships. This is not just a theoretical benefit but has been empirically observed, with heads learning to track distinct linguistic phenomena like noun-modifier or verb-object dependencies.
• Efficiency: The process is computationally efficient and highly parallelizable, making it well-suited for modern deep learning hardware.

You now understand the core computational block that makes up the Transformer. We have covered scaled dot-product attention, self-attention, and now multi-head attention. In our next lesson, we will put all these pieces together. We will examine the full Transformer architecture, incorporating positional encodings, feed-forward networks, and residual connections to see how these attention blocks are stacked to form a powerful encoder-decoder model.