Multi-Head Attention › Worked Example: Tracing Through Two Heads

Before turning to code, let's trace through a minimal example by hand to make the mechanics concrete. Suppose we have a sequence of three tokens: "cat chases mouse." We represent each token with a 4-dimensional embedding (a tiny model dimension chosen so the numbers stay readable), and we use $h = 2$ heads, giving each head $d_k = 2$ dimensions.

Our three token embeddings (rows of $\mathbf{X}$) are:

$$\mathbf{X} = \begin{bmatrix} 1 & 0 & 1 & 0 \\ 0 & 1 & 0 & 1 \\ 1 & 1 & 0 & 0 \end{bmatrix}$$

where:

• Row 1 ($[1, 0, 1, 0]$): the embedding for "cat"
• Row 2 ($[0, 1, 0, 1]$): the embedding for "chases"
• Row 3 ($[1, 1, 0, 0]$): the embedding for "mouse"

For head 1, suppose the learned projection matrices are:

$$\mathbf{W}^Q_1 = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 0 & 0 \\ 0 & 0 \end{bmatrix}, \quad \mathbf{W}^K_1 = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 0 & 0 \\ 0 & 0 \end{bmatrix}, \quad \mathbf{W}^V_1 = \begin{bmatrix} 0 & 0 \\ 0 & 0 \\ 1 & 0 \\ 0 & 1 \end{bmatrix}$$

These projections keep the first two embedding dimensions for queries and keys, and the last two dimensions for values. Computing $\mathbf{Q}_1 = \mathbf{X} \mathbf{W}^Q_1$:

$$\mathbf{Q}_1 = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix}, \quad \mathbf{K}_1 = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix}, \quad \mathbf{V}_1 = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 0 & 0 \end{bmatrix}$$

The raw attention scores for head 1 are $\mathbf{Q}_1 \mathbf{K}_1^T / \sqrt{d_k}$ with $d_k = 2$:

$$\mathbf{Q}_1 \mathbf{K}_1^T = \begin{bmatrix} 1\cdot1+0\cdot0 & 1\cdot0+0\cdot1 & 1\cdot1+0\cdot1 \\ 0\cdot1+1\cdot0 & 0\cdot0+1\cdot1 & 0\cdot1+1\cdot1 \\ 1\cdot1+1\cdot0 & 1\cdot0+1\cdot1 & 1\cdot1+1\cdot1 \end{bmatrix} = \begin{bmatrix} 1 & 0 & 1 \\ 0 & 1 & 1 \\ 1 & 1 & 2 \end{bmatrix}$$

Scaling by $\sqrt{2} \approx 1.414$ gives:

$$\text{Scores}_1 = \frac{1}{\sqrt{2}} \begin{bmatrix} 1 & 0 & 1 \\ 0 & 1 & 1 \\ 1 & 1 & 2 \end{bmatrix} \approx \begin{bmatrix} 0.71 & 0 & 0.71 \\ 0 & 0.71 & 0.71 \\ 0.71 & 0.71 & 1.41 \end{bmatrix}$$

After softmax (row-wise), row 3 ("mouse") ends up attending most strongly to itself because "mouse" has the largest self-similarity score. Row 1 ("cat") splits attention roughly equally between itself and "mouse." Row 2 ("chases") splits between "chases" and "mouse."

Notice that head 2, which uses different projection matrices, will compute a completely different set of scores from the same input $\mathbf{X}$. Even if head 2's projections emphasize different embedding dimensions, it still produces an $n \times d_v$ output that captures a distinct relationship pattern. When you concatenate both head outputs and apply $\mathbf{W}^O$, the final representation for each token can draw on both perspectives simultaneously.

In practice, trained models rarely use identity-like projection matrices. The projections are dense, full-rank matrices learned to maximize task performance. But the worked example illustrates the core mechanism: the same token sequence can be "seen" through multiple lenses, and those lenses are entirely determined by what the model learns to be useful.