Self-attention mechanism in detail: Exercises

This page tests everything from the parent note. Work each problem on paper first, then open the collapsible solution. Levels run from L1 Recognition (do you know the pieces?) up to L5 Mastery (can you reason about the whole machine under pressure?).

Every symbol used here was built in the parent note. A quick reminder of the cast, so line one makes sense:

Definition The symbols you will keep meeting

$n$ — sequence length, the number of words (tokens) in the input.
$d_{m o d e l}$ — embedding dimension, the length of each word-vector as it enters attention.
$X$ — the input, a stack of $n$ word-vectors, each of length $d_{m o d e l}$ . Shape $n \times d_{m o d e l}$ .
$Q, K, V$ — three "views" of $X$ : the query ("what I want"), the key ("what I offer"), the value ("what I actually carry"). Each is $X$ times a learned matrix.
$d_{k}$ — the length of each query and key vector (they must match, since we dot them together).
$d_{v}$ — the length of each value vector, i.e. the width of $V = X W^{V}$ . It may differ from $d_{k}$ in general, though many implementations (and every example below) take $d_{v} = d_{k}$ .
$Q K^{T}$ — the $n \times n$ table of scores: row $i$ , column $j$ = how much word $i$ cares about word $j$ .
$d_{k}$ — the shrink factor that keeps scores from blowing up.
$softmax$ — turns a row of scores into positive weights that add to $1$ .
$A$ — the $n \times n$ attention-weight matrix after softmax. Its single entries are written $α_{ij}$ : the weight word $i$ places on word $j$ . Each row $[α_{i 1}, \dots, α_{in}]$ sums to $1$ .

The one-line machine: $Attention (Q, K, V) = softmax (\frac{Q K ^{T}}{d _{k}}) V$

The figure below traces the shapes through that machine — keep it beside you while you solve, because most L1–L2 slips are really shape slips. Read it left to right: $X$ enters (an $n \times d_{m o d e l}$ block, so $n$ = sequence length, $d_{m o d e l}$ = embedding dimension), gets projected into $Q, K, V$ , the $Q K^{T}$ table appears ( $n \times n$ , one row per query), scaling and softmax reshape that table into $A$ , and $A \times V$ produces the output $Z$ . Notice the $n \times n$ blow-up in the middle — that is the source of the $O (n^{2} d)$ cost you will meet in Exercise 4.2.

Figure — Self-attention mechanism in detail

Prerequisite refreshers if any step feels shaky: 3.2.05-Attention-mechanism-basics, 4.1.01-Transformer-architecture-overview.

L1 — Recognition

Exercise 1.1 — Name the shapes

Let $n = 5$ (five words, so $n$ = sequence length) and $d_{m o d e l} = 8$ (embedding dimension), with $d_{k} = d_{v} = 4$ . State the shape of: (a) $X$ , (b) $W^{Q}$ , (c) $Q$ , (d) $Q K^{T}$ , (e) the softmax output $A$ , (f) the final output $Z$ .

Recall Solution 1.1

A shape is just (number of rows) $\times$ (number of columns). Read it off from what each object is.

(a) $X$ : five word-vectors of length $8$ → $5 \times 8$ .
(b) $W^{Q}$ : takes a length- $8$ vector and returns a length- $4$ vector → $8 \times 4$ .
(c) $Q = X W^{Q}$ : $(5 \times 8) (8 \times 4) = 5 \times 4$ .
(d) $Q K^{T}$ : $(5 \times 4) (4 \times 5) = 5 \times 5$ — every word against every word.
(e) $A = softmax (\dots)$ : softmax is applied entrywise per row, shape unchanged → $5 \times 5$ .
(f) $Z = A V$ : $(5 \times 5) (5 \times 4) = 5 \times 4$ .

Exercise 1.2 — Which matrix asks the question?

For the sentence "it was tired", the model is trying to figure out what "it" refers to. Which learned matrix produces the vector that "goes shopping" for the answer — $W^{Q}$ , $W^{K}$ , or $W^{V}$ ? And which produces the vector that "advertises" a word as findable?

Recall Solution 1.2

The word doing the looking is the query, so its shopping vector is $q_{i} = x_{i} W^{Q}$ → $W^{Q}$ asks the question. Every candidate word advertises through its key $k_{j} = x_{j} W^{K}$ → $W^{K}$ offers itself. ( $W^{V}$ is neither the question nor the advertisement — it is the content handed over once a match is made.)

L2 — Application

