Query, key, value matrices: Exercises — AI-ML deep dive

This page is a self-test ladder for Query, key, value matrices. Every problem states the goal cleanly, then hides a complete worked solution inside a collapsible callout. Reveal only after you try.

Before we begin, one promise: every symbol used below is defined here or in the parent note. We reuse only these:

Definition The Only Symbols We Use Here

$X$ — the input matrix. Each row is one token's embedding vector. Shape $n \times d_{m o d e l}$ (n rows = n tokens, each with $d_{m o d e l}$ numbers).
$W^{Q}, W^{K}, W^{V}$ — three separate learned matrices. Each takes a $d_{m o d e l}$ -long vector and outputs a shorter vector.
$Q = X W^{Q}$ , $K = X W^{K}$ , $V = X W^{V}$ — the query, key, value matrices. Row $i$ is token $i$ 's query/key/value vector.
$d_{k}$ — length of each query and key vector. $d_{v}$ — length of each value vector. These need not be equal — keys live in "matching space," values in "information space," and those can have different sizes.
$q_{i} \cdot k_{j}$ — the dot product: multiply matching entries, add them up. One number. Big = "these two align / are relevant."
$0$ — the zero vector: a vector whose every entry is $0$ (e.g. $[0, 0, 0]$ ). Bold zero always means "all-zeros," never the single number $0$ .
$softmax$ — turns a list of raw scores into positive weights that sum to $1$ (see Scaled Dot-Product Attention).
Two pieces of shorthand this page introduces (not in the parent, defined once here so the page is self-contained): a tilde, written $\tilde{s}$ , always means "the scaled version of score $s$ ," i.e. $\tilde{s} = s / d_{k}$ . And the capital Greek letter delta, $Δ$ , is just a name for "the gap between two scaled scores," $Δ = \tilde{s}_{first} - \tilde{s}_{second}$ — a single number. Both appear only in Exercise 4.1 and are re-explained there.

Level 1 — Recognition

Exercise 1.1 — Shape reading

Given $X \in R^{5 \times 8}$ and $W^{Q} \in R^{8 \times 3}$ , what is the shape of $Q = X W^{Q}$ ? How many query vectors are there, and how long is each?

Recall Solution

Matrix multiply rule: $(n \times m) (m \times p) = (n \times p)$ . The inner $8$ s must match — and they do. $Q = X W^{Q} : (5 \times 8) (8 \times 3) = 5 \times 3.$ So $Q \in R^{5 \times 3}$ . There are 5 query vectors (one per token, one per row), each of length 3 ( $= d_{k}$ ). What it looks like: 5 rows stacked, each row a short 3-number "search request."

Exercise 1.2 — Naming roles

In the library analogy, match each object to its job: (a) $q_{i}$ , (b) $k_{j}$ , (c) $v_{j}$ .

Recall Solution

(a) $q_{i}$ = your search request — "what token $i$ is looking for."
(b) $k_{j}$ = the book spine label of token $j$ — used only to decide whether $j$ is relevant.
(c) $v_{j}$ = the book's actual content of token $j$ — the information handed over once $j$ is judged relevant. Key point: $k_{j}$ and $v_{j}$ come from the same token $j$ but do different jobs, so they use different matrices ( $W^{K}$ vs $W^{V}$ ).

Exercise 1.3 — True/False

" $W^{Q}, W^{K}, W^{V}$ are fixed by the architecture and never change during training." True or false?

Recall Solution

False. They are learned parameters, initialised randomly and updated by Backpropagation via gradient descent. The architecture fixes only their shapes ( $d_{m o d e l} \times d_{k}$ etc.), not their values.

Level 2 — Application

Exercise 2.1 — Compute one query vector

Let $x_{1} = [1, 0, - 2]$ (so $d_{m o d e l} = 3$ ) and $W^{Q} = 210 0 - 1 3 (d_{k} = 2) .$ Compute $q_{1} = x_{1} W^{Q}$ .

Recall Solution

