Multi-head attention: Worked examples — AI-ML deep dive

This page is the drill floor for Multi-head attention. The parent note told you the machinery; here we run it against every kind of input the machine can meet — small dimensions, single tokens, one head vs. many, degenerate scores, cross-attention shapes, and an exam twist. Nothing is skipped, so you never hit a case you didn't already practise.

Before we start, one tiny reminder of the two ideas we lean on constantly (both come from Query-Key-Value Intuition and Self-Attention):

Recall What the softmax and the score actually mean

A score $s_{ij}$ measures "how much should token $i$ listen to token $j$ ?" It is a plain number (can be negative). The softmax turns a whole row of scores into positive weights that add up to $1$ — like slicing a pie so token $i$ decides what fraction of its attention each other token gets. $weight_{ij} = \frac{e ^{s_{ij}}}{\sum _{j^{'}} e ^{s_{i j^{'}}}}$

Every symbol below: $n$ = number of tokens (rows), $d_{model}$ = width of each token's vector, $h$ = number of heads, $d_{k} = d_{model} / h$ = width of the query/key space inside one head, and $d_{v}$ = width of the value space inside one head (the length of what each head actually outputs). In the standard transformer $d_{v} = d_{k} = d_{model} / h$ , and we use that throughout.

The scenario matrix

The cases below are chosen so that between them they hit every corner of multi-head attention: the tiniest dimensions, the boundary cases where a piece of the formula does nothing, the limiting values of the softmax, and a real-world / exam flavour.

#	Cell class	What makes it special	Example
A	Single head, tiny dims	$h = 1$ , $d_{k} = 2$ — strip away parallelism to see the core	Ex 1
B	Two heads, splitting	$h = 2$ — watch $d_{model}$ split and re-concatenate	Ex 2
C	Degenerate scores: all equal	softmax outputs uniform weights	Ex 3
D	Limiting scores: one huge	softmax → one-hot (a hard pick)	Ex 4
E	Zero / negative scores	shows scores may be negative, weights never are	Ex 5
F	Scaling matters	compare with vs. without $d_{k}$	Ex 6
G	Real-world word problem	pronoun co-reference, one head specialises	Ex 7
H	Cross-attention shapes	query length $\neq =$ key length (Cross-Attention)	Ex 8
I	Exam twist: count parameters	"how many learnable numbers?"	Ex 9

Example A — Single head, the bare core

Scores $Q_{1} K_{1}^{⊤}$ . Why this step? The dot product of query $i$ with key $j$ is how well they align — the raw "listen to $j$ " number. $Q_{1} K_{1}^{⊤} = [1001]$ Diagonal is $1$ (self-match), off-diagonal $0$ (orthogonal, no match).
Scale by $d_{k} = 2$ . Why? Large dot products push softmax into extreme regions where gradients vanish; dividing keeps them tame. Scores become $\frac{1}{2} \approx 0.707$ on the diagonal, $0$ off it.
Softmax each row. Why? Turn scores into a pie of attention that sums to $1$ . Row 1 = $softmax (0.707, 0)$ : $w_{11} = \frac{e ^{0.707}}{e ^{0.707} + e ^{0}} \approx 0.6698, w_{12} \approx 0.3302$ By symmetry row 2 = $(0.3302, 0.6698)$ .
Weighted sum of values. Why? The output of token $i$ is the pie-weighted blend of the value rows. $head_{1} [1] = 0.6698 (2, 0) + 0.3302 (0, 4) = (1.3396, 1.3208)$ $head_{1} [2] = 0.3302 (2, 0) + 0.6698 (0, 4) = (0.6604, 2.6792)$

Verify: each output row is a convex mix of $(2, 0)$ and $(0, 4)$ , so both coordinates lie inside $[0, 4]$ ✓. Row 1 leans toward $(2, 0)$ , row 2 toward $(0, 4)$ — exactly the "close to $V$ but blurred" forecast ✓.

Example B — Two heads split then re-join

Forecast: with a single token, attention can't gather anything from neighbours, so concatenation just rebuilds a 4-vector; with $W^{O} = I$ nothing changes → output $= (1, 2, 3, 4)$ .