Exercise 2.1 — Compute a score row by hand

Two query/key vectors live in $d_{k} = 3$ space: $q_{1} = [2, 0, 1], k_{1} = [1, 1, 0], k_{2} = [0, 2, 3], k_{3} = [2, 0, 1] .$ Compute the three raw scores $q_{1} \cdot k_{1}, q_{1} \cdot k_{2}, q_{1} \cdot k_{3}$ . Which key does $q_{1}$ align with best, and why does that make sense?

Recall Solution 2.1

A dot product is "multiply matching slots, add them up" — it is large when two vectors point the same way.

$q_{1} \cdot k_{1} = 2 (1) + 0 (1) + 1 (0) = 2$ .
$q_{1} \cdot k_{2} = 2 (0) + 0 (2) + 1 (3) = 3$ .
$q_{1} \cdot k_{3} = 2 (2) + 0 (0) + 1 (1) = 5$ . Best match is $k_{3}$ with score $5$ . That is no accident: $k_{3} = [2, 0, 1]$ is identical to $q_{1}$ , so they point in exactly the same direction — the strongest possible alignment. This is the geometric meaning of "high dot product = high relevance."

Exercise 2.2 — Softmax a scaled row

Take the scores from 2.1, scale by $d_{k} = 3 \approx 1.732$ , then apply softmax. Give the three attention weights to two decimals and confirm they sum to $1$ . (Recall: each output weight is one entry $α_{1 j}$ of the attention matrix $A$ — here we are computing the single row for query $q_{1}$ .)

Recall Solution 2.2

Scale first: $[2, 3, 5] /1.732 = [1.155, 1.732, 2.887]$ . Exponentiate: $e^{1.155} = 3.174$ , $e^{1.732} = 5.653$ , $e^{2.887} = 17.94$ . Sum $= 26.77$ . Divide each — the results are $α_{11}, α_{12}, α_{13}$ , the three entries of row 1 of $A$ : $[α_{11}, α_{12}, α_{13}] = [0.119, 0.211, 0.670] \approx [0.12, 0.21, 0.67] .$ They sum to $1.00$ . Notice how softmax amplified the gap: raw scores $2, 3, 5$ became weights $0.12, 0.21, 0.67$ — the top key now dominates far more than its raw lead suggested. That amplification is exactly why we use $exp$ and not plain division.

Exercise 2.3 — Why did we scale?

In 2.2 we divided by $3$ before softmax. Redo the softmax on the unscaled scores $[2, 3, 5]$ and compare the top weight. What did scaling buy us?

Recall Solution 2.3

Unscaled: $e^{2} = 7.389, e^{3} = 20.086, e^{5} = 148.41$ , sum $= 175.89$ . $α_{unscaled} = [0.042, 0.114, 0.844] .$ The top weight jumped from $0.67$ (scaled) to $0.84$ (unscaled) — the distribution got sharper. With larger $d_{k}$ this sharpening runs away: softmax becomes almost one-hot, and gradients in its flat tails vanish. Dividing by $d_{k}$ tames scores back toward unit variance so learning stays healthy. (See 5.3.03-Gradient-flow-in-deep-networks.)

L3 — Analysis

Exercise 3.1 — A full 2-word pass

Let $X = [1001], W^{Q} = W^{K} = [1001], W^{V} = [0110], d_{k} = 2.$ Compute $Q, K, V$ , then $Q K^{T}$ , scale by $2$ , softmax each row, and give the final output $Z = A V$ to two decimals.

Recall Solution 3.1

With identity projections, $Q = K = X = I$ . And $V = X W^{V} = W^{V} = [0110]$ . Scores: $Q K^{T} = I \cdot I = [1001]$ . Scale by $2 \approx 1.414$ : $[0.707 0 0 0.707]$ . Softmax row 0 on $[0.707, 0]$ : $e^{0.707} = 2.028, e^{0} = 1$ , sum $3.028$ → $[0.670, 0.330]$ . Row 1 is the mirror $[0.330, 0.670]$ . $A = [0.670 0.330 0.330 0.670] .$ $Z = A V = [0.670 0.330 0.330 0.670] [0110] = [0.330 0.670 0.670 0.330] .$ Row 0 output $\approx [0.33, 0.67]$ , row 1 $\approx [0.67, 0.33]$ . Because $W^{V}$ swaps the two coordinates, each word ends up carrying mostly the other word's original content — a clean demonstration that values decide what gets moved, independent of who matched whom.