Walk down each column of $W^{Q}$ , dot it with $x_{1} = [1, 0, - 2]$ — exactly the column rule pictured in figure s01.

Column 1 = $[2, 1, 0]$ : $1 (2) + 0 (1) + (- 2) (0) = 2$ .
Column 2 = $[0, - 1, 3]$ : $1 (0) + 0 (- 1) + (- 2) (3) = - 6$ . $q_{1} = [2, - 6] .$ What it looks like: the 3D token got squeezed to a 2D arrow pointing right-and-down in "search space" (figure s02).

Exercise 2.2 — Two tokens at once

Using the same $W^{Q}$ , add token $x_{2} = [0, 1, 1]$ . Stack $X = [1001 - 2 1]$ and compute $Q = X W^{Q}$ .

Recall Solution

Row 1 is $q_{1} = [2, - 6]$ from before. Row 2:

Col 1 $[2, 1, 0]$ with $[0, 1, 1]$ : $0 + 1 + 0 = 1$ .
Col 2 $[0, - 1, 3]$ with $[0, 1, 1]$ : $0 - 1 + 3 = 2$ . $Q = [21 - 6 2] .$ Each row is one token's query. The whole matrix is just Exercise 2.1 done in parallel for every token — that's why we batch as $X W^{Q}$ instead of looping.

Exercise 2.3 — Key for a token

With $x_{2} = [0, 1, 1]$ and $W^{K} = 12 - 1 101,$ compute $k_{2} = x_{2} W^{K}$ .

Recall Solution

Col 1 $[1, 2, - 1] \cdot [0, 1, 1] = 0 + 2 - 1 = 1$ .
Col 2 $[1, 0, 1] \cdot [0, 1, 1] = 0 + 0 + 1 = 1$ . $k_{2} = [1, 1] .$

Exercise 2.4 — When $d_{v} \neq = d_{k}$ : value projection to a different size

Same token $x_{1} = [1, 0, - 2]$ . Now suppose keys live in $d_{k} = 2$ but values live in $d_{v} = 3$ (a genuinely different-sized "information space"). With $W^{V} = 102 0 - 1 1 210 (3 \times 3, so d_{v} = 3),$ compute $v_{1} = x_{1} W^{V}$ and state its length.

Recall Solution

Column rule again (figure s01), $x_{1} = [1, 0, - 2]$ :

Col 1 $[1, 0, 2]$ : $1 (1) + 0 (0) + (- 2) (2) = 1 - 4 = - 3$ .
Col 2 $[0, - 1, 1]$ : $1 (0) + 0 (- 1) + (- 2) (1) = - 2$ .
Col 3 $[2, 1, 0]$ : $1 (2) + 0 (1) + (- 2) (0) = 2$ . $v_{1} = [- 3, - 2, 2], length d_{v} = 3.$ The edge case that matters: $q_{1}, k_{1}$ have length $d_{k} = 2$ but $v_{1}$ has length $d_{v} = 3$ . That is legal — keys are only used inside the dot product $q \cdot k$ (both length $d_{k}$ , so it works), while values are summed separately to build the output. The output vector then has length $d_{v} = 3$ , not $d_{k} = 2$ . Matching space and information space are decoupled.

Level 3 — Analysis

Exercise 3.1 — A relevance score

Take $q_{1} = [2, - 6]$ (Ex 2.1) and $k_{2} = [1, 1]$ (Ex 2.3). Compute the raw relevance score $score_{1, 2} = q_{1} \cdot k_{2}$ , then the scaled score $score_{1, 2} / d_{k}$ with $d_{k} = 2$ .

Recall Solution

Raw: $2 (1) + (- 6) (1) = 2 - 6 = - 4$ . Scaled: $- 4/ 2 = - 22 \approx - 2.828$ . Interpretation: negative score → token 1's query points against token 2's key. After softmax this will get a low attention weight. The $1/ d_{k}$ shrink keeps scores from blowing up as $d_{k}$ grows (which would push softmax into flat, gradient-starved regions).

Exercise 3.2 — Directionality of attention

