Study the paper

Attention Is All You Need

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

Multi-Head Attention Mechanism

Multi-Head Attention

Source equation

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

Multi-Head Attention allows the model to jointly attend to information from different representation subspaces at different positions. Instead of performing a single attention function with queries, keys, and values, the queries, keys, and values are projected multiple times with different, learned linear projections.

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}

The output of each individual attention head $\mathrm{head}_i$ is concatenated along the feature dimension, and then projected using a learned parameter matrix $W^O$ to produce the final multi-head attention output.

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}

Implementation detail

Illustrative Toy Calculation

Let us compute the final output of a multi-head attention layer with $h = 2$ heads, where each head produces a 2-dimensional output for a single token, and the final output projection matrix $W^O$ maps the concatenated representation back to a 2-dimensional space.

Inputs:
- $\mathrm{head}_1 = [1.0, 2.0]$
- $\mathrm{head}_2 = [3.0, 4.0]$
- $W^O = \begin{bmatrix} 0.5 & 1.0 \\ 0.0 & 0.5 \\ 1.0 & 0.0 \\ 0.0 & 1.0 \end{bmatrix}$
Step 1: Concatenation $\mathrm{Concat}(\mathrm{head}_1, \mathrm{head}_2) = [1.0, 2.0, 3.0, 4.0]$
Step 2: Linear Projection $\mathrm{MultiHead}(Q,K,V) = [1.0, 2.0, 3.0, 4.0] \times \begin{bmatrix} 0.5 & 1.0 \\ 0.0 & 0.5 \\ 1.0 & 0.0 \\ 0.0 & 1.0 \end{bmatrix}$ $\text{Output}_1 = 1.0(0.5) + 2.0(0.0) + 3.0(1.0) + 4.0(0.0) = 0.5 + 0.0 + 3.0 + 0.0 = 3.5$ $\text{Output}_2 = 1.0(1.0) + 2.0(0.5) + 3.0(0.0) + 4.0(1.0) = 1.0 + 1.0 + 0.0 + 4.0 = 6.0$ $\mathrm{MultiHead}(Q,K,V) = [3.5, 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 representing the queries · Matrix of shape (N, d_k)
$K$: Key matrix representing the keys · Matrix of shape (M, d_k)
$V$: Value matrix representing the values · Matrix of shape (M, d_v)
$head_i$: The output of the i-th attention head · Matrix of shape (N, d_v)
$h$: Number of attention heads · Scalar integer
$Concat$: Concatenation operation along the feature dimension · Function mapping h matrices of shape (N, d_v) to a single matrix of shape (N, h * d_v)
$W^O$: Learned output projection parameter matrix · Matrix of shape (h * d_v, d_model)
$MultiHead(Q,K,V)$: The final multi-head attention output · Matrix of shape (N, d_model)