Single-token softmax. Why? With $n = 1$ each row has one score; $softmax (s) = \frac{e ^{s}}{e ^{s}} = 1$ . Attention is forced to be the identity here — a real degenerate case, so each head just outputs its own value vector.
Read off each head's output. Why? Because attention is the identity, $head_{1} = v_{1} = (1, 2)$ and $head_{2} = v_{2} = (3, 4)$ — exactly the slices the select-matrices picked.
Concatenate heads. Why? We stack the two 2-D outputs side by side to rebuild $d_{model} = 4$ : $[(1, 2) ∣ (3, 4)] = (1, 2, 3, 4)$ .
Apply $W^{O} = I$ . Why? $W^{O}$ mixes heads; the identity is the "do-nothing" mix, useful to isolate the concat step. Output $= (1, 2, 3, 4)$ .

Verify: dimension check $h \cdot d_{v} = 2 \cdot 2 = 4 = d_{model}$ ✓. Single token means self-attention is an identity, so the value passes through unchanged ✓.

Example C — Degenerate scores: everything equal

Forecast: identical scores mean "no token is preferred" — attention must be uniform ( $\frac{1}{3}$ each), so the output is the plain average $= 20$ .

Softmax of equal scores. Why? When all $s_{j}$ match, $e^{s_{j}}$ are equal and cancel: $w_{j} = \frac{e ^{2}}{3 e ^{2}} = \frac{1}{3}$ . This is the maximum-entropy, "I don't know who matters" case.
Weighted sum. Why? Blend the values: $\frac{1}{3} (10) + \frac{1}{3} (20) + \frac{1}{3} (30) = 20$ .

Verify: weights sum to $1$ ( $3 \times \frac{1}{3}$ ) ✓. Output $20$ is the arithmetic mean, which is what uniform attention must give ✓.

Example D — Limiting scores: one huge value

Forecast: the third score dwarfs the others, so softmax should collapse to "almost all attention on token 3" → output near $30$ .

Softmax. Why? Exponentials magnify gaps. $e^{0} = 1, e^{0} = 1, e^{20} \approx 4.85 \times 1 0^{8}$ . $w_{3} = \frac{e ^{20}}{2 + e ^{20}} \approx 0.99999999588$
Weighted sum. Why? Blend: $\approx 0.0000000021 (10) + 0.0000000021 (20) + 0.99999999588 (30) \approx 30.000000$ .

Verify: as the top score $\to \infty$ , softmax $\to (0, 0, 1)$ (a hard "pick token 3"), giving exactly $v_{3} = 30$ . Our number $29.99999994$ rounds to $30$ ✓. This is the hard-attention limit.

The figure below plots exactly this collapse: as the gap between the top score and the others grows, the weight on the dominant token climbs toward $1$ . The magenta dot marks Example D's gap of $20$ , already glued to the ceiling.

Example E — Zero and negative scores

Forecast: token 2's score is higher, so it wins most attention; despite the negative score, both weights stay positive → output between 100 and 200, nearer 200.

Softmax with a negative score. Why? $e^{- 3} \approx 0.0498$ is small but positive — that's the whole point of the exponential: it maps any real number (even $- 3$ ) to a positive weight. $w_{1} = \frac{e ^{- 3}}{e ^{- 3} + e ^{1}} \approx 0.01799, w_{2} \approx 0.98201$
Weighted sum. Why? Blend: $0.01799 (100) + 0.98201 (200) = 198.20$ .

Verify: weights positive and sum to $1$ ✓. Result $198.20 \in (100, 200)$ , close to $200$ as forecast ✓. A negative score just means "listen less," never "listen a negative amount."

Example F — Why the $d_{k}$ scaling exists

Forecast: unscaled the gap $32$ is huge → softmax nearly one-hot on token 1 (very "peaky", gradients die). Scaled, the gap shrinks to $4$ → softer, healthier weights.

Unscaled softmax. Why? Show the pathology. $w_{1} = \frac{e ^{40}}{e ^{40} + e ^{8}}$ . The gap $32$ gives $w_{1} \approx 1 - e^{- 32} \approx 0.99999999999999$ .
Scaled softmax. Why? Divide scores by $8$ : $(5, 1)$ , gap $4$ . $w_{1} = \frac{e ^{5}}{e ^{5} + e ^{1}} \approx 0.9820$ .