Is $score_{1, 2}$ (token 1 querying token 2) necessarily equal to $score_{2, 1}$ ? Compute $score_{2, 1} = q_{2} \cdot k_{1}$ to test it. Use $q_{2} = [1, 2]$ (Ex 2.2) and $k_{1} = x_{1} W^{K}$ with $x_{1} = [1, 0, - 2]$ .

Recall Solution

First get $k_{1} = x_{1} W^{K}$ , $x_{1} = [1, 0, - 2]$ :

Col 1 $[1, 2, - 1] \cdot [1, 0, - 2] = 1 + 0 + 2 = 3$ .
Col 2 $[1, 0, 1] \cdot [1, 0, - 2] = 1 + 0 - 2 = - 1$ . $k_{1} = [3, - 1] .$ Now $score_{2, 1} = q_{2} \cdot k_{1} = [1, 2] \cdot [3, - 1] = 3 - 2 = 1$ . Compare: $score_{1, 2} = - 4$ but $score_{2, 1} = 1$ . They are different. Why: attention is directional. Because $W^{Q} \neq = W^{K}$ , "1 looking at 2" and "2 looking at 1" use different projections — exactly the asymmetry the parent note stressed ("cat querying sat" ≠ "sat querying cat"). Figure s03 draws all four arrows ( $q_{1}, q_{2}, k_{1}, k_{2}$ ) so you can see that $q_{1}$ points away from $k_{2}$ (negative) while $q_{2}$ leans toward $k_{1}$ (positive).

Exercise 3.3 — Why not one shared matrix?

Suppose an engineer sets $W^{Q} = W^{K}$ . Show that self-scores $score_{i, i} = q_{i} \cdot k_{i}$ then equal $∥ q_{i} ∥^{2} \geq 0$ , and explain what capacity is lost.

Recall Solution

If $W^{Q} = W^{K}$ then $k_{i} = x_{i} W^{K} = x_{i} W^{Q} = q_{i}$ . So $score_{i, i} = q_{i} \cdot k_{i} = q_{i} \cdot q_{i} = ∥ q_{i} ∥^{2} \geq 0.$ A token can now never assign itself a negative raw score and the score matrix becomes symmetric ( $q_{i} \cdot k_{j} = q_{j} \cdot k_{i}$ ). Lost capacity: the model can no longer express asymmetric relations ("A should attend to B strongly but not vice-versa"). Two separate matrices restore that freedom.

Level 4 — Synthesis

Exercise 4.1 — Full single-head attention output for one token

Given (all computed above or here): $Q = [21 - 6 2], K = [31 - 1 1], V = [1002] .$ Compute the attention output for token 1: (1) raw scores $s_{1, 1}, s_{1, 2}$ , (2) scale by $1/ 2$ , (3) softmax into weights $α_{1, 1}, α_{1, 2}$ , (4) output $= α_{1, 1} v_{1} + α_{1, 2} v_{2}$ .

Recall Solution

Notation for this problem: we write a plain $s$ for a raw score $q \cdot k$ , and a tilde $\tilde{s}$ for its scaled version $\tilde{s} = s / d_{k}$ . The tilde is just "same score, after dividing by $d_{k}$ " — no new maths, only a short label so we don't rewrite the division every line.

Step 1 — raw scores ( $q_{1} = [2, - 6]$ ):

$s_{1, 1} = q_{1} \cdot k_{1} = [2, - 6] \cdot [3, - 1] = 6 + 6 = 12$ .
$s_{1, 2} = q_{1} \cdot k_{2} = [2, - 6] \cdot [1, 1] = 2 - 6 = - 4$ .

Step 2 — scale by $1/ 2 \approx 0.7071$ (this produces the tilde versions): $\tilde{s}_{1, 1} = 12/ 2 = 62 \approx 8.4853, \tilde{s}_{1, 2} = - 4/ 2 = - 22 \approx - 2.8284.$

