Nucleus Sampling: How LLMs Choose the Next Token

The Core Idea and Its Invariant

A language model defines a conditional distribution P(xₓ | x₁…xₓ₋₁) over a fixed vocabulary V. Decoding is the problem of turning that distribution into a token. Nucleus sampling is a truncation-then-sample strategy: rather than sampling from all of V (too noisy) or committing to the argmax (too rigid), it restricts sampling to the model's most confident region.

Formally, sort the vocabulary by probability in descending order to get p₁ ≥ p₂ ≥ … ≥ p_|V|. The nucleus V⁽⁾⁽ is the smallest prefix of that sorted list whose cumulative probability reaches the threshold p:

V⁽⁾⁽ = smallest set such that  Σᵤ∈V⁽⁾⁽ P(v) ≥ p

The defining invariant: the kept set is exactly the shortest sorted prefix crossing the cumulative-mass threshold p. Everything outside the nucleus has its probability zeroed; the surviving probabilities are renormalized to sum to 1, and the next token is drawn from that renormalized distribution. The key property — the one top-k lacks — is that |V⁽⁾⁽| is a function of the distribution's shape. When the model is certain (e.g., after "United States of"), the nucleus may contain a single token; when it is uncertain (e.g., after "My favorite color is"), the nucleus may span dozens. The truncation self-tunes to the local entropy.

Step by Step, With Temperature

Given the model's raw logits z ∈ ℝ^|V| for the current position, one decoding step is:

• Apply temperature. Divide logits by τ > 0: z′ = z / τ. Lower τ sharpens (mass concentrates on the top token), higher τ flattens. τ = 1 leaves logits unchanged; τ → 0 approaches argmax.
• Softmax. Compute pᵢ = exp(z′ᵢ) / Σ₃ exp(z′₃) in a numerically stable way (subtract max first).
• Sort descending by probability, tracking original indices.
• Prefix-scan. Walk the sorted list accumulating a running sum; stop at the first index where the cumulative probability ≥ p. That prefix is the nucleus.
• Mask & renormalize. Zero every token past the cutoff; divide the survivors by their new sum so they form a valid distribution.
• Sample. Draw one token from the renormalized nucleus (inverse-CDF / multinomial draw).

In PyTorch-flavored pseudocode:

probs = softmax(logits / temperature)          # O(V)
sorted_probs, idx = sort(probs, desc=True)     # O(V log V)
cum = cumsum(sorted_probs)                      # O(V)
# keep tokens up to and including the first that crosses p
mask = cum - sorted_probs < p                   # shift-by-one to always keep ≥1 token
kept = sorted_probs * mask
kept = kept / kept.sum()                         # renormalize, O(V)
token = idx[multinomial(kept, 1)]                # O(V)

The cum - sorted_probs < p trick (comparing the cumulative sum before adding the current token) guarantees the top-1 token always survives even when pᴹ ≥ p, so the nucleus is never empty.

Complexity Analysis

Let V = |vocabulary| (tens of thousands to hundreds of thousands). Per generated token:

• Softmax: Θ(V) time, Θ(V) space.
• Full sort: the dominant cost. A comparison sort is Θ(V log V) worst and average. This is the naive, most-common implementation.
• Prefix scan + mask + renormalize: Θ(V) each.
• Multinomial sample: Θ(V) to build the CDF and binary-search into it (O(log V) for the search itself once the CDF exists).

So a straightforward implementation is O(V log V) time, O(V) space per token. Over a sequence of length n this is O(n · V log V) — but in practice the softmax-and-decode step is dwarfed by the transformer forward pass (O(n² · d) attention, plus the large matrix multiplies), so nucleus sampling is essentially free relative to inference.

You can shave the log factor. Because you only need the sorted prefix that reaches mass p, and that prefix is usually tiny, a partial selection beats a full sort. torch.topk or a heap-based selection finds the top-k in O(V + k log k); iteratively doubling k (k = 8, 16, 32, …) until the cumulative mass crosses p gives an expected O(V) when the nucleus is small — the common case. Quickselect-style partitioning around a probability threshold reaches O(V) average as well. The catch: the threshold that bounds the nucleus isn't known in advance, so these fast paths add branch complexity and are worthwhile mainly at very large V or high batch throughput (as in vLLM's fused sampling kernels).

Why the Fat Tail Wrecks Naive Decoding

The paper's central empirical finding: a well-trained LM's tail is unreliably long. Summed over 50,000+ tokens, the low-probability mass (each individual token near-zero) becomes substantial in aggregate — Holtzman et al. measured that pure sampling routinely draws from tokens the model assigned near-zero, tanking coherence. Two clean failure regimes bracket the design space:

• Maximization (greedy/beam) repeatedly picks high-probability tokens and falls into neural text degeneration — verbatim loops like "I don't know. I don't know. I don't know." Human text is not the maximum-likelihood sequence; real language has surprises, and always taking the mode strips them out. This is a positive-feedback trap: repeating a phrase raises its own conditional probability.
• Pure sampling respects the surprise but also samples the garbage tail, yielding topic drift and gibberish.

