Scaled dot-product attention: Exercises — AI-ML deep dive

Before we begin, one shared vocabulary card, so no symbol is used before it is named:

See Query-Key-Value Model, Softmax Function, and Self-Attention for the ideas these problems lean on.

Level 1 — Recognition

Can you name and shape the pieces?

Exercise 1.1

$Q$ is $3 \times 8$ and $K$ is $5 \times 8$ . What is the shape of $Q K^{⊤}$ ? What does the entry in row 2, column 4 physically mean?

Recall Solution

What we do: $Q K^{⊤}$ multiplies a $3 \times 8$ matrix by an $8 \times 5$ matrix (the $⊤$ flips $K$ from $5 \times 8$ to $8 \times 5$ ). Inner dimensions $8$ and $8$ match, so the product exists and has shape outer dims = $3 \times 5$ .

Why the widths must match: each score is a dot product $q_{i} \cdot k_{j} = \sum_{t = 1}^{8} q_{i, t} k_{j, t}$ — you can only pair up numbers if both vectors have the same length $d_{k} = 8$ .

Meaning of entry $(2, 4)$ : it is $q_{2} \cdot k_{4}$ , the raw compatibility score between the 2nd query and the 4th key — "how much does question 2 line up with label 4."

Answer: shape $3 \times 5$ ; entry $(2, 4)$ = alignment of query 2 with key 4.

Exercise 1.2

In $softmax (\frac{Q K ^{⊤}}{d _{k}}) V$ , which single operation guarantees that each query's weights are non-negative and sum to exactly $1$ ?

Recall Solution

Answer: the softmax step.

Why: softmax on a row $z = [z_{1}, \dots, z_{m}]$ returns $\frac{e ^{z_{j}}}{\sum _{j^{'}} e ^{z_{j^{'}}}}$ . Each $e^{z_{j}} > 0$ (exponentials are never negative), and dividing by the sum of all of them forces the row to add to $1$ . Neither $Q K^{⊤}$ , the scaling, nor $\times V$ can promise this — only softmax does.

Exercise 1.3

If $d_{k} = 100$ , what number do we divide the scores by before softmax?

Recall Solution

Answer: $d_{k} = 100 = 10$ .

Why not $d_{k} = 100$ itself? The spread (standard deviation) of a dot product grows like $d_{k}$ , not $d_{k}$ . To cancel that spread you divide by exactly $d_{k}$ . Dividing by $100$ would over-shrink and flatten every score toward zero.

Level 2 — Application

Can you turn the crank end-to-end?

Exercise 2.1

Let $d_{k} = 4$ , and take one query and two keys: $q = [1, 1, 1, 1], k_{1} = [1, 1, 1, 1], k_{2} = [1, - 1, 1, - 1] .$ Compute the two scaled scores $\frac{q \cdot k _{1}}{d _{k}}$ and $\frac{q \cdot k _{2}}{d _{k}}$ .

Recall Solution

Step 1 (raw scores): $q \cdot k_{1} = 1 + 1 + 1 + 1 = 4$ . $q \cdot k_{2} = 1 - 1 + 1 - 1 = 0$ . Why: $q$ is parallel to $k_{1}$ (identical), so alignment is maximal; $q$ and $k_{2}$ are orthogonal (dot product $0$ ), so no alignment.

Step 2 (scale): $d_{k} = 4 = 2$ , so scaled scores $= 4/2 = 2$ and $0/2 = 0$ .

Answer: $[2, 0]$ .

Exercise 2.2

Take the scaled scores $[2, 0]$ from 2.1. Compute the attention weights via softmax.

Recall Solution

Step: $softmax ([2, 0]) = [\frac{e ^{2}}{e ^{2} + e ^{0}}, \frac{e ^{0}}{e ^{2} + e ^{0}}]$ . Numerically $e^{2} \approx 7.389$ , $e^{0} = 1$ , sum $\approx 8.389$ . $[\frac{7.389}{8.389}, \frac{1}{8.389}] \approx [0.881, 0.119] .$ Interpretation: the query pours $\approx 88%$ of its attention onto key 1 (the matching one) and $\approx 12%$ onto the orthogonal key 2. Weights sum to $1$ . ✓

Exercise 2.3

With weights $A = [0.881, 0.119]$ and values $v_{1} = [10, 0], v_{2} = [0, 10],$ compute the attention output.

Recall Solution

Step (weighted blend): $out = 0.881 [10, 0] + 0.119 [0, 10] = [8.81, 1.19] .$ Why: the output is a soft mixture of the value rows, weighted by how much attention each key received. Because key 1 won most of the weight, the output leans strongly toward $v_{1} = [10, 0]$ . Answer: $[8.81, 1.19]$ .