Step 3 — softmax, shown from scratch. Softmax of two numbers $\tilde{s}_{1, 1}, \tilde{s}_{1, 2}$ is by definition $α_{1, 1} = \frac{e ^{\tilde{s}_{1, 1}}}{e ^{\tilde{s}_{1, 1}} + e ^{\tilde{s}_{1, 2}}} .$ Why we can simplify: divide numerator and denominator by $e^{\tilde{s}_{1, 1}}$ (allowed — it is a nonzero common factor): $α_{1, 1} = \frac{e ^{\tilde{s}_{1, 1}} / e ^{\tilde{s}_{1, 1}}}{e ^{\tilde{s}_{1, 1}} / e ^{\tilde{s}_{1, 1}} + e ^{\tilde{s}_{1, 2}} / e ^{\tilde{s}_{1, 1}}} = \frac{1}{1 + e ^{\tilde{s}_{1, 2} - \tilde{s}_{1, 1}}} = \frac{1}{1 + e ^{- Δ}} .$ Here we gave the gap between the two scaled scores a short name: $Δ = \tilde{s}_{1, 1} - \tilde{s}_{1, 2}$ (a single number). We introduce it only because it makes the formula show clearly that only the difference of scores matters, not their absolute size — this compact $1/ (1 + e^{- Δ})$ is the logistic (sigmoid) form that any two-way softmax reduces to. Now $Δ = 82 \approx 11.3137$ , so $α_{1, 1} \approx \frac{1}{1 + e ^{- 11.3137}} \approx 0.99999, α_{1, 2} = 1 - α_{1, 1} \approx 1.2 \times 1 0^{- 5} .$ Token 1 attends almost entirely to itself.

Step 4 — output ( $v_{1} = [1, 0], v_{2} = [0, 2]$ ): $out_{1} = α_{1, 1} [1, 0] + α_{1, 2} [0, 2] \approx [0.99999, 0.0000244] \approx [1.0, 0.0] .$ What it looks like: because token 1's query aligns strongly with its own key and clashes with token 2's, the output is essentially $v_{1}$ — it copied its own value. This is one row of $softmax (Q K^{⊤} / d_{k}) V$ from Scaled Dot-Product Attention.

Exercise 4.2 — Same machinery, cross-attention

In Cross-Attention the queries come from a decoder sequence but keys/values from an encoder sequence. If decoder has 4 tokens and encoder has 7 tokens, with $d_{k} = d_{v} = 6$ , give the shapes of $Q$ , $K$ , $V$ , the score matrix $Q K^{⊤}$ , and the final output.

Recall Solution

$Q$ : from 4 decoder tokens → $4 \times 6$ .
$K$ : from 7 encoder tokens → $7 \times 6$ .
$V$ : from 7 encoder tokens → $7 \times 6$ .
$Q K^{⊤}$ : $(4 \times 6) (6 \times 7) = 4 \times 7$ — one score per (decoder token, encoder token) pair.
Output $= softmax (Q K^{⊤} / 6) V$ : $(4 \times 7) (7 \times 6) = 4 \times 6$ — one vector per decoder token. The only change from Self-Attention is where the rows of $X$ come from; the algebra is identical.

Level 5 — Mastery

Exercise 5.1 — Multi-head parameter count

Multi-Head Attention uses $h$ heads, each with its own $W^{Q}, W^{K}, W^{V}$ of shape $d_{m o d e l} \times d_{k}$ (and $d_{v} = d_{k}$ ), with $d_{k} = d_{m o d e l} / h$ . For $d_{m o d e l} = 512$ , $h = 8$ : (a) find $d_{k}$ , (b) count parameters in one head's three matrices, (c) count across all $h$ heads, (d) compare to a single head that uses the full width $d_{k} = d_{m o d e l} = 512$ .

Recall Solution

(a) $d_{k} = 512/8 = 64$ . (b) One head: three matrices each $512 \times 64 = 32768$ params → $3 \times 32768 = 98304$ . (c) All $8$ heads: $8 \times 98304 = 786432$ params. (d) Single full-width head, $d_{k} = 512$ : three matrices each $512 \times 512 = 262144$ → $3 \times 262144 = 786432$ . They are equal: $786432$ . Multi-head doesn't cost more parameters — it splits the same budget into $h$ independent subspaces, each learning a different relation pattern.

