Multi-head attention: Exercises — AI-ML deep dive

These problems build from "just read the definition" up to "design and reason about the whole block." Every problem has a full solution hidden in a collapsible callout — try first, then reveal. Symbols used here all come from the parent note; if you meet one you forgot, that's your cue to peek back.

Prerequisite ideas live in Self-Attention, Query-Key-Value Intuition, and Linear Projections in Neural Networks.

Level 1 — Recognition

Problem 1.1

State the formula for the per-head dimension $d_{k}$ in terms of $d_{model}$ (the input width per token) and $h$ (the number of heads). Then compute $d_{k}$ for $d_{model} = 768$ , $h = 12$ .

Recall Solution

The parent note defines $d_{k} = d_{v} = \frac{d _{model}}{h} .$ The word "dimension" here just means how many numbers describe one token inside a head. We split the full width $d_{model}$ evenly across the $h$ heads. $d_{k} = \frac{768}{12} = 64.$

Problem 1.2

A multi-head attention block has input $Q, K, V \in R^{n \times d_{model}}$ . Here $n$ is the number of tokens (rows) and $d_{model}$ the width (columns). What is the shape of the final output of the whole block, and why is it the same as the input width?

Recall Solution

Output shape is $R^{n \times d_{model}}$ — same width as the input. Why: concatenating $h$ heads each of width $d_{v}$ gives width $h \cdot d_{v} = h \cdot \frac{d _{model}}{h} = d_{model}$ . The final matrix $W^{O} \in R^{d_{model} \times d_{model}}$ maps that back to $d_{model}$ . Same width in, same width out — so blocks can be stacked without reshaping.

Level 2 — Application

Problem 2.1

With $d_{model} = 512$ , $h = 8$ , $n = 10$ : give the shapes of $W_{i}^{Q}$ , $Q_{i}$ , the score matrix $Q_{i} K_{i}^{T}$ , and $head_{i}$ .

Recall Solution

$d_{k} = 512/8 = 64$ .

$W_{i}^{Q} \in R^{512 \times 64}$ — it takes a 512-wide token and squeezes it to 64.
$Q_{i} = Q W_{i}^{Q} \in R^{10 \times 64}$ .
$Q_{i} K_{i}^{T} \in R^{10 \times 10}$ — one score for every (token, token) pair. Look at Figure s01: each cell is "how much does row-token look at column-token."
$head_{i} = softmax (\dots) V_{i} \in R^{10 \times 64}$ .

Problem 2.2

For a single head with $d_{k} = 64$ , the raw dot product $q^{T} k = 40$ . Compute the scaled score. Why do we divide by $d_{k}$ and not, say, $d_{k}$ ?

Recall Solution

$score = \frac{q ^{T} k}{d _{k}} = \frac{40}{64} = \frac{40}{8} = 5.0.$ Why $d_{k}$ : if $q$ and $k$ have entries with variance $1$ , the dot product $q^{T} k$ is a sum of $d_{k}$ such products, so its variance grows like $d_{k}$ and its typical size like $d_{k}$ . Dividing by $d_{k}$ cancels exactly that growth, keeping scores at a stable size no matter how big $d_{k}$ is. Dividing by $d_{k}$ would over-shrink them (killing the signal); dividing by $d_{k}$ is the Goldilocks amount.

Problem 2.3

Three tokens give scaled scores $s_{1} = 4.0$ , $s_{2} = 2.0$ , $s_{3} = 5.0$ for query token 3. Compute the softmax attention weights. Which token wins, and by roughly how much?

Recall Solution

Softmax turns scores into probabilities: $w_{j} = \frac{e ^{s_{j}}}{e ^{s_{1}} + e ^{s_{2}} + e ^{s_{3}}} .$ Numerators: $e^{4} = 54.598$ , $e^{2} = 7.389$ , $e^{5} = 148.413$ . Sum $= 210.400$ . $w_{1} = \frac{54.598}{210.400} = 0.2595, w_{2} = \frac{7.389}{210.400} = 0.0351, w_{3} = \frac{148.413}{210.400} = 0.7054.$ Token 3 (attending to itself) wins with $\approx 0.71$ . Because softmax uses the exponential, a score gap of $1.0$ (from $4$ to $5$ ) already makes token 3 roughly $e \approx 2.7$ times more likely than token 1.

