Nucleus Sampling › A Worked Example

Abstract formulas come alive with concrete numbers. Let's trace through nucleus sampling step by step with a realistic example.

Suppose a language model predicts the next token after "The weather is" and produces the following probability distribution:

This distribution is already sorted by probability. The model strongly favors positive weather descriptions, with "nice" and "good" together accounting for 60% of the mass.

With $p = 0.9$, we walk through the tokens from highest to lowest probability, keeping a running sum:

At step 5, our cumulative sum hits exactly 0.90, meeting the threshold. We stop here. The nucleus is {"nice", "good", "bad", "great", "terrible"}, containing 5 tokens.

Notice what happened: we didn't need to decide in advance how many tokens to include. The distribution itself determined that 5 tokens were needed to capture 90% of the probability mass. If we had used top-k=5, we would get the same result here, but only by coincidence. Change the distribution slightly and top-k=5 would fail; nucleus sampling at $p = 0.9$ would adapt.

Notice also what got excluded: "wonderful", "cold", and "hot" are all perfectly reasonable weather descriptions in English. The model just assigned them less probability in this context, so they fall below the threshold. This is the right call. The model's distribution reflects its training, and if it thinks "wonderful" is less likely than "terrible" here, we should respect that judgment.

The five tokens in our nucleus have probabilities summing to 0.90. To sample from a valid probability distribution, we divide each by this sum:

You can verify: $0.389 + 0.278 + 0.167 + 0.111 + 0.056 = 1.001$ (the small discrepancy is rounding). We now have a valid probability distribution over just five tokens.

The left panel shows how cumulative probability grows as we add tokens. The curve rises steeply at first (the top tokens contribute most of the mass) then flattens. We stop when we cross $p = 0.9$. The right panel shows how renormalization scales up each probability proportionally, preserving relative rankings while creating a valid distribution.

In this particular case, top-k=5 would include the same tokens. But consider what happens if the model were 90% confident in a single token. Say "Paris" has 0.90 probability and everything else shares the remaining 0.10.

With nucleus sampling at $p = 0.9$, we'd include just "Paris" (0.90 meets the threshold immediately). With top-k=5, we'd include "Paris" plus four tokens that collectively have only 10% probability. Those four tokens are noise that nucleus sampling correctly excludes.

This is the adaptive behavior that makes nucleus sampling effective: it contracts when the model is confident and expands when the model is uncertain.

Let's revisit our weather example to see how the nucleus changes as we vary $p$. With $p = 0.6$, we only need "nice" and "good" (cumulative sum: 0.60). With $p = 0.75$, we need to add "bad" as well. With $p = 0.85$, we add "great". With $p = 0.9$, we add "terrible". With $p = 0.94$, we bring in "wonderful" as well.

Each threshold draws the line at a different point in the distribution. The choice of $p$ is ultimately a statement about how cautious you want to be: how much of the distribution's mass must you capture before you feel comfortable sampling? A $p$ of 0.9 says "I need the 90% core before sampling." A $p$ of 0.99 says "I need nearly everything." The first is conservative and produces coherent text; the second approaches pure sampling and produces more varied but potentially less coherent text.