Part 6 of 7 · 3 min read
Exercises
Try each exercise before opening its hint.
Easy
Easy 1: Compare by hand (20 minutes)
Using the toy table, compute the cosine similarity of tea with coffee, cat and python on paper, then confirm with code.
Done when: the three numbers match your hand calculation to two decimals.
Hint
the dot product of two vectors that share no non-zero axis is exactly 0.
Easy 2: Read a weight row (15 minutes)
Run the heat-map example and look at the row for the fourth token. Write one sentence saying which earlier tokens it relies on and by how much, and why no weight appears for the fifth token.
Done when: your sentence mentions the mask.
Hint
the fourth row of the heat map is the token "the" (the second one).
Easy 3: Memory of a conversation (20 minutes)
Use kv_cache_bytes to compute the cache for 16,000 tokens with the 7B-class configuration, then again with 8 key-value heads. Express both in GiB.
Done when: you can state the ratio between the two (4) and the reason for it.
Medium
Medium 1: Build your own toy embedding space (45 minutes)
Make a table of 12 words of your own across three or four named axes (for example food, sport, music). Include two ambiguous words. Write tests asserting that related words are closer than unrelated ones, and that your two ambiguous words sit between two groups.
Done when: all tests pass and at least one ambiguous word resolves differently under two contexts.
Medium 2: Prove the causal property (45 minutes)
Write a test that, for random inputs of length 6, changes each position i in turn and checks that outputs before i are unchanged under causal attention. Then show that with the mask removed at least one earlier output changes.
Done when: the loop covers every position and both cases are asserted.
Hint
use np.allclose and compare only the first i rows.
Medium 3: Add heads (60 minutes)
Implement multi_head_attention(x, wq, wk, wv, wo, heads): split the width into heads equal parts, run attention in each, concatenate and apply an output projection wo. Test that with one head it equals the single-head result (when wo is the identity) and that the output shape is unchanged.
Done when: both tests pass for 1, 2 and 4 heads.
Hint
x @ wq has shape (n, d); reshape to (n, heads, d // heads) and transpose.
Hard
Hard 1: Retrieval from scratch (90 minutes)
Write 30 short sentences about three topics. Build a bag-of-words vector for each (counts over a shared vocabulary), search with cosine similarity, and measure how often the top result is on the right topic for 10 questions you write. Then identify two questions where it fails and explain why meaning-based embeddings would help.
Done when: you report accuracy and have two documented failure cases.
Hint
build the vocabulary from all sentences; lower-case and strip punctuation first.
Hard 2: Sliding-window attention (90 minutes)
Add a window parameter so each token may attend only to the previous window tokens (plus itself). Test it against the full causal version on short inputs, and compute how the number of compared pairs scales with length for a fixed window. Explain what the model loses and gains.
Done when: pair counts are linear in length and a test shows that distant tokens no longer affect an output.
Hard 3: Cache the keys and values (120 minutes)
Implement incremental decoding: process tokens one at a time, storing each token's key and value in a cache and computing only the new token's query. Show that the outputs equal full causal attention on the same input, and count the number of key-value computations in both approaches.
Done when: outputs match to 6 decimals and the cached version does n key-value computations instead of n(n+1)/2.
Hint
the new token's query must still attend to all cached keys including its own.