Part 6 of 7 · 4 min read
Exercises
This module is optional, so the exercises are harder than usual. All are doable in NumPy. Try each before opening its hint.
Easy
Easy 1: Count the parameters (20 minutes)
For init_block(d, heads), count the parameters as a function of d. Check against the arrays. Then compute the count for d = 4096.
Done when: your formula is 4d^2 for attention plus 8d^2 for the MLP, giving 12d^2, and the code agrees.
Hint
attention has four d x d matrices; the MLP has d x 4d and 4d x d.
Easy 2: Cache arithmetic (15 minutes)
Using kv_cache_gib, compute the cache for 64 layers, 8 key-value heads, 128-wide heads, 100,000 tokens and 2 bytes per value. Then for 16 key-value heads.
Done when: you have both numbers (24.4 and 48.8 GiB) and the ratio (2).
Easy 3: Active parameters (15 minutes)
An MoE layer has 16 experts of 100 million parameters each and uses 2 per token. How many expert parameters exist? How many are active per token? What fraction is that?
Done when: you answer 1.6 billion, 200 million and one eighth, and say what memory you need.
Hint
expert parameters exist for all 16 even though only 2 run.
Medium
Medium 1: Stack blocks (45 minutes)
Write forward(x, blocks) that applies a list of blocks in turn, and a test that a stack of 4 blocks preserves shape and causality. Measure how the output norm changes with depth for random weights, and describe what you see.
Done when: both tests pass and you have a short observation about norms.
Medium 2: Ablate the residual (45 minutes)
Make a copy of transformer_block without the two residual additions. Compare the output's relationship to the input for random weights and 12 stacked blocks: compute the correlation between the input and output of the stack with and without residuals.
Done when: you report both correlations and explain the difference.
Hint
flatten both arrays and use np.corrcoef.
Medium 3: Long-range RoPE (60 minutes)
Plot the score q . k as a function of distance m - n (0 to 200) for random q and k, using rope. Do it with base = 10,000 and base = 1,000,000 (pass base explicitly, or the default is used). Describe what changes and why a larger base helps with long contexts.
Done when: you have two curves and a two-sentence explanation.
Hint
the rotation speed of the slowest pair is set by base; a larger base makes the slowest pairs rotate more slowly.
Hard
Hard 1: Implement grouped-query attention (120 minutes)
Extend causal_multi_head to take kv_heads: compute keys and values with that many heads and let each group of query heads share one. Test that kv_heads = heads equals the original, that kv_heads = 1 runs, and that the parameter count matches attention_params.
Done when: three tests pass, including equality with the original implementation.
Hint
use np.repeat along the head axis to expand each key-value head to its group of query heads.
Hard 2: Add a load-balancing loss (120 minutes)
Compute a balance measure for an MoE layer: the sum over experts of fraction_of_tokens x mean_router_probability, scaled by the number of experts (a common form). Show that it is smallest when the load is uniform. Then use gradient-free search over the router weights (random perturbations) to reduce it on random data.
Done when: you show the measure falls from its starting value and the expert load becomes closer to uniform.
Hint
the router probabilities come from softmax(x @ router_w); average them over tokens.
Hard 3: Predict a bigger model (120 minutes)
Train tiny character-level n-gram or bigram models of different capacities on text (for example, vary the context length from 1 to 6 or the amount of training text), record held-out loss, fit a power law with fit_power_law and predict the loss at the next size up. Then actually run it and report the prediction error.
Done when: you report fit, prediction, truth and a sentence on why the extrapolation may have been wrong.
Hint
loss is cross-entropy in nats on held-out text; make sure the held-out text is never used for counting.