Explaining Multi-Head Attention (MHA)

Benjamin

3 days ago

To understand Multi-Head Attention (MHA), we first need to take a step back and examine the limitation of applying the attention equation (Scaled Dot-Product Attention) only once.

The Problem with a Single “Head”

Human language is complex and possesses multiple layers of meaning that occur simultaneously. Consider the sentence:

“The student delivered the book to the library yesterday.”

To understand the meaning of the word “delivered”, we need to capture several syntactic and semantic relationships at the same time:

Subject: Who delivered? (The student)
Direct Object: What was delivered? (The book)
Location: Where was it delivered? (To the library)
Time: When was it delivered? (Yesterday)

If a model uses a single self-attention mechanism (a single “head”), it has to squeeze all these distinct relationships into a single probability distribution. If it allocates most of its attention probability to connecting “delivered” to “student”, it will struggle to pay adequate attention to “book”, “library”, and “yesterday” in the same step.

The Solution: Multi-Head Attention

Multi-Head Attention resolves this by applying a “divide and conquer” strategy.

Instead of running a single attention mechanism on full-dimensional vectors, MHA splits the model’s representations into multiple smaller subspaces, allowing several attention mechanisms (“heads”) to run in parallel.

Imagine the original model has a total dimension of dmodel=512d_{\text{model}} = 512. If the model uses h = 8 heads:

It projects Q, K, and V into 8 smaller versions (dimension dk=512/8=64d_k = 512 / 8 = 64).
Each of the 8 heads runs its own version of the attention equation independently.

This allows each head to specialize in a different facet of language:

Head 1: Focuses on subject-verb agreement (subject-verb).
Head 2: Focuses on objects of the action (verb-object).
Head 3: Focuses on time and location information.
Head 4: Focuses on pronoun references (anaphora).

The Mathematics of Multi-Head Attention

The mathematics of MHA expands the original attention equation by integrating linear projections, parallel computation, and the merging of results.

A. Linear Projections

Before calculating attention, each head $i$ receives its own projected version of Queries, Keys, and Values through trainable weight matrices (WiQ,WiK,WiVW_i^Q, W_i^K, W_i^V):

Qi=Q⋅WiQQ_i = Q \cdot W_i^Q

Ki=K⋅WiKK_i = K \cdot W_i^K

Vi=V⋅WiVV_i = V \cdot W_i^V

B. Parallel Attention Computation

Each head computes its scaled dot-product attention independently:

headi=Attention(Qi,Ki,Vi)=softmax(QiKiTdk)Vi\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

C. Concatenation and Final Projection

At the end, the outputs from all h heads are concatenated (placed side-by-side) to form a full-sized vector again. Finally, this concatenated vector is projected back to the original model dimension (dmodeld_{\text{model}}) using a final trainable weight matrix (WOW^O):

MultiHead(Q,K,V)=Concat(head1,head2,…,headh)⋅WO\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, \dots, \text{head}_h) \cdot W^O

In short, Multi-Head Attention gives Large Language Models their deep capacity for interpretation by enabling them to process multiple nuances of human language simultaneously.