Exercises with solutions — Q/K/V attention: from projections to causal output
General instruction: every answer must show data, transformation, result, one check, and one limitation. A bare number or copied definition is insufficient.
Starting data: Reduced trace: q=[2,1], keys k_Maya=[1,1], k_book=[0,2]. The real head has d_head=4; our vectors show only 2 coordinates (a truncated illustration), but the scaling keeps the true dimension. Raw scores: 3 and 2; division by √4=2 gives 1.5 and 1. Softmax ≈ [0.62,0.38]. Output mixes 62% of v_Maya and 38% of v_book. For an earlier position, the causal mask would remove every future key.
Boundary to retain: Attention maps alone do not prove a causal explanation of the model’s overall behavior.
Exercise 1 — Calculated trace — Three learned projections
Reproduce and annotate the chain Q=XW_Q, K=XW_K, V=XW_V. Increase q’s first component from 2 to 2.4 without changing the keys. Recompute both scores divided by √4 and the attention weights.
Deliverable: a data → operation → result → interpretation table, plus two sentences about the changed value.
Worked solution
Reduced trace: q=[2,1], keys k_Maya=[1,1], k_book=[0,2]. The real head has d_head=4; our vectors show only 2 coordinates (a truncated illustration), but the scaling keeps the true dimension. Raw scores: 3 and 2; division by √4=2 gives 1.5 and 1. Softmax ≈ [0.62,0.38]. Output mixes 62% of v_Maya and 38% of v_book. For an earlier position, the causal mask would remove every future key.
Solved variant: Scaled scores move from [1.5,1] to [1.7,1]. Softmax weights move from about [0.622,0.378] to [0.668,0.332]: the Maya key receives more mass without becoming certain.
Each representation x produces q=xW_Q, k=xW_K, and v=xW_V. Query expresses what the position seeks; key describes how it can be found; value carries information to mix. Dot product q·k measures alignment. In “Maya put down the book… She picked it up”, one head can learn that the query for “She” aligns with the key for “Maya”. The minimum check covers dimensions, sign, and order of magnitude. If the observed change contradicts the prediction, locate the first operation whose direction changes instead of fixing only the final line.
Rubric Exercise 1 — /10
| Criterion | Points |
|---|---|
| Explicit data and shapes | 2 |
| Traceable calculation | 3 |
| Prediction before variation | 2 |
| Interpretation and check | 2 |
| Named limitation | 1 |
Exercise 2 — Diagnose a seductive explanation — Scaling
A colleague claims: « Scaling proves the system will be accurate, fast, and stable in every context. »
- Separate mechanism, assumption, observation, and conclusion.
- Name two correct lesson elements and two unsupported extrapolations.
- Propose a bounded experiment with controlled variable, metric, and stop threshold.
- Rewrite the claim as one defensible sentence.
Reasoned solution
As d_head grows, dot products can become large and saturate softmax. Dividing by √d_head keeps a more stable scale. Before softmax, future positions receive −∞. Softmax turns each permitted row into positive weights summing to 1. Attention maps alone do not prove a causal explanation of the model’s overall behavior.
The claim mixes a local relation with a global guarantee. A defensible version states only the observed mechanism, test conditions, and measured metric. Stop the test if shapes become invalid, the metric crosses the declared degradation threshold, or another variable changes.
Rubric Exercise 2 — /10
2 points per element: separation, lesson grounding, extrapolations, protocol, and rewrite.
Exercise 3 — Architecture decision and transfer — Weighted value mixture
You must reproduce the worked case “Reduced trace: q=[2,1], keys k_Maya=[1,1], k_book=[0,2]. The real head has d_head=4; our vectors show only 2 coordinates (a truncated illustration), but the scaling keeps the true dimension. Raw scores: 3 and 2; division by √4=2 gives 1.5 and 1. Softmax ≈ [0.62,0.38]. Output mixes 62% of v_Maya and 38% of v_book. For an earlier position, the causal mask would remove every future key.” under two conditions. Option A uses the full chain through “Weighted value mixture.” Option B is a transparent baseline that retains “Three learned projections,” calculates the expected output directly, and does not use the compression or adjustment mechanism studied. Build a decision record containing:
- the workload and dominant constraint;
- each option’s mechanism, without slogans;
- one quality, memory, or latency prediction;
- one case where your preferred procedure loses;
- an A/B protocol, metrics, and rollback threshold;
- a bounded verdict: choose, defer, or reject.
Elements of a strong solution
Output is AV: weights choose how much of each value passes. Scores are not themselves the retrieved content. Multiple heads compute different relations, their outputs are concatenated then projected by W_O. During decoding, past K/V are cached; the new query reads that cache without recomputing the prefix. Established mechanisms; numerical simplifications are pedagogical.
A strong answer does not make the newer mechanism the default winner. It retains a measurable baseline, sets thresholds before testing, and separates component cost from whole-system behavior. The verdict names what remains uncertain and the next evidence that could change it.
Rubric Exercise 3 — /15
| Criterion | Points |
|---|---|
| Framing and baseline | 3 |
| Compared causal chains | 4 |
| Protocol and metrics | 4 |
| Rollback threshold | 2 |
| Bounded verdict | 2 |
Extension
Repeat Exercise 3 after reversing the dominant constraint. If you optimized memory, impose a strict quality floor; if you optimized fidelity, halve the memory budget. Identify the first part of the verdict that changes and the evidence required.
Review before submission
Review the packet as if another group had to reproduce it without speaking to you. Are all starting values or assumptions present? Are shapes or roles stated before operations? Does the prediction truly precede the observation? Is the result translated into behavior rather than left as an isolated number? Did you test a boundary value and identify a stop condition? Does the procedure or architecture choice retain a measurable baseline and a rollback threshold set before the test? Finally, highlight one sentence that states what is established, one that remains a hypothesis, and one measurement that could change your verdict. If any element is missing, the work is not reproducible.
Reference appendix for correction
1. Three learned projections
Each representation x produces q=xW_Q, k=xW_K, and v=xW_V. Query expresses what the position seeks; key describes how it can be found; value carries information to mix.
Working formula:
Q=XW_Q, K=XW_K, V=XW_V
2. Query-key compatibility
Dot product q·k measures alignment. In “Maya put down the book… She picked it up”, one head can learn that the query for “She” aligns with the key for “Maya”.
3. Scaling
As d_head grows, dot products can become large and saturate softmax. Dividing by √d_head keeps a more stable scale.
Working formula:
S=QKᵀ/√d_head
4. Causal mask and softmax
Before softmax, future positions receive −∞. Softmax turns each permitted row into positive weights summing to 1.
Working formula:
A=softmax(S+causal mask)
5. Weighted value mixture
Output is AV: weights choose how much of each value passes. Scores are not themselves the retrieved content.
Working formula:
O=AV
6. Multi-head, projection, and cache
Multiple heads compute different relations, their outputs are concatenated then projected by W_O. During decoding, past K/V are cached; the new query reads that cache without recomputing the prefix.
Working formula:
MHA(X)=Concat(head₁…head_h)W_O
Sources and evidence boundary
- Vaswani et al., “Attention Is All You Need”, NeurIPS (2017).
- Course source packet supplied by the owner; named-product details remain source-reported until primary verification.
Scope: Established mechanisms; numerical simplifications are pedagogical. These references support the session frame; they do not turn a reported product choice into an independently verified result.