Level 3 — Analysis

Can you reason about why the numbers behave as they do?

Exercise 3.1

A model uses $d_{k} = 64$ . Query and key entries are independent with mean $0$ and variance $1$ . What are the mean and variance of a raw dot product $q \cdot k$ ? After dividing by $d_{k}$ , what is the new variance?

Recall Solution

Mean: $E [q \cdot k] = \sum_{i = 1}^{64} E [q_{i}] E [k_{i}] = \sum 0 \cdot 0 = 0$ (independence lets us split the expectation of a product).

Variance: each term $q_{i} k_{i}$ has variance $E [q_{i}^{2}] E [k_{i}^{2}] = 1 \cdot 1 = 1$ , and $64$ independent terms add up: $Var (q \cdot k) = 64$ .

After scaling: dividing a random variable by a constant $c$ divides its variance by $c^{2}$ . Here $c = 64 = 8$ , so new variance $= 64/ 8^{2} = 64/64 = 1$ .

Answer: mean $0$ , raw variance $64$ ; scaled variance $1$ . The whole point of $d_{k}$ is to force this last number back to $1$ .

Exercise 3.2

Two scores are $z = [8.0, 7.5]$ (unscaled, $d_{k} = 64$ ). Compute softmax both without scaling and with scaling (divide by $8$ ). Then compute the softmax "self-derivative" $σ_{i} (1 - σ_{i})$ for the first entry in each case. Which case gives a healthier gradient?

Recall Solution

Without scaling. $softmax ([8, 7.5])$ : differences of exponentials, $e^{8} / (e^{8} + e^{7.5})$ . Since $e^{8 - 7.5} = e^{0.5} \approx 1.6487$ , $σ_{1} = \frac{1.6487}{1.6487 + 1} \approx 0.6225.$ Self-derivative $σ_{1} (1 - σ_{1}) \approx 0.6225 \times 0.3775 \approx 0.2350$ .

With scaling (divide by $8$ ): scores become $[1.0, 0.9375]$ . Difference $e^{1 - 0.9375} = e^{0.0625} \approx 1.0645$ , $σ_{1} = \frac{1.0645}{1.0645 + 1} \approx 0.5156.$ Self-derivative $\approx 0.5156 \times 0.4844 \approx 0.2498$ .

Which is healthier? The scaled case: $0.2498 > 0.2350$ . Its self-derivative sits nearer the maximum possible value ( $0.25$ , reached when $σ = 0.5$ ). Bigger derivative = stronger learning signal flowing back through Gradient Descent. The figure below makes this concrete.

How to read the figure below. The blue curve is the softmax self-derivative $σ_{1} (1 - σ_{1})$ — the size of the learning signal — plotted against the softmax output $σ_{1}$ . Notice it is a hill that peaks at $σ_{1} = 0.5$ (dashed gray line, height $0.25$ ) and slides toward $0$ at both ends. The red dot is the unscaled case ( $σ_{1} \approx 0.6225$ ): it has slid partway down the hill to a weaker signal $0.235$ . The green dot is the scaled case ( $σ_{1} \approx 0.5156$ ): it sits almost exactly on the peak, signal $0.250$ . The pedagogical point: scaling drags the softmax output back toward the middle where the gradient is strongest, so the model keeps learning; leaving scores large pushes softmax toward $0$ or $1$ , off the hill, where learning stalls.

Level 4 — Synthesis

Can you assemble pieces you were never shown together?

Exercise 4.1

You are given a causal Attention Mask: query position $2$ is only allowed to attend to keys $1$ and $2$ (not key $3$ , which is "in the future"). The scaled scores for query 2 are $[1.0, 2.0, 5.0]$ across keys $[1, 2, 3]$ . Apply the mask by setting the forbidden score to $- \infty$ , then softmax. What weights does query 2 get?

Recall Solution

Step 1 (mask): replace the score for the forbidden key 3 with $- \infty$ : scores become $[1.0, 2.0, - \infty]$ . Why $- \infty$ and not $0$ ? After softmax, $e^{- \infty} = 0$ , so the forbidden key gets exactly zero weight and vanishes from the sum. A $0$ score would leave $e^{0} = 1$ , still leaking attention.

Step 2 (softmax over surviving keys): $e^{1} \approx 2.718$ , $e^{2} \approx 7.389$ , $e^{- \infty} = 0$ ; sum $\approx 10.107$ . $[\frac{2.718}{10.107}, \frac{7.389}{10.107}, 0] \approx [0.269, 0.731, 0] .$ Answer: $[0.269, 0.731, 0]$ . Notice key 3 — which had the highest raw score $5.0$ — contributes nothing, because the mask forbade it.