Level 3 — Analysis

Problem 3.1

A student claims: "Since attention cost is $O (n^{2} d)$ , using $h = 8$ heads makes the block $8 \times$ slower than single-head attention." Is this correct? Show the FLOP comparison.

Recall Solution

Incorrect. The key is that each head runs in the reduced width $d_{k} = d_{model} / h$ , not the full $d_{model}$ .

Single head, full width: cost $\propto n^{2} d_{model}$ .
$h$ heads, each width $d_{k}$ : cost $\propto h \cdot n^{2} d_{k} = h \cdot n^{2} \cdot \frac{d _{model}}{h} = n^{2} d_{model}$ . The $h$ and the $\frac{1}{h}$ cancel. So multi-head is the same order of cost as one full-width head, yet buys $h$ different views. That is precisely the design trick.

Problem 3.2

Below is a broken forward pass. Find the error.

"For $d_{model} = 512, h = 8$ , I set each $W_{i}^{V} \in R^{512 \times 64}$ , computed 8 heads of width 64, then summed them to get a $10 \times 64$ output, and applied $W^{O} \in R^{64 \times 512}$ ."

Recall Solution

The error is summing instead of concatenating.

Summing collapses the 8 distinct 64-wide views into a single 64-wide vector — the specialized patterns each head learned get averaged away, destroying the whole point of multiple heads.
The correct step is concatenation: $[head_{1} ∣ \dots ∣ head_{8}] \in R^{10 \times 512}$ , then $W^{O} \in R^{512 \times 512}$ . (The stated $W^{O} \in R^{64 \times 512}$ is a downstream symptom of the same mistake.)

Problem 3.3

Explain why every head must use a different learned $W_{i}^{Q}, W_{i}^{K}, W_{i}^{V}$ . What exactly happens if two heads share identical projection matrices?

Recall Solution

Each projection is like a different question asked of the same tokens (see Query-Key-Value Intuition). If head $a$ and head $b$ share $W^{Q}, W^{K}, W^{V}$ , then $Q_{a} = Q_{b}$ , $K_{a} = K_{b}$ , $V_{a} = V_{b}$ , so they produce byte-for-byte identical outputs. In the concatenation those two 64-wide blocks are duplicates — you paid for 2 heads but got 1 distinct view. Different matrices let head $a$ track, say, syntax while head $b$ tracks co-reference. Diversity of projections is the source of multi-head power.

Level 4 — Synthesis

Problem 4.1

Derive the total parameter count of a multi-head attention block (all $W_{i}^{Q}, W_{i}^{K}, W_{i}^{V}$ plus $W^{O}$ ; ignore biases). Show it does not depend on $h$ . Then evaluate for $d_{model} = 512$ .

Recall Solution

Per head: three matrices of shape $d_{model} \times d_{k}$ . Across $h$ heads: $3 \cdot h \cdot (d_{model} \cdot d_{k}) = 3 h d_{model} \cdot \frac{d _{model}}{h} = 3 d_{model}^{2} .$ The $h$ cancels — splitting into more heads doesn't add parameters. Plus $W^{O} \in R^{d_{model} \times d_{model}}$ adds $d_{model}^{2}$ . Total: $4 d_{model}^{2} .$ For $d_{model} = 512$ : $4 \cdot 51 2^{2} = 4 \cdot 262144 = 1048576$ parameters ( $\approx 1.05$ M).

Problem 4.2

Set up (symbolically) the full forward pass for a causal single query at the last position of an $n = 4$ sequence, one head, $d_{k} = 2$ . Given projected vectors $q = (1, 0), k_{1} = (1, 0), k_{2} = (0, 1), k_{3} = (1, 1), k_{4} = (2, 0),$ and value scalars packed as $v_{1} = (1, 0), v_{2} = (0, 2), v_{3} = (3, 0), v_{4} = (0, 4)$ , compute the output for the query at position 4 (it may attend to all 4 keys).

Recall Solution

