Top-k and Top-p Sampling

Top-k and top-p (nucleus) sampling are two ways of cutting off the unlikely tail of a language model’s next-token distribution before choosing a token at random. Softmax gives every token in the vocabulary some probability, so with a vocabulary of tens of thousands of tokens, pure sampling occasionally picks a token the model considers nearly impossible. A single such token early in a response then conditions everything after it. Truncation removes the tail, renormalizes what remains, and samples from that.

Top-k

Keep the \(k\) most probable tokens, set the rest to zero, and divide the survivors by their sum. Fan et al. used this for story generation. Its weakness is that \(k\) is fixed. If the model is nearly certain, \(k=40\) still admits 39 tokens that do not belong. If many continuations are reasonable, a small \(k\) removes legitimate choices.

Top-p (nucleus)

Holtzman et al. propose keeping the smallest set of top tokens whose cumulative probability reaches a threshold \(p\), then renormalizing. The size of the set adapts to the distribution: one or two tokens when it is peaked, many when it is flat.

A worked example uses six hypothetical tokens with logits 3.0, 2.2, 1.5, 1.0, 0.2, -2.0 and softmax at temperature 1. The probabilities, sorted, with running totals:

token	probability	cumulative
the	0.533	0.533
a	0.240	0.773
this	0.119	0.892
that	0.072	0.964
one	0.032	0.996
zebra	0.004	1.000

With top-p = 0.9 the loop adds tokens until the cumulative sum reaches 0.9, which happens at “that” (0.964). That token is included because it is the one that crosses the threshold, so the kept set is {the, a, this, that}. Divide by 0.964 and the renormalized probabilities are 0.553, 0.249, 0.123, 0.075. With top-k = 3 the kept set is {the, a, this}, renormalized to 0.598, 0.269, 0.133.

The adaptivity shows on two other distributions. For a peaked one (largest probability 0.968), top-p 0.9 keeps one token while top-k = 3 keeps three. For a flat one (largest probability 0.211), top-p 0.9 keeps all six tokens while top-k = 3 still keeps three.

Interaction with temperature

Temperature divides the logits by \(T\) before softmax, which sharpens (\(T<1\)) or flattens (\(T>1\)) the distribution without changing the ranking. Truncation then works on whichever distribution it receives, so the order matters. For the logits above with \(T=2\) and top-p 0.9, applying temperature first keeps five tokens (probabilities 0.363, 0.243, 0.171, 0.133, 0.089), because the flattened distribution needs more tokens to reach 0.9. Applying top-p first (keeping four tokens) and then temperature gives 0.398, 0.267, 0.188, 0.147. Libraries differ in this order and in how they combine top-k and top-p, so the library’s documentation, not the parameter names, determines behavior.

Other logit adjustments, such as repetition penalties, are applied before softmax and change which tokens are in the nucleus.

What it can and cannot do

• Truncation reduces the chance of a bad token; it does not make the remaining tokens correct. The model’s knowledge is unchanged.
• It makes output less diverse than pure sampling. Holtzman et al. observe that always taking the highest-probability continuation (see Greedy Decoding and Beam Search) produces bland, repetitive open-ended text, and truncated sampling is a middle path between that and unrestricted sampling.
• Top-p = 1.0 and top-k at the vocabulary size both mean no truncation. Top-k = 1 is greedy decoding.
• For tasks with one right answer (extraction, classification), low temperature or greedy decoding is the usual choice and truncation adds little.

References

• Holtzman et al., The Curious Case of Neural Text Degeneration — https://arxiv.org/abs/1904.09751