Module 9 — Multi-Head Attention

Multi-Head Attention (MHA) is an extension of the Self-Attention mechanism introduced in the original Transformer architecture. Instead of computing attention once, the model computes multiple attention operations in parallel using different learned projections.

Each attention head focuses on different relationships within the input sequence, enabling the model to learn richer contextual representations.

Topics

Multiple Attention Heads
Head Projection
Concatenation
Output Projection

1. Input Embedding

The input embedding matrix is

X \in \mathbb{R}^{n\times d_{model}}

where

$n$ = Sequence length
$d_{model}$ = Embedding dimension

2. Query Projection

Each attention head learns its own Query projection matrix.

For head $i$

Q_i = XW_i^Q

where

W_i^Q \in \mathbb{R}^{d_{model}\times d_k}

3. Key Projection

Each head has its own Key projection.

K_i = XW_i^K

where

W_i^K \in \mathbb{R}^{d_{model}\times d_k}

4. Value Projection

Each head has its own Value projection.

V_i = XW_i^V

where

W_i^V \in \mathbb{R}^{d_{model}\times d_v}

5. Attention Head

Each head computes Scaled Dot-Product Attention independently.

head_i = Attention(Q_i,K_i,V_i)

Expanding the attention function,

head_i = Softmax \left( \frac {Q_iK_i^T} {\sqrt{d_k}} \right) V_i

6. Multiple Attention Heads

Assume there are

h

attention heads.

The outputs are

head_1, head_2, \ldots, head_h

Each head learns different semantic relationships between tokens.

7. Concatenation

The outputs of all heads are concatenated along the feature dimension.

H = Concat (head_1, head_2, \ldots, head_h)

Output Shape

H \in \mathbb{R}^{n\times hd_v}

where

$h$ = Number of heads
$d_v$ = Value dimension

8. Output Projection

The concatenated output is projected back to the model dimension.

MultiHead = Concat(head_1,\ldots,head_h) W_O

where

W_O \in \mathbb{R}^{hd_v\times d_{model}}

Output Shape

MultiHead \in \mathbb{R}^{n\times d_{model}}

9. Complete Multi-Head Attention Formula

The complete Multi-Head Attention equation is

MultiHead(Q,K,V) = Concat (head_1, \ldots, head_h) W_O

where

head_i = Softmax \left( \frac {Q_iK_i^T} {\sqrt{d_k}} \right) V_i

10. Matrix Dimensions

Suppose

Sequence Length = $n$
Model Dimension = $d_{model}$
Number of Heads = $h$
Head Dimension = $d_{k} = d_{v} = \frac{d_{m o d e l}}{h}$

Then

Input

X \in \mathbb{R}^{n\times d_{model}}

Each Query

Q_i \in \mathbb{R}^{n\times d_k}

Each Key

K_i \in \mathbb{R}^{n\times d_k}

Each Value

V_i \in \mathbb{R}^{n\times d_v}

Each Attention Output

head_i \in \mathbb{R}^{n\times d_v}

Concatenated Output

Concat \in \mathbb{R}^{n\times hd_v}

Final Output

MultiHead \in \mathbb{R}^{n\times d_{model}}

11. Why Multi-Head Attention?

Multiple heads allow the Transformer to learn different relationships simultaneously.

Different heads can capture

Local context
Long-range dependencies
Syntax
Semantics
Entity relationships
Sentence structure

This improves representation learning compared to a single attention head.

Applications

Multi-Head Attention is the core component of

Transformer
BERT
RoBERTa
ALBERT
GPT
GPT-2
GPT-3
GPT-4
T5
ViT
LLaMA
Gemma
Qwen
DeepSeek
Mistral
Phi
Modern Large Language Models (LLMs)

Summary

Concept	Formula
Query Projection	$Q_i=XW_i^Q$
Key Projection