Exercise 4.2

Design a tiny numeric $Q, K$ (with $d_{k} = 2$ , one query, two keys) so that after scaling the softmax gives exactly equal weights $[0.5, 0.5]$ . Explain the geometric condition.

Recall Solution

Requirement: equal softmax weights happen iff the two scaled scores are equal, i.e. $q \cdot k_{1} = q \cdot k_{2}$ . Geometric meaning: the query must be equally aligned with both keys — e.g. it bisects them, or is orthogonal to both.

A clean construction: let $q = [1, 0]$ , $k_{1} = [1, 1]$ , $k_{2} = [1, - 1]$ . Then $q \cdot k_{1} = 1$ and $q \cdot k_{2} = 1$ : equal. Scaling by $2$ keeps them equal at $1/ 2$ each. Softmax of two equal numbers is always $[0.5, 0.5]$ . Check: $softmax ([0.707, 0.707]) = [0.5, 0.5]$ . ✓

Level 5 — Mastery

Can you invent, defend, and connect?

Exercise 5.1

A colleague proposes scaling by $d_{k}$ (not $d_{k}$ ) "to be safe." Using $d_{k} = 64$ and query/key entries of variance $1$ , compute the resulting variance of the scaled scores and argue what goes wrong during training.

Recall Solution

Variance after dividing by $d_{k} = 64$ : raw variance $64$ , divided by $6 4^{2} = 4096$ , gives $64/4096 = 0.015625 = 1/64$ . What goes wrong: the scaled scores now have standard deviation $1/64 = 1/8 = 0.125$ — far too small. Almost every score sits in a razor-thin band near $0$ , so softmax outputs are nearly uniform ( $\approx 1/ m$ for every key) no matter what the query asks. The model becomes blind to which key is actually relevant. Over-shrinking is just as harmful as not shrinking. The Goldilocks point: $d_{k}$ makes the variance exactly $1$ — scores neither explode (dead gradients) nor collapse (uniform, uninformative attention).

Exercise 5.2

Split one $d_{k} = 64$ head into $4$ heads of $d_{k} = 16$ each (as in Multi-Head Attention). For a single head, what is the correct scaling factor now, and what variance does each head's scaled score have? Give one reason multi-head is preferred over a single wide head.

Recall Solution

Scaling per head: each head has its own $d_{k} = 16$ , so it divides by $16 = 4$ . Variance: raw dot product over $16$ terms of variance $1$ has variance $16$ ; dividing by $16 = 4$ gives $16/16 = 1$ . Every head is independently normalised to variance $1$ . Why multi-head wins: four narrow heads can each learn a different kind of relationship (syntax, position, coreference, …) and their outputs are concatenated. One wide head must cram all relationships into a single attention pattern. Same total width, more diverse behaviour, and each head keeps a well-conditioned variance- $1$ score distribution.

Exercise 5.3 (capstone)

Full pipeline, no scaffolding. $d_{k} = 2$ , $d_{v} = 2$ . One query, two keys: $q = [2, 0], k_{1} = [2, 0], k_{2} = [0, 2], v_{1} = [1, 0], v_{2} = [0, 1] .$ Compute the final attention output, showing scaled scores, weights, and blend. Then state in one sentence what the output "means."

Recall Solution

Step 1 — raw scores: $q \cdot k_{1} = 4 + 0 = 4$ ; $q \cdot k_{2} = 0 + 0 = 0$ . Step 2 — scale by $2$ : $[4/ 2, 0] = [2.8284, 0]$ . Step 3 — softmax: $e^{2.8284} \approx 16.918$ , $e^{0} = 1$ , sum $\approx 17.918$ . $A \approx [0.9442, 0.0558] .$ Step 4 — blend values: $0.9442 [1, 0] + 0.0558 [0, 1] = [0.9442, 0.0558]$ . Answer: output $\approx [0.944, 0.056]$ . Meaning: the query aligns strongly with key 1 (parallel) and not at all with key 2 (orthogonal), so it retrieves almost entirely $v_{1}$ — a soft "lookup" that returns mostly the content tied to the matching key.

Recall One-line self-test before you close the page

Softmax runs along which axis of the score matrix? ::: Along each row — one row is one query spending its full 100% of attention across all keys.

Connections

Parent: Scaled Dot-Product Attention
Prerequisites & neighbours: Softmax Function · Query-Key-Value Model · Self-Attention · Cross-Attention · Multi-Head Attention · Attention Mask · Positional Encoding · Gradient Descent