Top-k Sampling › Top-k Sampling

When a language model generates text, it produces a probability distribution over tens of thousands of possible next tokens at each step. The question of how to turn that distribution into an actual token choice is called the decoding strategy, and it turns out this choice matters enormously for the quality and character of the resulting text.

Greedy decoding always picks the single most probable token, which sounds reasonable but in practice produces repetitive, bland text. If the most likely word at every step is always selected, the model converges on safe, common phrasing rather than engaging with the richness of language. You might ask a language model to write a story and receive paragraph after paragraph of "The man walked to the store. He bought some things. He went home." Grammatically correct, statistically probable, and deeply uninteresting.

Temperature scaling offers one way out. By dividing logits by a temperature parameter before computing softmax, we can flatten the distribution, giving lower-ranked tokens a better chance. But temperature has a important weakness: it raises the probability of every token, including the nonsensical ones buried in the long tail. A large vocabulary contains tokens like unusual punctuation sequences, rare Unicode characters, and word fragments that should never appear in coherent text. Raising temperature to improve diversity simultaneously makes these garbage tokens more likely to sneak through.

Top-k sampling offers a different solution: rather than redistributing probability across all tokens, first eliminate the ones that cannot possibly be good choices, then sample from what remains. By keeping only the $k$ most probable tokens and discarding everything else, we define a bounded sampling space that excludes the incoherent long tail while preserving meaningful variation among plausible candidates.

The beauty of this approach lies in its simplicity. You are not changing the model's beliefs about what is likely, only restricting which possibilities you are willing to consider. The model still says "Paris" is thirty times more likely than "café" for completing "The capital of France is", and that ratio is preserved after truncation. You are acting like a careful editor who says "I'll accept anything the model finds plausible, but I refuse to publish something the model itself considers highly unlikely."

This chapter explores top-k sampling in depth. We examine the mathematical foundation of truncated sampling, implement the algorithm from scratch and verify it behaves correctly, study how to choose $k$ for different contexts and tasks, see how temperature and top-k interact when combined, and analyze the basic limitations that motivated the development of more sophisticated alternatives like nucleus sampling. By the end, you will have both the conceptual understanding to reason about decoding strategy design and the practical implementation knowledge to apply top-k in real generation pipelines.