Top-k sampling restricts each next-token choice to the k tokens with the highest logits, excluding every other vocabulary token. section 5.3.2
The generation flow is:
- Select the k largest next-token logits with PyTorch’s torch.topk. section 5.3.2
- Replace logits below the kth-largest value with negative infinity, making their softmax probabilities zero. section 5.3.2
- Apply softmax, then sample among the surviving tokens, optionally after temperature scaling. section 5.3.2
Worked top-three example
For a nine-token vocabulary, logits 6.75, 6.28, and 4.51 survive when k = 3. After masking and softmax, their probabilities are 0.5775, 0.3610, and 0.0615; all others receive zero. section 5.3.2
Figure 5.15 Using top-k sampling with k = 3, we focus on the three tokens associated with the highest logits and mask out all other tokens with negative infinity (–inf) before applying the softmax function. This results in a probability distribution with a probability value 0 assigned to all non-top-k tokens. (The numbers in this figure are truncated to two digits after the decimal point to reduce visual clutter. The values in the “Softmax” row should add up to 1.0.)
The figure traces top-k filtering from selected logits through negative-infinity masking to a normalized probability distribution. section 5.3.2