Exercise 3.2 — Reading an attention matrix

A trained head gives, for "The cat which was fluffy sat": $A_{sat} = [0.02, 0.71, 0.05, 0.03, 0.07, 0.12]$ (the row for the query word "sat", columns = The, cat, which, was, fluffy, sat). What has this head learned, and which parent-note claim does it confirm?

Recall Solution 3.2

The weights add to $1.00$ . The overwhelming mass ( $0.71$ ) sits on "cat", the grammatical subject, even though "fluffy" ( $0.07$ ) is physically closer to "sat". This is a subject–verb link learned by $W^{Q}$ / $W^{K}$ : the verb's query aligns with the noun's key. It confirms the parent's Example 2 claim — self-attention captures long-range dependencies directly, with no decay over distance, unlike an RNN that must carry the subject through every intervening step.

L4 — Synthesis

Exercise 4.1 — Permutation blindness

Show that plain self-attention treats "dog bites man" and "man bites dog" as the same set of relationships in the sense that permuting the input rows permutes the output rows identically — i.e. attention is permutation-equivariant. Then explain which mechanism from the chapter fixes this.

Recall Solution 4.1

Let $P$ be a permutation matrix (it just reorders rows; note $P^{T} P = I$ ). Feed $P X$ instead of $X$ :

$Q^{'} = P X W^{Q} = P Q$ , likewise $K^{'} = P K$ , $V^{'} = P V$ .
Scores: $Q^{'} K^{' T} = P Q (P K)^{T} = P Q K^{T} P^{T} = P S P^{T}$ — the score table is reindexed on both axes.
Why softmax commutes with the reindexing. Softmax acts one row at a time and only reorders columns within a row. Multiplying $S$ on the left by $P$ shuffles which rows we have; multiplying on the right by $P^{T}$ shuffles the columns inside each row identically. Since softmax of a row is a function of that row's entries only (and reindexing a row's entries just reindexes softmax's outputs the same way), applying softmax after the shuffle equals shuffling after softmax. Crystallized as the identity: $softmax (P S P^{T}) = P softmax (S) P^{T} = P A P^{T} .$ Concretely: row $i$ of $P S P^{T}$ is row $π (i)$ of $S$ with its columns permuted by $P$ , so its softmax is row $π (i)$ of $A$ with columns permuted by $P$ — exactly row $i$ of $P A P^{T}$ .
Output: unpack the multiplication carefully, $Z^{'} = A^{'} V^{'} = (P A P^{T}) (P V) = P A (P^{T} P) V = P A V = P Z,$ using $P^{T} P = I$ in the middle. So permuting inputs merely permutes outputs the same way — the operation has no idea about order. The fix in this chapter is 4.1.04-Positional-encoding: add position-dependent signals to $X$ so that $x_{i}$ and $x_{j}$ differ even when the words are identical, breaking the symmetry.

Exercise 4.2 — Cost of a longer document

Self-attention costs $O (n^{2} d)$ . You double the sequence length $n$ and also double the model dimension $d$ . By what factor does the attention cost grow? Then contrast with the memory for storing the score matrix.

Recall Solution 4.2

Compute cost scales as $n^{2} d$ . New cost $\propto (2 n)^{2} (2 d) = 4 n^{2} \cdot 2 d = 8 n^{2} d$ → 8×. The score matrix $A$ has shape $n \times n$ , so its memory scales as $n^{2}$ (independent of $d$ ): doubling $n$ alone gives 4× the stored scores. This $n^{2}$ term is the wall that motivates efficient attention variants — see 4.2.02-Computational-complexity-of-transformers.

L5 — Mastery

Exercise 5.1 — Prove the scaling variance claim

The parent asserts: if each $q_{i d}$ and $k_{j d}$ are independent with mean $0$ and variance $1$ , then $Var (q_{i} \cdot k_{j}) = d_{k}$ . Prove it, stating each independence assumption you use, and conclude why $d_{k}$ (not $d_{k}$ ) is the right divisor.

Recall Solution 5.1

Write the dot product as a sum: $q_{i} \cdot k_{j} = \sum_{d = 1}^{d_{k}} q_{i d} k_{j d}$ . Step 1 — mean of one term. By independence of $q_{i d}, k_{j d}$ : $E [q_{i d} k_{j d}] = E [q_{i d}] E [k_{j d}] = 0 \cdot 0 = 0$ . Step 2 — variance of one term. By definition $Var (Y) = E [Y^{2}] - (E [Y])^{2}$ . Since $q_{i d}$ has mean $0$ , $E [q_{i d}^{2}] = Var (q_{i d}) + 0^{2} = Var (q_{i d}) = 1$ ; identically $E [k_{j d}^{2}] = 1$ . Now $Var (q_{i d} k_{j d}) = E [(q_{i d} k_{j d})^{2}] - 0^{2} = E [q_{i d}^{2} k_{j d}^{2}] = indep. E [q_{i d}^{2}] E [k_{j d}^{2}] = 1 \cdot 1 = 1.$ (The zero-mean fact is exactly what turns $E [q^{2}]$ into $Var (q)$ ; without it we'd have an extra $(E [q])^{2}$ term.) Step 3 — variance of the whole sum. The $d_{k}$ product-terms are independent of one another (each uses a different coordinate $d$ ), and for independent quantities variances add: $Var (\sum_{d = 1}^{d_{k}} q_{i d} k_{j d}) = \sum_{d = 1}^{d_{k}} Var (q_{i d} k_{j d}) = \sum_{d = 1}^{d_{k}} 1 = d_{k} .$ This is the claimed result: $Var (q_{i} \cdot k_{j}) = d_{k}$ . Step 4 — the divisor. Variance is $d_{k}$ , so the standard deviation (the typical size of a score) is $d_{k}$ . To rescale a random quantity back to unit spread you divide by its standard deviation, not its variance. Dividing by $d_{k}$ gives $Var (\frac{q _{i} \cdot k _{j}}{d _{k}}) = \frac{Var ( q _{i} \cdot k _{j} )}{( d _{k} ) ^{2}} = \frac{d _{k}}{d _{k}} = 1,$ restoring unit variance so softmax stays well-behaved. Dividing by $d_{k}$ instead would over-shrink to variance $1/ d_{k}$ . ∎

Exercise 5.2 — Softmax gradient vanishing

Consider a single query with two keys and raw (scaled) scores $s = [c, 0]$ . Let $α_{1}$ be the softmax weight on the first key. Compute $α_{1}$ as a function of $c$ , then $\frac{d α _{1}}{d c}$ , and show it collapses to $0$ as $c \to + \infty$ . Explain in one sentence why this is the vanishing-gradient danger the scaling factor guards against.

Recall Solution 5.2

Softmax of $[c, 0]$ : $α_{1} = \frac{e ^{c}}{e ^{c} + e ^{0}} = \frac{e ^{c}}{e ^{c} + 1} = \frac{1}{1 + e ^{- c}}$ — the logistic sigmoid $σ (c)$ . Its derivative: $\frac{d α _{1}}{d c} = α_{1} (1 - α_{1}) = σ (c) (1 - σ (c))$ . As $c \to + \infty$ : $σ (c) \to 1$ , so $σ (c) (1 - σ (c)) \to 1 \cdot 0 = 0$ . (Symmetrically $\to 0$ as $c \to - \infty$ .) One sentence: once scores are pushed far apart the softmax saturates, its slope goes to zero, and gradients cannot flow back through the attention weights — precisely why unscaled large- $d_{k}$ scores kill learning, and why $d_{k}$ keeps $c$ modest.

Exercise 5.3 — Design a masking rule

You want a causal attention where word $i$ may only attend to words $j \leq i$ (no peeking ahead). Describe exactly what you add to the score matrix $S^{'}$ before softmax, and prove that a future position receives weight $0$ after softmax.

Recall Solution 5.3

Add a mask matrix $M$ where $M_{ij} = 0$ for $j \leq i$ and $M_{ij} = - \infty$ for $j > i$ , forming $S^{'} + M$ before softmax. Proof that future weight vanishes: for a forbidden entry, the scaled score becomes $s_{ij}^{'} - \infty = - \infty$ . In softmax the numerator is $e^{- \infty} = 0$ , so $α_{ij} = \frac{e ^{- \infty}}{\sum _{j^{'}} e ^{s_{i j^{'}}^{'} + M_{i j^{'}}}} = \frac{0}{(finite positive)} = 0.$ Thus every future position gets exactly zero weight while the allowed positions still form a valid distribution summing to $1$ (the $- \infty$ entries contribute $0$ to the denominator too). In practice $- \infty$ is a large negative number like $- 1 0^{9}$ , giving weight indistinguishable from $0$ .

Related deep dives: 4.1.03-Multi-head-attention, 4.1.04-Positional-encoding, 4.2.02-Computational-complexity-of-transformers.