Top-k Sampling › Summary

Top-k sampling provides a practical solution to the long-tail problem in language model decoding. By truncating the vocabulary to the $k$ most likely tokens and renormalizing, it eliminates improbable tokens while preserving meaningful diversity among plausible choices.

Key takeaways from this chapter:

• Truncation principle: Top-k zeros out all tokens outside the $k$ highest-probability candidates, then renormalizes the remaining probabilities. This prevents sampling from the incoherent long tail while letting variety among reasonable options.

• Temperature interaction: Temperature scaling reshapes the distribution before truncation. Low temperature concentrates probability in fewer tokens; high temperature spreads it more evenly. Applying temperature first affects which tokens make the top-$k$ cut.

• Parameter selection: Common values range from $k$=40 to $k$=100. Lower $k$ produces more focused text; higher $k$ allows more diversity. The optimal choice depends on task, context confidence, and user preferences.

• Fixed-k limitation: A constant $k$ applies regardless of context, including too many tokens when the model is confident and potentially too few when it's uncertain. This motivates adaptive approaches like nucleus sampling.

• Computational efficiency: Top-k selection adds negligible overhead, with $O(n \log k)$ time complexity where $n$ is vocabulary size and $k$ is the number of tokens kept. This is fast compared to the model forward pass, which makes it practical for production use.

• Edge cases: Production implementations should handle edge cases including $k$ larger than vocabulary size, near-zero temperature, and numerical instability in logits. A reliable fallback to greedy decoding ensures reliable behavior across all inputs.

Top-k sampling strikes a useful balance between greedy decoding (deterministic but repetitive) and pure sampling (diverse but sometimes incoherent). While nucleus sampling offers more adaptive truncation, top-k remains widely used due to its simplicity and interpretability. The next chapter explores nucleus sampling, which addresses the fixed-k limitation by adapting the threshold to each distribution's shape rather than maintaining a fixed count.