Top-K Sampling & Repetition Control
Filters vocabulary down to the K highest probability tokens using PyTorch torch.topk.
Sets tokens outside the top K to -inf before softmax.
Guarantees the model never picks nonsensical tail tokens.
The Big Idea
Imagine you are hosting a high-stakes dinner party. If you invite anyone who walks past the door, even someone with a 0.01% chance of being civil, eventually an unruly guest will enter and ruin the entire evening.
In pure multinomial sampling, every token in the vocabulary has a non-zero probability. If an irrelevant punctuation mark or bizarre symbol has a 0.001% probability, rolling the dice thousands of times means it WILL eventually get picked. And once an absurd character enters the context, the autoregressive model tries to make sense of it, derailing the entire story into gibberish!
Top-K Sampling, introduced in Hierarchical Neural Story Generation (Fan et al., 2018), steps in like a club bouncer: it keeps only the top K most likely candidates (e.g. K = 40) and overwrites every other logit with -inf. Because e^(-inf) = 0, those low-probability tail tokens are mathematically erased from consideration before applying temperature-scaled Softmax. Combining Top-K (K = 40) with Temperature (T = 0.8) from Lesson 4.2 is the gold standard used by modern frontier models.
Jargon Buster: Plain-English Translations
- Top-K Sampling: Restricting the sampling distribution to only the K tokens with the highest logits.
- Long-Tail Distribution: The large group of low-probability candidates that individually have small chances, but collectively can introduce errors.
- Negative Infinity Masking: Overwriting excluded logits with -inf so their probability becomes exactly 0.0 after softmax.
- Repetition Penalty: Slightly discounting the logits of recently generated tokens to discourage the model from getting stuck in monotonous phrases.
Architecture Blueprint & Tensor Mechanics
The Top-K sampling algorithm locates the K-th largest logit score, replaces all lower values with -inf, and normalizes the remaining candidates with softmax.
import torch
import torch.nn.functional as F
@torch.no_grad()
def generate_top_k(model, idx, max_new_tokens, block_size, temperature=0.8, top_k=40):
model.eval()
for _ in range(max_new_tokens):
# 1. Crop context
idx_cond = idx if idx.size(1) <= block_size else idx[:, -block_size:]
logits, _ = model(idx_cond)
# 2. Extract final step and scale by temperature
logits = logits[:, -1, :] / max(temperature, 1e-5)
# 3. Top-K filtering: keep only top_k values, mask out the rest with -inf
if top_k is not None:
# Find the threshold value at position K
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
# Any logit smaller than the K-th value is set to -inf
logits[logits < v[:, [-1]]] = float('-inf')
# 4. Softmax converts remaining top_k logits into valid probabilities
probs = F.softmax(logits, dim=-1)
# 5. Sample from the pruned distribution
idx_next = torch.multinomial(probs, num_samples=1)
idx = torch.cat((idx, idx_next), dim=1)
return idxBeginner Traps & Gotchas
1. Setting K larger than the vocabulary size
The Gotcha: Calling torch.topk(logits, 100) when vocabulary size is 65 throws a PyTorch RuntimeError: k out of range.
The Fix: Clamp K dynamically using min(top_k, logits.size(-1)) before calling torch.topk.
2. Applying Top-K filtering after softmax instead of before
The Gotcha: Zeroing out probabilities after softmax means the remaining probabilities no longer sum up to 1.0, breaking multinomial sampling assumptions.
The Fix: Apply Top-K filtering to raw logits by setting them to float('-inf') before calling F.softmax.