Top-P Sampling

Top-P (nucleus) sampling fixes top-K's fixed-count limitation: instead of always keeping exactly \(K\) tokens, it keeps however many tokens are needed for their cumulative probability to reach a threshold \(P\) — adapting automatically to how peaked or flat the distribution actually is at each step.

The Algorithm

1. Sort tokens by probability, highest first.
2. Accumulate probabilities from the top down until the running total first reaches or exceeds \(P\).
3. Keep exactly that "nucleus" of tokens (however many that turned out to be); discard the rest.
4. Renormalize the kept probabilities, then sample from this restricted set.

Numerical Example — A Sharp Distribution

Sorted probabilities: \([0.7, 0.15, 0.08, 0.04, 0.03]\), with \(P=0.9\): cumulative sum reaches \(0.7\), then \(0.85\), then \(0.93 \ge 0.9\) — stop here, keeping just the top 3 tokens (\(0.7+0.15+0.08=0.93\)).

Numerical Example — A Flat Distribution

Sorted probabilities: \([0.25, 0.22, 0.20, 0.18, 0.15]\), with the same \(P=0.9\): cumulative sum reaches \(0.25, 0.47, 0.67, 0.85, 1.0\) — the threshold isn't reached until essentially all 5 tokens are included. Notice that with the exact same \(P\) value, top-P automatically kept far more tokens here than in the sharp-distribution example — this adaptivity is exactly top-K's fixed-count approach lacks.

Code

import torch
import torch.nn.functional as F

def top_p_sampling(logits, p=0.9):
    probs = F.softmax(logits, dim=-1)
    sorted_probs, sorted_indices = torch.sort(probs, descending=True)
    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

    cutoff = (cumulative_probs >= p).nonzero()[0].item()   # first index where cumsum reaches p
    kept_probs = sorted_probs[:cutoff + 1]
    kept_indices = sorted_indices[:cutoff + 1]

    kept_probs = kept_probs / kept_probs.sum()   # renormalize
    sampled = torch.multinomial(kept_probs, num_samples=1)
    return kept_indices[sampled]

logits = torch.tensor([2.0, 1.5, 1.0, 0.2, -1.0])
next_token = top_p_sampling(logits, p=0.9)
print(next_token)

Top-K vs Top-P — Directly Compared

Both are frequently used together in practice — applying top-K first as a rough cutoff, then top-P for finer, distribution-aware refinement, alongside temperature scaling.

Common Mistakes

• Setting \(P\) too close to 1.0 — this barely restricts anything, defeating the purpose of excluding the implausible tail; too low a \(P\) can overly restrict variety, similar to top-K's failure modes.
• Assuming top-P alone eliminates the need for temperature — the two serve complementary roles (temperature reshapes the whole distribution's sharpness; top-P restricts which tokens are eligible at all) and are commonly combined.

Interview Relevance

Q: "How does top-P sampling adapt better to varying model confidence than top-K sampling?" Top-K always keeps a fixed number of candidate tokens, regardless of how confident or uncertain the model's actual distribution is at that step. Top-P instead keeps however many tokens are needed to reach a cumulative probability threshold — automatically keeping very few tokens when the model is confident (a peaked distribution) and more tokens when the model is uncertain (a flatter distribution), directly adapting to the shape of the distribution rather than using a one-size-fits-all count.

Practice Question

For sorted probabilities \([0.6, 0.25, 0.1, 0.05]\) and \(P=0.8\), which tokens would top-P sampling keep?

Related DL Notes

• Top-K Sampling
• Temperature (Sampling)
• LLM Perplexity