Nucleus Sampling › The Top-p Formulation

The insight behind nucleus sampling is simple: instead of asking "how many tokens should I consider?", we ask "how much probability mass should I capture?" This shift in perspective leads to an adaptive algorithm that naturally handles both high-confidence and uncertain predictions.

Think about what we really want when sampling the next token. We want to include all tokens that have a reasonable chance of being correct, and exclude all tokens that are essentially noise. But "reasonable" depends on context. After "The sky is", the token "blue" might have 40% probability, meaning we should consider other options. After "2 + 2 =", the token "4" might have 99% probability, meaning alternatives are probably mistakes.

The key insight is that probability itself tells us what's reasonable. If we capture 90% of the total probability mass, we've included essentially all the tokens the model considers plausible. Everything in the remaining 10% is, by definition, something the model thinks is unlikely. We're not guessing how many tokens to include; we're reading the distribution and respecting what it says.

Notice that this reframing changes what the hyperparameter means. With top-k, you are saying "always consider exactly $k$ options." With nucleus sampling, you are saying "always capture at least $p$ fraction of the probability mass." The second statement is much easier to reason about. Setting $p = 0.9$ means "I want to include the options in the 90% probability core." That is a semantically meaningful statement that stays valid regardless of whether the model is confident or uncertain.

This leads us to define the nucleus as the smallest set of highest-probability tokens that together account for at least $p$ fraction of the total probability. Formally, given a probability distribution $P$ over vocabulary $V$, the nucleus $V_p$ is the minimal set such that:

$$\sum_{x \in V_p} P(x) \geq p$$

where:

• $V_p$: the nucleus, the minimal set of highest-probability tokens we'll sample from
• $P(x)$: the probability assigned to token $x$ by the model
• $p$: the cumulative probability threshold, typically set between 0.9 and 0.95

How do we find this minimal set? The algorithm is straightforward once you see the logic:

1. Sort tokens by probability in descending order. This puts the most likely tokens first.

2. Walk through the sorted list, accumulating probabilities. Keep adding tokens until your running sum reaches or exceeds $p$.

3. Stop as soon as you cross the threshold. The tokens you've accumulated form the nucleus. Discard everything else.

4. Renormalize so that the remaining probabilities sum to 1, giving you a valid distribution to sample from.

Let's express this mathematically. After sorting, we have tokens ordered so that $P(x_1) \geq P(x_2) \geq \ldots \geq P(x_{|V|})$, where $|V|$ is the vocabulary size. We find the smallest $k$ such that:

$$\sum_{i=1}^{k} P(x_i) \geq p$$

where:

• $x_i$: the token with the $i$-th highest probability
• $k$: the number of tokens in the nucleus (determined dynamically based on the distribution)
• $|V|$: the total vocabulary size

The nucleus is then $V_p = \{x_1, x_2, \ldots, x_k\}$, containing exactly the top $k$ tokens needed to reach the probability threshold. The key property of this formulation is that $k$ emerges from the distribution itself. A peaked distribution yields small $k$; a flat distribution yields large $k$.

In practice, this means the nucleus can be as small as 1 token (when the model is extremely confident) or as large as thousands of tokens (when the model is almost completely uncertain across its vocabulary). Top-k cannot span this range with a single setting; nucleus sampling spans it automatically.

After truncating to the nucleus, the probabilities no longer sum to 1. If our nucleus contains 90% of the probability mass, the probabilities inside it sum to 0.90, not 1.0. To sample correctly, we need to renormalize.

The renormalized probability for each token in the nucleus is simply its original probability divided by the total probability mass in the nucleus:

$$P'(x) = \frac{P(x)}{\sum_{x' \in V_p} P(x')} \quad \text{for } x \in V_p$$

where:

• $P'(x)$: the renormalized probability used for sampling
• $P(x)$: the original probability of token $x$
• $\sum_{x' \in V_p} P(x')$: the sum of original probabilities over all tokens in the nucleus, which is the normalizing constant

This renormalization preserves the relative ordering of tokens. If "nice" was twice as likely as "good" before, it remains twice as likely after. We're simply scaling everything up so the probabilities form a valid distribution that sums to 1.

The renormalization step also has a subtle effect on quality: it redistributes the probability mass that was assigned to tail tokens back into the nucleus. This slightly increases the probability of every nucleus token, which is exactly what we want. When the model was "uncertain" about whether to pick from the tail, we're saying: no, allocate that residual probability to the plausible options instead.