Top-k and Nucleus Sampling

Why It Matters

How it Works

The Challenge of Text Generation

When an LLM generates text, it does not "know" the answer; it calculates a probability distribution over its entire vocabulary. If you ask a model to complete the sentence "The cat sat on the...", the model might assign 40% probability to "mat", 20% to "floor", 5% to "sofa", and 0.0001% to "refrigerator". If we always pick the highest probability (greedy decoding), the model becomes predictable and prone to repeating the same phrases. If we pick randomly based on the full distribution, the model will eventually pick the 0.0001% chance word, leading to nonsensical output. Sampling strategies like Top-k and Nucleus sampling are the "filters" we apply to this distribution to ensure the model picks from a "sensible" subset of words.

Top-k Sampling: The Fixed Filter

Top-k sampling is a static approach. We define a hyperparameter (e.g., ). At every generation step, the model identifies the 50 tokens with the highest probability. It then sets the probability of all other tokens in the vocabulary to zero and re-normalizes the remaining 50 probabilities so they sum to 1.0. The model then draws a sample from this truncated distribution. The primary advantage is that it effectively prunes the "long tail" of the distribution—the thousands of irrelevant words that would make the sentence incoherent. However, is rigid. If the model is very confident (e.g., "The capital of France is..."), might be too wide, forcing the model to consider unlikely words. If the model is uncertain, might be too narrow, excluding valid options.

Nucleus Sampling (Top-p): The Dynamic Filter

Nucleus sampling, introduced by Holtzman et al. (2019), addresses the rigidity of Top-k. Instead of picking a fixed number of tokens, we pick a probability threshold (e.g., ). The model sorts all tokens by probability and calculates the cumulative sum. It stops adding tokens to the "nucleus" as soon as the sum of their probabilities exceeds . If the model is confident, the nucleus might contain only one or two tokens. If the model is uncertain, the nucleus might expand to include hundreds of tokens. This dynamic behavior allows the model to be creative when the context allows for it and precise when the context demands it.

Edge Cases and Interaction

In practice, these methods are often combined. One might apply Top-k first to remove extreme noise and then apply Top-p to the remaining tokens. A critical edge case occurs when the model is extremely confident in a single token; if is set too low, the model might be forced to choose a token that is technically within the "nucleus" but contextually poor. Furthermore, when the model enters a "repetition loop," these sampling methods can sometimes struggle to break the cycle because the probability of the repeated token remains high in the model's internal state. Advanced practitioners often use "frequency penalties" or "presence penalties" alongside sampling to explicitly discourage the model from reusing tokens, effectively modifying the logits before the sampling stage occurs.

Common Pitfalls

Sample Code

import torch
import torch.nn.functional as F

def sample_next_token(logits: torch.Tensor, top_k: int = 50,
                      top_p: float = 0.9, temperature: float = 1.0) -> int:
    """Combined top-k + nucleus (top-p) sampling on a 1-D logit vector."""
    logits = logits / temperature
    probs  = F.softmax(logits, dim=-1)           # shape: [vocab_size]

    # ── Top-k: zero out all but the k highest-probability tokens ──
    top_k_vals, top_k_idx = torch.topk(probs, top_k)
    filtered = torch.zeros_like(probs)
    filtered.scatter_(0, top_k_idx, top_k_vals)  # restored to vocab-space

    # ── Top-p (nucleus): within those k tokens keep the smallest
    #    cumulative-probability set that still exceeds top_p ──
    sorted_vals, sorted_ord = torch.sort(filtered, descending=True)
    cum = torch.cumsum(sorted_vals, dim=-1)
    # Remove tokens whose *running* total already exceeds top_p,
    # but keep at least one token (shift mask right by 1)
    remove = (cum - sorted_vals) >= top_p
    filtered.scatter_(0, sorted_ord, sorted_vals * (~remove).float())

    filtered /= filtered.sum()
    return torch.multinomial(filtered, num_samples=1).item()

torch.manual_seed(7)
logits = torch.randn(50000)               # simulated one-step LLM output
token  = sample_next_token(logits, top_k=50, top_p=0.9)
print(f"Sampled token index: {token}")
# Output: Sampled token index: 23174

Key Terms

Related Concepts