Exercise 5.2 — Degenerate input: the zero token

Suppose after Embedding Layer + Positional Encoding a token's vector is exactly $0$ (the zero vector — all entries $0$ ). Compute its $q, k, v$ for any $W^{Q}, W^{K}, W^{V}$ . Then, in a 2-token sequence where token 2 is this zero token, what is its raw score against any query, and does it still receive attention weight?

Recall Solution

$q = 0 \cdot W^{Q} = 0$ , and likewise $k = 0$ , $v = 0$ . A zero row in maps to a zero row out (linear maps send $0 \to 0$ ). Raw score of any query $q_{i}$ against this key: $q_{i} \cdot 0 = 0$ . A score of $0$ is not $- \infty$ , so after softmax the zero token still gets nonzero weight ( $e^{0} = 1$ in the numerator). But its value is $0$ , so it contributes nothing to the weighted sum — it dilutes but does not corrupt. (True masking uses $- \infty$ scores, not zero inputs — a different mechanism.)

Exercise 5.3 — Limiting behaviour of the scale factor

Assume the entries of $q$ and $k$ are independent random numbers, each with mean $0$ and variance $1$ (a standard initialisation assumption), and that $q$ 's entries are independent of $k$ 's. The raw dot product is $q \cdot k = \sum_{c = 1}^{d_{k}} q_{c} k_{c}$ . Show its variance is $d_{k}$ , so typical magnitude $\sim d_{k}$ , and explain why dividing by $d_{k}$ keeps the scaled score's magnitude roughly constant — and what softmax does without it.

Recall Solution

Why each term has mean $0$ and variance $\approx 1$ : since $q_{c}$ and $k_{c}$ are independent with mean $0$ , $E [q_{c} k_{c}] = E [q_{c}] E [k_{c}] = 0 \cdot 0 = 0.$ And for independent zero-mean variables, $Var (q_{c} k_{c}) = E [q_{c}^{2} k_{c}^{2}] - 0 = E [q_{c}^{2}] E [k_{c}^{2}] = 1 \cdot 1 = 1.$ So each product term has mean $0$ , variance $1$ . Sum of $d_{k}$ independent such terms: variances of independent variables add, so $Var (q \cdot k) = \sum_{c = 1}^{d_{k}} Var (q_{c} k_{c}) = d_{k},$ giving standard deviation $\approx d_{k}$ . Dividing by $d_{k}$ : $Var (\frac{q \cdot k}{d _{k}}) = \frac{d _{k}}{( d _{k} ) ^{2}} = \frac{d _{k}}{d _{k}} = 1.$ So the scaled score has $O (1)$ spread regardless of $d_{k}$ . Without the scale, scores grow like $d_{k}$ ; feeding huge values into softmax makes one weight $\to 1$ and the rest $\to 0$ — a near one-hot, almost flat gradient. Scaling keeps softmax in a responsive, learnable regime.

Recall Self-check summary

What makes attention directional (asymmetric)? ::: Separate $W^{Q}$ and $W^{K}$ , so $q_{i} \cdot k_{j} \neq = q_{j} \cdot k_{i}$ in general. Why divide scores by $d_{k}$ ? ::: To keep scaled-score variance $\approx 1$ so softmax stays responsive and gradients don't vanish. Can $d_{v}$ differ from $d_{k}$ ? ::: Yes — keys are used only inside $q \cdot k$ (both length $d_{k}$ ), while values build the output separately, so the output has length $d_{v}$ , which may differ. Does multi-head cost extra parameters vs one full head? ::: No — same total, split into $h$ subspaces each of size $d_{k} = d_{m o d e l} / h$ . What does a zero input token contribute to attention output? ::: Nothing ( $v = 0$ ), though it still receives nonzero softmax weight.