Modern LLMs20 / 20

Module 20. The complete LLM: from tokens to training and generation

Try each exercise first. Expand worked answers or project guidance when needed; their open state is shared between RU/EN.

Exercises 1–3 belong to the first route and Exercises 4–5 to the second. Read Section 19 before attempting Exercise 6: the SSM and diffusion comparison appears there to close the module with the scope of the construction.

Download course files (.zip) 2588.1 KB

The archive preserves the source folder structure, both languages, and supporting files.

Exercise 1. Reconstruct the final attention row

Use $\widetilde q_2$ , the keys, and the values in Section 5 to calculate $A_{2,:}$ and $Z_2$ . Why is the final key unavailable to row 1 even though the complete array already exists in memory?

Show answer and solution

Solution. Divide the three dot products by $\sqrt2$ , then normalize their exponentials. The result is $(0.201972,0.397396,0.400632)$ . Their weighted value sum is approximately $(0.343132,0.962634)$ . Multiplication by $0.25I$ followed by addition of $(1,1)$ gives $(1.085783,1.240658)$ . For row 1, the score of key 2 is replaced by $-\infty$ . Visibility is a property of the causal mask, not of whether an array has already been allocated.

Exercise 2. Count parameters and cached values

Use $V=4,d=4,m=3,L=1,h_q=2,h_{kv}=1$ with an independent output projection. How many K/V scalars are stored after three tokens? What changes if $h_{kv}=2$ ?

Show answer and solution

Solution. The two vocabulary tables contribute 32 parameters, Q and O contribute 32, K and V contribute 16, the FFN contributes 36, and the norms contribute 12: 128 in total. The cache holds $2\cdot1\cdot3\cdot1\cdot2=12$ scalars. A second K/V head adds 16 parameters, bringing the total to 144, and doubles the cache to 24 scalars. The concatenated width of the query-head outputs stays unchanged.

Exercise 3. Separate the model from its decoder

Apply top-2 and temperature 0.5 to the final logit row in Section 9. Does the model assign zero probability to EOS? What is the loss if only the EOS target is selected before any weight update?

Show answer and solution

Solution. The two retained logits differ by about $0.126910$ , or $0.253819$ after temperature scaling. The probability of A within the pair is therefore $\sigma(0.253819)=0.563116$ , leaving $0.436884$ for B. The decoder assigns zero to EOS, while the model assigns 0.183502. With only that target selected, the loss is $-\log0.183502\approx1.695530$ ; the mask denominator is one.

Exercise 4. Follow the gradient into a shared table

With tied embeddings, why can the EOS row receive a gradient when EOS is absent from the inputs $(0,1,2)$ ? Where do the two computational roles meet?

Show answer and solution

Solution. The output operation $RE^\top$ uses every row of $E$ , including EOS, so its logit derivative contributes to the EOS row. Input lookup contributes only to rows selected by input IDs. Both paths accumulate into one parameter tensor before the optimizer acts. The notebook's finite-difference check for tied embeddings includes both roles simultaneously.

Exercise 5. Explain the first update and the router

Why is the first AdamW step with zero initial moments independent of $\beta_1,\beta_2\in[0,1)$ after bias correction? Why does that reasoning fail later? Check the selected experts and weights in Section 15.

Show answer and solution

Solution. The factors $1-\beta_1$ and $1-\beta_2$ cancel in the bias-corrected first moments, leaving $\widehat m_1=\widetilde g_1$ and $\widehat v_1=\widetilde g_1^2$ . Later moments contain earlier gradients, so the same cancellation is unavailable. The MoE offsets select experts 2 and 0, but the mixture normalizes the original scores 0.268941 and 0.5. Their weights are 0.349755 and 0.650245. Normalizing offset-adjusted scores would change the routing rule.

Exercise 6. Compare the two changes in Section 19

Calculate the states of $s_p=0.5s_{p-1}+u_p$ with $s_{-1}=0$ and inputs $(1,0,-1)$ . Then verify that the final probabilities of the two-step reverse chain in Section 19 sum to one. Which index visits token positions, and which index refines an entire sequence?

Show answer and solution

Solution. The states are $(1,0.5,-0.75)$ ; $p$ visits successive input positions. The reverse-chain probabilities $(0.2128,0.3472,0.1672,0.2728)$ sum to one. Its index $s$ changes the complete two-position sequence. Both computations repeat an operation, but they do not organize generation in the same way.