Raw scores $q^{T} k_{j}$ : $q^{T} k_{1} = 1$ , $q^{T} k_{2} = 0$ , $q^{T} k_{3} = 1$ , $q^{T} k_{4} = 2$ . Scale by $d_{k} = 2 \approx 1.4142$ : $s_{1} = 0.7071, s_{2} = 0, s_{3} = 0.7071, s_{4} = 1.4142$ . Softmax: $e^{s_{j}} = 2.0281, 1.0, 2.0281, 4.1133$ ; sum $= 9.1695$ . $w = (0.2212, 0.1091, 0.2212, 0.4486) .$ Weighted sum of values: $out = \sum_{j} w_{j} v_{j} = 0.2212 (1, 0) + 0.1091 (0, 2) + 0.2212 (3, 0) + 0.4486 (0, 4) .$ First component: $0.2212 + 0 + 0.6635 + 0 = 0.8847$ . Second component: $0 + 0.2181 + 0 + 1.7942 = 2.0123$ . $out \approx (0.8847, 2.0123) .$

Level 5 — Mastery

Problem 5.1

Steel-man the claim "more heads is always better," then refute it with a limiting-case argument. What happens as $h \to d_{model}$ ?

Recall Solution

Steel-man: more heads $\Rightarrow$ more independent subspaces $\Rightarrow$ more relationship types captured, at no extra cost (Problem 3.1) and no extra parameters (Problem 4.1). Sounds free. Refutation via limit: $d_{k} = d_{model} / h$ . As $h$ grows, $d_{k}$ shrinks. At the extreme $h = d_{model}$ , each head has $d_{k} = 1$ : its query and key are single numbers, so the score $q k$ can encode almost nothing — the head is nearly blind. Expressiveness per head collapses. There is a sweet spot (empirically $d_{k} \approx 64$ ): enough width for a rich subspace, enough heads for diversity. So "more heads" trades diversity against per-head capacity — not a free lunch.

Problem 5.2

Consider the degenerate input where all $n$ tokens have identical value vectors $v_{1} = v_{2} = \dots = v_{n} = v$ . Prove that every head's output at every position equals $v$ , regardless of the attention weights. What does this tell you about where attention's power actually comes from?

Recall Solution

For any position $i$ , the output is $\sum_{j} w_{ij} v_{j}$ . With all $v_{j} = v$ : $\sum_{j} w_{ij} v_{j} = v \sum_{j} w_{ij} = v \cdot 1 = v,$ because softmax weights in each row sum to exactly $1$ . The attention pattern becomes irrelevant. Lesson: attention only ever produces a convex combination (weighted average) of the value vectors. If the values carry no variety, the fanciest attention map buys nothing. Power comes from (a) the diversity of the projected values and (b) which values get emphasized — not from the weights alone.

Problem 5.3

In Cross-Attention, queries come from a decoder sequence of length $n_{q}$ and keys/values from an encoder sequence of length $n_{k} \neq = n_{q}$ . What shape is the score matrix $Q_{i} K_{i}^{T}$ , and what is the output length? Why does multi-head still work unchanged?

Recall Solution

$Q_{i} \in R^{n_{q} \times d_{k}}$ , $K_{i} \in R^{n_{k} \times d_{k}}$ , so $Q_{i} K_{i}^{T} \in R^{n_{q} \times n_{k}} .$ Softmax runs over each of the $n_{q}$ rows (across the $n_{k}$ keys), then $\times V_{i} \in R^{n_{k} \times d_{v}}$ gives $head_{i} \in R^{n_{q} \times d_{v}}$ . Output length $= n_{q}$ (one enriched vector per query token). Multi-head machinery is identical — concatenate over $d_{v}$ , apply $W^{O}$ — because concatenation and $W^{O}$ act on the width axis, which never depended on whether $n_{q} = n_{k}$ . See Figure s02.

Recall Quick self-test

What cancels the $h$ in the FLOP count? ::: Each head runs in width $d_{k} = d_{model} / h$ , so $h \cdot n^{2} d_{k} = n^{2} d_{model}$ . Why concatenate heads instead of summing? ::: Summing averages away each head's distinct learned pattern; concatenation preserves all $h$ views. Total parameters of MHA (no biases)? ::: $4 d_{model}^{2}$ , independent of $h$ . Output length of cross-attention? ::: $n_{q}$ (the number of query tokens).