Study the paper

Attention Is All You Need

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

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 allows the model to jointly attend to information from different representation subspaces at different positions. Instead of performing a single attention function with $d_{\text{model}}$ -dimensional queries, keys, and values, we project them $h$ 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}

Each head is computed as an independent scaled dot-product attention: $\mathrm{head}_i = \mathrm{Attention}(Q W_i^Q, K W_i^K, V W_i^V)$ where the projections are parameter matrices $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}}}$ .

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 consider a simplified setup with $h = 2$ heads, $d_v = 2$ , and $d_{\text{model}} = 4$ .

Suppose the outputs of the two heads for a single token are:

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

Concatenation: $\mathrm{Concat}(\mathrm{head}_1, \mathrm{head}_2) = [1.0, 2.0, 3.0, 4.0]$
Linear Projection: Let the output projection matrix $W^O$ be a $4 \times 4$ identity matrix: $W^O = \begin{pmatrix} 1 & 0 & 0 & 0 \\ 0 & 1 & 0 & 0 \\ 0 & 0 & 1 & 0 \\ 0 & 0 & 0 & 1 \end{pmatrix}$
Final Output: $\mathrm{MultiHead}(Q, K, V) = [1.0, 2.0, 3.0, 4.0] \times W^O = [1.0, 2.0, 3.0, 4.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 containing representation vectors · Matrix of shape (N, d_model)
$K$: Key matrix containing representation vectors · Matrix of shape (M, d_model)
$V$: Value matrix containing representation vectors · Matrix of shape (M, d_model)
$head_i$: The output of the i-th attention head · Matrix of shape (N, d_v)
$h$: Number of attention heads · Scalar integer
$W^O$: Learned output projection parameter matrix · Matrix of shape (h * d_v, d_model)
$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)