Verify: scaled $w_{1} \approx 0.982$ leaves real weight ( $0.018$ ) on token 2, so gradients flow to it; unscaled leaves essentially $0$ , freezing learning ✓. This is exactly why the parent's formula divides by $d_{k}$ .

Example G — Real-world: pronoun co-reference

Forecast: "cat" has by far the top score, so this head should place most of "it"'s attention on "cat" — that's the head doing its co-reference job.

Softmax over the 6 scores. Why? Convert this head's opinions into a probability of who "it" refers to. $denominator = e^{3.0} + e^{0.5} + e^{0.2} + e^{1.0} + e^{0.3} + e^{0.6}$ $w_{cat} = \frac{e ^{3.0}}{denominator} \approx 0.6963$
Interpret. Why? $\approx 70%$ of "it"'s attention lands on "cat" — the head has learned co-reference. Other heads meanwhile handle syntax or position (that's the Self-Attention division of labour the parent describes).

Verify: all six weights sum to $1$ ✓; "cat" is the argmax and dominates ✓. This is what Attention Visualization would show as a bright line from "it" → "cat."

Example H — Cross-attention: different lengths

Forecast: scores form a "who-listens-to-whom" grid with one row per query token and one column per key token → $3 \times 5$ . Output has one row per query token → $3 \times 64$ .

Project. Why? Query comes from the decoder, key/value from the encoder: $Q_{i} \in R^{3 \times 64}$ , $K_{i}, V_{i} \in R^{5 \times 64}$ .
Scores $Q_{i} K_{i}^{⊤}$ . Why? Each of the 3 query rows dots against each of the 5 key rows → shape $3 \times 5$ . (Not square — that's the hallmark of cross-attention.)
Softmax over the key axis (the 5 columns) → still $3 \times 5$ , each row sums to 1. Why? Each target token distributes its attention over the 5 source tokens.
Times $V_{i}$ ( $3 \times 5$ times $5 \times 64$ ) $\to$ $head_{i} \in R^{3 \times 64}$ . Why? Inner dimension $5$ cancels; output length follows the query side.

Verify: shape algebra $(3 \times 5) \cdot (5 \times 64) = (3 \times 64)$ ✓, and $8 \cdot 64 = 512 = d_{model}$ after concat ✓. In self-attention $n_{q} = n_{k}$ makes this grid square; cross-attention is the general case.

Example I — Exam twist: count the parameters

Forecast: the four big matrices are each effectively $512 \times 512$ , so the total should be about $4 \times 51 2^{2} \approx 1 0^{6}$ .

Per-head projections. Why? Each $W_{i}^{Q} \in R^{512 \times 64}$ has $512 \cdot 64 = 32768$ entries. Same for $W_{i}^{K}, W_{i}^{V}$ . Across $8$ heads: $3 \times 8 \times (512 \times 64) = 3 \times 8 \times 32768 = 786432$ Note $8 \times 64 = 512$ , so all heads' $W^{Q}$ together equal one $512 \times 512$ matrix — the "stacked projection" trick.
Output projection $W^{O}$ . Why? $W^{O} \in R^{512 \times 512}$ adds $262144$ weights.
Total. Why? Sum: $786432 + 262144 = 1048576$ .

Verify: $1048576 = 4 \times 51 2^{2} = 2^{20}$ exactly ✓ — a clean " $4 d_{model}^{2}$ " rule of thumb you can quote in an exam. (Compare Linear Projections in Neural Networks for where these matrices live.)

Recall

Reveal after guessing:

Equal scores for all tokens produce what attention weights?

Uniform,

\frac{1}{n}

each — the plain average of the values.

As one score

\to + \infty

while others stay fixed, softmax approaches what?

A one-hot vector — hard attention onto the single dominant token.

Can a softmax weight ever be negative?

No —

e^{s} > 0

for every real

s

, so all weights are strictly positive and sum to

1

In cross-attention with 3 query tokens and 5 key tokens, the score matrix shape is?

3 \times 5

(rows = queries, columns = keys).

Total learnable weights of one MHA block with

d_{model} = 512

(no biases)?

4 \times 51 2^{2} = 1048576

Why divide scores by

d_{k}

To stop large dot products from making softmax nearly one-hot, which would kill gradients.

See also: Transformer Encoder Layer, Positional Encoding, and the Hinglish walkthrough.