Unit 3.6 · about 4 h
309 / 324Decoding policies: from greedy to nucleus sampling
The model outputs probabilities; the policy decides the text. The same logits can produce a memo or gibberish depending on truncation, so the policy must be owned separately from the weights.
- derive
- implement
You can now
- Derive the renormalized distribution after temperature, top-\(k\) and top-\(p\) with their edge cases.
- Implement greedy, temperature, top-\(k\) and top-\(p\) sampling with deterministic seeds.
- Compare decoding policies on matched logits with fixed output budgets, including beam search and repetition control.
Covers
- Logits, greedy decoding and deterministic argmax
- Temperature and distribution renormalization
- Top-k and nucleus top-p truncation
- Beam-search objective and cost vocabulary
- Repetition and length penalties
Study, in this order
4 steps
- CS336 Assignment 1 handout: decoding with temperature and top-\(p\)Stanford CS336 Assignment 1 - training and decoding contract
Section 6, Generating text (pp. 37 to 38): equations 21 and 22 for last-position sampling, temperature scaling in equation 23 with the \(\operatorname{argmax}\) limit, and nucleus sampling in equation 24 as the smallest set with cumulative probability at least \(p\)
- LLMs from Scratch ch05: temperature, top-\(k\) and the `generate` functionLLMs from Scratch companion - training and generation
Sections 5.3.1 (
multinomialversus \(\operatorname{argmax}\), temperature 0.1 and 5.0 behavior), 5.3.2 (top-\(k\) masking with negative infinity) and 5.3.3 (thegeneratefunction combining temperature, top-\(k\) andeos_id) - Modern LLM Notebook 20-generation: truncation arithmeticModern LLM Notebook - configuration, loss and generation
Sections 3 (greedy as \(\operatorname{argmax}\)), 4 (temperature table), 5.1 (top-\(k\) with the 0.71 and 0.29 renormalization), 5.2 (top-\(p\) nucleus with \(p\) 0.9) and 5.3 (industry default of top-\(k\) plus top-\(p\))
- Modern LLM Notebook 20-generation: beam search and repetitionModern LLM Notebook - configuration, loss and generation
Sections 6 (beam search keeps \(K\) paths, fits translation and summarization, bores creative writing) and 7 (repetition penalty divides positive logits and multiplies negative ones)
Depth · Read for objective and cost vocabulary only; full beam and penalty tuning belong to later tracks.
Focus on
- Temperature approaching zero recovers \(\operatorname{argmax}\), and top-\(k\) with \(k\) 1 is greedy by another name.
- Truncation zeroes the tail first, then renormalizes, then samples. Sampling before renormalizing draws from the wrong distribution.
- Top-\(p\) adapts the candidate count to the distribution shape, which is why a fixed \(k\) fails on both sharp and flat distributions.
- Beam search optimizes sequence score over \(K\) paths; it is not a sampling method and it trades diversity for safety.
Reproduce closed-book
Traps
- Applying top-\(p\) to raw logits. Cumulative probability needs probabilities, so
softmaxmust come before sorting and accumulation. - Forgetting renormalization after masking. The kept tokens must sum to one before
multinomialsees them. - Expecting diversity from beam search. Wider beams concentrate on safe sequences; temperature and nucleus control diversity instead.
Practice
Modern LLM Notebook 20-generation Exercises 1 to 3: temperature, top-\(k\), top-\(p\)
Exercises section: Exercise 1 (scaled logits assert), Exercise 2 (top-\(k\) positions assert) and Exercise 3 (top-\(p\) accumulation); each cell asserts the exact answer
Solve each by hand first, then confirm with the cell. Afterwards reproduce fixed-seed greedy, temperature, top-\(k\) and top-\(p\) outputs from your own sampler on matched logits.
Self-test
-
Logits
[2.0, 1.0, 0.5]with temperature 0.5. What are the scaled logits?[4.0, 2.0, 1.0]. Temperature divides the logits, which sharpens the distribution toward the \(\operatorname{argmax}\). -
Logits
[0.1, 2.0, 0.5, 3.0, 0.01]withtop_k2. Which positions survive?Positions 3 (value 3.0) and 1 (value 2.0). Everything else is set to negative infinity beforesoftmax. -
Probabilities
[0.4, 0.3, 0.15, 0.1, 0.05]with \(p\) 0.8. Which tokens stay?The first three. Sorted accumulation gives 0.4, then 0.7, then 0.85, which first reaches 0.8 at the third token. -
The distribution is sharp and \(k\) is 40. What goes wrong, and what fixes it?
Top-\(k\) drags 39 implausible tail tokens into the candidate set. Top-\(p\) fixes it by stopping accumulation once the nucleus reaches \(p\), keeping only the concentrated head.
Completion criterion
Stop when
Your filters pass all three exercise asserts, and your sampler reproduces fixed-seed outputs for each policy on matched logits within a fixed budget.
Mark passed, then open 3.7 Only mark this when you can do it without notes.Leads to
- 3.7 Stopping rules and reproducible generation Bound the policies with stopping rules and exact reproducibility.
- 5.1 Why cache: recomputation in autoregressive decoding See why serial generation needs a cache before measuring it.
Comes after
- 3.5 Checkpoints that really resumeGenerate from resumed weights with owned sampling policies.
Resources for this unit
- LLMs from Scratch companion - training and generation Use the ch05 training and generation code for the small pretraining loop. Add full checkpoint state, explicit loss masks and resume-equivalence tests that the base example does not supply.
- Stanford CS336 Assignment 1 - training and decoding contract Use cross_entropy, adamw, adamw_accounting, learning_rate_schedule, gradient_clipping, data_loading, checkpointing, training_together and decoding specifications plus public tests. Debug on CPU and implement independently in your own repository.
- Modern LLM Notebook - configuration, loss and generation Use English notebooks 09-model-config, 10-training-loss and 20-generation to visualize configuration ownership, shifted loss and decoding policies. Change edge cases and cross-check probability calculations against your own implementation.
Also in this section
- PyTorch saving, loading and reproducibility references Use the state-dict and reproducibility notes to enumerate all state needed for the supported resume guarantee. Test model, optimizer, schedule, scaler, data order and RNG state across an interruption.
- Hugging Face LLM Course - library translation Only after the blank-file loop passes, map tokenization, model configuration, data collation, generation configuration and stopping behavior into Transformers APIs. Keep the from-scratch tests as the authority.
- PyTorch Automatic Mixed Precision recipe Use the
autocast,GradScaler, unscale-before-clip, scaler checkpointing and inf-NaN troubleshooting sections to add mixed precision to the blank-file loop. Keep the CS336 clipping contract as the authority for the clipping rule.