Study the paper

Attention Is All You Need

Lessons, visuals, quizzes, flashcards, and resources—organized in teaching order.

Multi-Head Attention Mechanics

Multi-Head Attention

Source equation

$\displaystyle\mathrm{MultiHead}(Q,K,V)=\mathrm{Concat}(\mathrm{head_{1}},...,\mathrm{head_{h}})W^{O}$

The Multi-Head Attention mechanism projects queries, keys, and values into multiple lower-dimensional subspaces, performs attention on each subspace in parallel, and then concatenates and projects the results back to the original dimension.

Sources

equation

MultiHead(Q,K,V)\displaystyle\mathrm{MultiHead}(Q,K,V) =Concat(head1,…,headh)WO\displaystyle=\mathrm{Concat}(\mathrm{head_{1}},...,\mathrm{head_{h}})W^{O}

\displaystyle\mathrm{MultiHead}(Q,K,V)=\mathrm{Concat}(\mathrm{head_{1}},...,\mathrm{head_{h}})W^{O}

$\displaystyle\mathrm{MultiHead}(Q,K,V)=\mathrm{Concat}(\mathrm{head_{1}},...,\mathrm{head_{h}})W^{O}$

Sources

equation

MultiHead(Q,K,V)\displaystyle\mathrm{MultiHead}(Q,K,V) =Concat(head1,…,headh)WO\displaystyle=\mathrm{Concat}(\mathrm{head_{1}},...,\mathrm{head_{h}})W^{O}

\displaystyle\mathrm{MultiHead}(Q,K,V)=\mathrm{Concat}(\mathrm{head_{1}},...,\mathrm{head_{h}})W^{O}

Illustrative Toy Calculation

Let us consider a simplified scenario with $h = 2$ heads, where each head output is a 1D vector of size 2:

$\text{head}_1 = [1.0, 2.0]$
$\text{head}_2 = [3.0, 4.0]$

Concatenation: $\text{Concat}(\text{head}_1, \text{head}_2) = [1.0, 2.0, 3.0, 4.0]$
Linear Projection: Let the output projection matrix $W^O$ be a $4 \times 2$ matrix: $W^O = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 0 \\ 0 & 1 \end{bmatrix}$
Result: $\text{MultiHead}(Q, K, V) = [1.0, 2.0, 3.0, 4.0] \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 0 \\ 0 & 1 \end{bmatrix} = [1.0 \cdot 1 + 3.0 \cdot 1, \; 2.0 \cdot 1 + 4.0 \cdot 1] = [4.0, 6.0]$

Sources

equation

MultiHead(Q,K,V)\displaystyle\mathrm{MultiHead}(Q,K,V) =Concat(head1,…,headh)WO\displaystyle=\mathrm{Concat}(\mathrm{head_{1}},...,\mathrm{head_{h}})W^{O}

\displaystyle\mathrm{MultiHead}(Q,K,V)=\mathrm{Concat}(\mathrm{head_{1}},...,\mathrm{head_{h}})W^{O}

$Q$: Query matrix · (N, d_model)
$K$: Key matrix · (M, d_model)
$V$: Value matrix · (M, d_model)
$head_i$: Output of the i-th attention head · (N, d_v)
$h$: Number of attention heads · scalar
$W^O$: Output projection parameter matrix · (h * d_v, d_model)
$MultiHead(Q, K, V)$: Final multi-head attention output · (N, d_model)