Benefits of Multiple Attention Heads

Multi-Head Attention (MHA) operates by running Scaled Dot-Product Attention multiple times in parallel. This design offers notable advantages compared to using a single, larger attention mechanism. The benefits arise primarily from MHA's ability to allow the model to process information from different perspectives simultaneously.

Focusing on Different Representation Subspaces

The core idea behind MHA is that projecting the original Query ($Q$), Key ($K$), and Value ($V$) matrices into lower-dimensional spaces for each head ($h$) allows each head to potentially learn different types of relationships or focus on different aspects of the sequence.

Recall that for each head $i$, we compute:

$\text{head}_i = \text{Attention}(QW^Q_i, KW^K_i, VW^V_i)$

Here, $W^Q_i$, $W^K_i$, and $W^V_i$ are distinct, learned weight matrices for head $i$. These matrices project the full-dimensional $Q, K, V$ vectors into subspaces specific to that head. Because these projection matrices are learned independently during training, each head can specialize.

For example, consider the sentence: "The tired cat slept on the warm mat." One attention head might learn projections that help it focus on syntactic dependencies, perhaps linking "cat" (subject) to "slept" (verb). Another head might capture longer-range semantic relationships, linking "tired" to "slept". A third could focus on positional proximity, strongly weighting adjacent words.

Parallel attention heads process input embeddings projected into different subspaces, capturing diverse relationships before results are combined.

This ability to attend to different "representation subspaces" means the model isn't forced to average conflicting attention needs into a single set of weights. Instead, it can leverage specialized heads for different tasks.

Integrating Diverse Information

After each head computes its attention output, the results are concatenated:

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^O$

This concatenation combines the outputs from all the specialized heads. The final linear projection, learned by the weight matrix $W^O$, then integrates this rich, multi-faceted information into a single output tensor. This projection layer learns how to best combine the insights gathered from the different heads.

Enhanced Model Expressiveness

Using multiple heads, each operating on a lower-dimensional projection, doesn't necessarily increase the total computation significantly compared to a single attention head with the full dimensionality (assuming the dimensions are chosen appropriately, e.g., $d_{k,head} = d_k / h$). However, it provides a more expressive way to capture complex dependencies within the sequence. It enriches the model's capacity to understand relationships between words, leading to better performance on downstream tasks.

In essence, Multi-Head Attention provides a mechanism for the model to look at the input sequence from multiple viewpoints simultaneously, aggregate these views, and make more informed decisions about how different parts of the sequence relate to each other. This parallel, multi-perspective processing is a significant factor contributing to the success of Transformer models.