Why Multi-Head Attention

What is Multi Head Attention?

Single head attention calculates a single weighted average across tokens.

However, a word in a sentence often has multiple different relationships simultaneously:

In the sentence "The cat ate the fish because it was hungry":

1. Relationship A (Grammar): "ate" connects to subject "cat".
2. Relationship B (Pronoun Coreference): "it" connects to "cat".
3. Relationship C (Action Target): "ate" connects to object "fish".

A single attention head averages these competing signals into a single blurred score.

Multi Head Attention projects Queries, Keys, and Values into $h$ independent subspace heads (for example, $h = 8$ or $h = 16$), allowing the network to attend to multiple relationship types in parallel.

  Input Token Embeddings [N x 512]
            │
            ├─► [ Head 1 Projection ] ──► Scaled Dot Product Attention ──► Output 1 [N x 64] (Grammar Head)
            ├─► [ Head 2 Projection ] ──► Scaled Dot Product Attention ──► Output 2 [N x 64] (Pronoun Head)
            │   ...
            └─► [ Head 8 Projection ] ──► Scaled Dot Product Attention ──► Output 8 [N x 64] (Topic Head)
                                                        │
                                                        ▼
                               Concatenate [ Head 1; ...; Head 8 ] ──► [ Output Linear W_O ] ──► Final [N x 512]

Mathematical Formulation

Given model dimension $d_{\text{model}}$ and $h$ heads:

$$d_k = d_v = \frac{d_{\text{model}}}{h}$$

For each head $i \in {1, \dots, h}$:

$$\text{head}_i = \text{Attention}\left( Q W_i^Q, K W_i^K, V W_i^V \right)$$

Concatenate all heads and project:

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

Where $W_i^Q \in \mathbb{R}^{d_{\text{model}} \times d_k}$, $W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$, $W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$, and $W^O \in \mathbb{R}^{h d_v \times d_{\text{model}}}$.

Does Multi Head Attention Cost More FLOPs?

Counter intuitively, No!

Because each head operates on a reduced sub dimension $d_k = d_{\text{model}} / h$:

$$h \times \left( \text{Cost for dimension } \frac{d_{\text{model}}}{h} \right) = \text{Cost for full dimension } d_{\text{model}}$$

Processing 8 parallel heads of dimension 64 takes the exact same computational FLOPs as processing 1 single head of dimension 512.

You get rich multi aspect feature representation learning at zero extra computational cost!

Say this out loud

Multi Head Attention splits Queries, Keys, and Values into h parallel heads operating on reduced sub dimensions d_k = d_model / h. This allows the model to attend to multiple distinct relationship types simultaneously, like grammar, pronouns, and topic associations. Outputs from all heads are concatenated and linearly projected, costing the exact same compute as single head attention.

Followups to expect

1. What is Grouped Query Attention (GQA)? A memory efficient variant used in LLaMA 3 where multiple Query heads share a single Key and Value head group, cutting KV cache memory requirements during inference.
2. What do individual attention heads learn in practice? Probing studies show specific heads specialize in distinct jobs: positional lookup (prev/next token), syntactic parsing, coreference resolution, and Induction Heads for in context pattern repetition.