How Multi-Head Attention Works

Scaled Dot-Product Attention enables models to prioritize different tokens within a sequence. However, performing this calculation only once can compel the attention mechanism to average diverse types of relationships. For example, consider a sentence such as 'The tired animal didn't cross the street because it was too wide.' A single attention mechanism might struggle to simultaneously capture both the 'tired animal' relationship and the 'street width' relationship effectively when focusing on the word 'it'.

Multi-Head Attention addresses this by running the Scaled Dot-Product Attention process multiple times in parallel, each with different learned transformations of the original queries, keys, and values. This allows each "head" to potentially focus on different aspects or representation subspaces of the information.

Here's the step-by-step process:

1. Linear Projections: Instead of using a single set of Query (Q), Key (K), and Value (V) matrices, Multi-Head Attention first creates $h$ different sets of these matrices, where $h$ is the number of attention heads (a hyperparameter). For each head $i$ (from $1$ to $h$), the original input Q, K, and V matrices (often derived from the same input sequence embeddings in the case of self-attention) are projected using learned weight matrices: $W^Q_i$, $W^K_i$, and $W^V_i$.

   • $Q_i = Q W^Q_i$
   • $K_i = K W^K_i$
   • $V_i = V W^V_i$

   Typically, the dimensions of these projected matrices are smaller than the original embedding dimension ($d_{model}$). If the input embedding dimension is $d_{model}$, each head often works with dimensions $d_k = d_v = d_{model} / h$. This ensures that the total computational cost is similar to a single head attention with full dimensions. These weight matrices ($W^Q_i, W^K_i, W^V_i$) are unique for each head and are learned during the training process.

2. Parallel Attention Calculations: Each of these projected sets ($Q_i, K_i, V_i$) is then fed into its own Scaled Dot-Product Attention mechanism simultaneously. This results in $h$ separate output matrices, let's call them $\text{head}_i$:

   $$\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$ matrix captures attention information based on the specific projections learned by head $i$. Because the projections differ ($W^Q_i, W^K_i, W^V_i$ are different for each $i$), each head can potentially learn to focus on different types of relationships or features within the input sequence.

3. Concatenation: The outputs from all $h$ attention heads are concatenated together along the feature dimension. If each $\text{head}_i$ has dimension $d_v$, the concatenated matrix will have dimension $h \times d_v$. Since we typically set $d_v = d_{model} / h$, the dimension of the concatenated matrix becomes $d_{model}$, matching the original input embedding dimension.

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

4. Final Linear Projection: This concatenated output is then passed through one final linear projection layer, parameterized by another learned weight matrix $W^O$. This projection mixes the information learned by the different heads and produces the final output of the Multi-Head Attention layer, which typically has the dimension $d_{model}$.

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

This entire Multi-Head Attention block can then be used as a component within the larger Transformer architecture, replacing the single Scaled Dot-Product Attention mechanism.

The following diagram illustrates the flow of information through a Multi-Head Attention block with $h$ heads.

This diagram shows how input Q, K, and V matrices are first projected independently for each of the $h$ attention heads. Scaled Dot-Product Attention is then applied to each projected set in parallel. The resulting attention outputs are concatenated and passed through a final linear layer to produce the Multi-Head Attention output.

By allowing different heads to learn different projection matrices ($W^Q_i, W^K_i, W^V_i, W^O$), Multi-Head Attention enables the model to jointly attend to information from different representation subspaces at different positions, leading to a richer and more effective representation compared to using a single attention mechanism.