How does top-k sampling work?

Top-k sampling works by keeping only the k highest-probability next tokens, renormalizing their probabilities, and randomly selecting one of them. The process repeats for every generated token.

Step-by-step example

Suppose an LLM predicts the following next-token probabilities:

A → 0.40
B → 0.30
C → 0.15
D → 0.10
E → 0.05

Step 1: Choose k. Set top-k = 3.

Step 2: Keep the top three. Tokens A, B, and C stay; D and E are excluded from selection.

Step 3: Renormalize. The retained probabilities sum to 0.40 + 0.30 + 0.15 = 0.85. Divide each retained probability by 0.85:

A → 0.471
B → 0.353
C → 0.176

The new probabilities sum to approximately 1.

Step 4: Sample. Draw one token using these probabilities. A is most likely, but B and C can also be chosen.

Step 5: Repeat. Append the chosen token to the context, obtain the next-token probabilities, select the new top-k candidates, and sample again.

predict probabilities
        ↓
keep k highest
        ↓
renormalize
        ↓
sample one token
        ↓
append and repeat

What changes when k changes?

• k = 1: Only the highest-probability token remains, effectively giving greedy decoding.
• Small k: Fewer choices, generally less variety.
• Large k: More choices, generally more variety and a greater chance of weak continuations.
• k covering the vocabulary: No tokens are excluded by top-k filtering.

Important limitation

Top-k uses a fixed number of candidate tokens, even when the model is extremely confident or highly uncertain. For example, k = 50 keeps up to 50 candidates at each step regardless of how their probability mass is distributed.

This differs from top-p sampling, which adjusts the number of candidates based on cumulative probability.

Top-k can also be combined with temperature to adjust probabilities before sampling.

In short: top-k sampling filters out all but the k most likely tokens, rescales the remaining probabilities, and samples one token at each generation step.