Nucleus sampling threads the needle: it truncates the unreliable tail (killing gibberish) while preserving the reliable head's diversity (killing loops). Crucially, it does so adaptively. Holtzman et al. showed that top-p generations match human text on perplexity, self-BLEU (diversity), and repetition statistics far better than top-k or beam search, and human raters preferred them. The recommended default landed around p = 0.95.

Trade-offs: top-p vs top-k vs Temperature

The three knobs interact and are frequently confused:

• Top-k fixes the count of candidates. Its flaw is the whole reason nucleus sampling exists: a fixed k is simultaneously too wide for peaked distributions (after "United States of", k = 40 lets in 39 junk tokens) and too narrow for flat ones (a genuinely open choice with 200 reasonable continuations gets clipped to 40). Nucleus sampling makes the count a consequence of the mass, so it dilates and contracts automatically.
• Top-p fixes the cumulative probability. Trade-off: choosing p is choosing how much tail risk to accept. p near 1.0 approaches pure sampling (incoherence risk); p too low collapses toward greedy (blandness, loops).
• Temperature reshapes the distribution before truncation, so it is orthogonal and usually applied first. High τ + low p can cancel out; the two are best tuned together.

When to use which: for open-ended creative or conversational generation, top-p (often with temperature) is the default. For tasks wanting the single best-supported answer (translation, factual QA, code with a canonical output), greedy or beam search can win. Modern stacks combine methods: min-p, typical sampling, and η-sampling (eta) are refinements; many APIs also expose top-k and top-p together, applying the intersection of both filters.

Where It Runs in Production

Nucleus sampling is the industry default sampler:

• Hugging Face transformers implements it in TopPLogitsWarper (paired with TopKLogitsWarper and TemperatureLogitsWarper), chained in generate() when you pass do_sample=True, top_p=0.9.
• vLLM and TensorRT-LLM fuse temperature + top-k + top-p into GPU sampling kernels so the truncation happens without a full host-side sort — critical when serving thousands of concurrent sequences, since a naive O(V log V) sort per token per request would be a real cost at scale.
• OpenAI, Anthropic, and Google APIs expose top_p / temperature as first-class request parameters; the vendors' guidance is generally to tune one or the other, not both, to avoid compounding effects.
• Structured/constrained decoding (JSON schemas, grammars) composes with nucleus sampling by masking illegal tokens to probability 0 before the top-p cut, so the nucleus is formed only from grammar-valid continuations.

The reason it's ubiquitous is deployment ergonomics: one scalar, p, gives a monotone dial from safe-and-focused (low p) to diverse-and-risky (high p), with a defensible statistical meaning — "trust the model's top p fraction of belief."

Pitfalls, Edge Cases, and Variants

Nucleus sampling has sharp edges worth knowing:

• Off-by-one / empty nucleus. If pᴹ alone already exceeds p, a naive "keep while cumsum ≤ p" keeps zero tokens. Always keep at least the top token — the standard fix is to compare the cumulative sum excluding the current element, guaranteeing |V⁽⁾⁽| ≥ 1.
• Interaction order matters. Temperature must be applied before computing softmax probabilities used for the top-p cut; applying it after truncation changes the semantics. Renormalization is mandatory — sampling from unnormalized survivors silently biases toward whichever tokens survived.
• Numerical stability. Subtract the max logit before exp to avoid overflow; a runaway logit can otherwise produce inf/NaN and a degenerate nucleus of one token.
• Determinism / reproducibility. Because it samples, output varies run-to-run unless you fix the RNG seed. Ties in probability under a sort can also break differently across hardware/kernels.
• The p→1 trap. At p = 1.0 you're doing pure sampling and inherit the fat-tail incoherence the method was built to prevent; many bugs are "someone set top_p=1 and temperature=1.5."

Variants address top-p's own weakness — that a fixed cumulative threshold still admits many low-probability tokens when the head is flat. Min-p keeps tokens with probability ≥ (min_p × pᴹ), scaling the floor to the top token. Typical sampling (Meister et al.) keeps tokens whose information content is near the distribution's conditional entropy rather than its most probable. η-sampling and locally typical decoding blend an entropy-aware threshold with an absolute floor. All share nucleus sampling's DNA: truncate an unreliable tail, then sample the rest.

Nucleus sampling (top-p) vs. its main decoding alternatives, per generated token.

Method	Truncation set	Adaptivity	Time / step	Failure mode
Greedy / argmax	1 token	None	O(V)	Repetitive loops, degeneration
Beam search (width k)	k best sequences	None	O(k·V log(k·V))	Blandness, repetition on open text
Top-k sampling	Fixed k tokens	None (k constant)	O(V + k log k)	Too wide when peaked, too narrow when flat
Top-p (nucleus)	Smallest set, cumprob ≥ p	Yes — adapts to entropy	O(V log V)	Runaway if p→1; too flat if p low
Pure sampling (p=1)	Whole vocabulary	N/A	O(V)	Incoherence from the fat tail