NLP Foundations
Top-k Sampling
Truncating the vocabulary to a fixed-size shortlist of the k most probable tokens before sampling, and why a fixed k is systematically the wrong size for at least some fraction of every distribution.
beginner · 6 min read
Temperature reshapes a distribution's spread but never removes any tokens from consideration, so a long tail of low-probability nonsense is always technically samplable (see temperature-sampling). Top-k sampling is the earliest widely adopted fix: before sampling, discard every token except the k with the highest probability, renormalise what remains, and sample only from that shortlist. It was popularised for open-ended neural story generation by Fan, Lewis, and Dauphin, 2018, who used it to keep generation coherent while still allowing variation.
The mechanism
Sort the vocabulary by probability, keep the top k, zero out everything else, and renormalise the remaining probabilities so they sum to one:
candidates = top_k(p, k)
p_i' = p_i / sum(candidates) if i in candidates, else 0
Sample from p'. k is a fixed integer chosen ahead of time (values in roughly the 20-100 range are common defaults), independent of what the actual distribution looks like at any given step. That independence is both the appeal, top-k is trivial to implement and cheap to compute, and its central flaw.
The fixed-size problem
A language model's distribution is not the same shape at every step. Sometimes the model is extremely confident: after "The capital of France is", the correct continuation "Paris" might carry 95% of the probability mass, with the remaining 5% smeared across everything else. Sometimes the model is genuinely uncertain: after "My favourite colour is", dozens of colours are all roughly equally plausible.
A fixed k handles neither case well. When the model is confident, k=50 drags 49 essentially irrelevant, low-probability tokens into the candidate pool, giving them a real (if small) chance of being sampled even though the distribution says they are almost certainly wrong. When the model is uncertain and the genuinely reasonable continuations number more than k, top-k arbitrarily chops off some of them, narrowing diversity exactly when the model's own distribution says diversity is warranted. In both cases the fixed shortlist is the wrong size, not because k was chosen badly but because no single k is right for every distribution shape a model produces across a generation.
This observation, that truncation should adapt to the model's confidence rather than use a fixed count, is exactly the motivation behind nucleus-sampling-top-p, which keeps a variable-size set defined by cumulative probability instead of a fixed rank cutoff.
Where it still holds up
Despite the fixed-size critique, top-k remains a common default, often layered together with top-p and temperature in the same request (a model API accepting top_k, top_p, and temperature simultaneously is typical, applied in some fixed order, commonly top-k first as a coarse filter, then top-p, then temperature reshaping what remains). Its simplicity, one integer comparison per candidate, no cumulative sum required, makes it cheap even at very large vocabularies, and as a coarse first-pass filter combined with a second, distribution-aware truncation it still does useful work.
When it falls down
- Wrong size by construction. A fixed
kis too large when the model is confident and too small when it is uncertain; it cannot adapt to the shape of the distribution at each individual step. - Sensitive to vocabulary size and domain. A
ktuned for a general chat model can behave very differently on a narrow-domain fine-tune where the "reasonable" candidate set is naturally smaller or larger. - Small k with a wrong top token is unrecoverable. If the single highest-probability token is wrong (the model is confidently mistaken) and
kis small, top-k cannot help; the flawed token is still in the shortlist regardless ofk. - Order with temperature is a common misconception. Because temperature is a positive monotonic rescale of the logits, it preserves token ranking, so applying it before or after top-k truncation selects the same k tokens and, after renormalisation, yields the same sampling distribution; readers who assume otherwise are conflating top-k with top-p, where order genuinely does matter (see nucleus-sampling-top-p). What does interact non-trivially with top-k's ordering is a logit bias or repetition penalty applied elsewhere in the pipeline, since those can change relative ranking and therefore which tokens land in the shortlist depending on when they run.
Further reading
- Fan, Lewis, Dauphin, 2018, Hierarchical Neural Story Generation, arXiv:1805.04833 - introduces top-k sampling for long-form open-ended generation.
- Holtzman et al., 2019, The Curious Case of Neural Text Degeneration, arXiv:1904.09751 - the paper that motivated top-p as a fix for top-k's fixed-size problem.