Top-K sampling restricts token generation to only the \(K\) most probable candidates at each step, discarding the long tail of unlikely tokens entirely before sampling — a simple, effective safeguard against occasionally sampling a wildly implausible token.
The Algorithm
- Compute the full probability distribution over the vocabulary (via softmax, possibly temperature-scaled).
- Keep only the \(K\) tokens with the highest probability; discard everything else.
- Renormalize the kept probabilities so they sum to 1 again.
- Sample the next token from this restricted, renormalized distribution.
Numerical Example
Full distribution over 5 tokens: \([0.4, 0.3, 0.15, 0.1, 0.05]\). With \(K=3\): keep only the top 3, \([0.4, 0.3, 0.15]\), which sum to \(0.85\). Renormalize: \(\left[\frac{0.4}{0.85}, \frac{0.3}{0.85}, \frac{0.15}{0.85}\right] \approx [0.471, 0.353, 0.176]\) — the bottom two tokens (originally 0.1 and 0.05 probability) are now completely excluded from sampling, no matter how the dice roll.
Why This Helps
Even a well-trained model's raw probability distribution has a long tail of extremely unlikely but non-zero-probability tokens. Sampling directly from the full, un-truncated distribution occasionally — rarely, but not never — produces a genuinely implausible or nonsensical token purely by chance. Top-K sampling eliminates this risk entirely by construction, never allowing sampling from that unlikely tail at all.
Code
import torch
import torch.nn.functional as F
def top_k_sampling(logits, k=3):
top_k_values, top_k_indices = torch.topk(logits, k)
top_k_probs = F.softmax(top_k_values, dim=-1)
sampled_index_in_topk = torch.multinomial(top_k_probs, num_samples=1)
return top_k_indices[sampled_index_in_topk]
logits = torch.tensor([2.0, 1.5, 1.0, 0.2, -1.0])
next_token = top_k_sampling(logits, k=3)
print(next_token) # sampled ONLY from the top 3 tokens by logit valueThe Fixed-K Limitation
Top-K's core weakness: \(K\) is a fixed number, regardless of how "peaked" or "flat" the actual distribution is at a given step. Sometimes only 1–2 tokens are truly plausible (a fixed \(K=40\) would then include many genuinely bad options); other times dozens of tokens are all reasonably plausible (a fixed \(K=40\) might then cut off legitimately good options). This exact shortcoming is what motivates Top-P Sampling, the next note.
Common Mistakes
- Setting \(K\) too small — this can overly restrict variety, even when the model's distribution genuinely has many reasonable candidate tokens at a given step.
- Setting \(K\) too large — this weakens the safeguard against implausible tail tokens, providing little benefit over unrestricted sampling.
Interview Relevance
Q: "What problem does top-K sampling solve, and what's its main limitation?" It prevents sampling from the long tail of extremely unlikely tokens in a distribution, which can otherwise occasionally produce nonsensical output by chance. Its main limitation is that \(K\) is fixed regardless of how confident or uncertain the model actually is at each specific step — a distribution that's very sharply peaked (few plausible tokens) or very flat (many plausible tokens) both get truncated to the same fixed count, which isn't always appropriate.
Practice Question
For the distribution \([0.5, 0.05, 0.05, 0.05, 0.05, 0.3]\) (unsorted), which tokens would top-K sampling with \(K=2\) keep?
