Nucleus Sampling › Nucleus Sampling

When GPT models generate text, they face a fundamental challenge: how do you sample from a probability distribution over thousands of possible next tokens in a way that stays coherent while retaining variety? We've seen how temperature scaling adjusts the sharpness of the distribution and how top-k sampling restricts choices to the k most likely tokens. But top-k has a flaw: the number of reasonable next tokens varies dramatically depending on context. Sometimes only one or two tokens make sense; other times, dozens are equally valid. A fixed k cannot adapt to this variation, so it inevitably fails in one direction or the other.

Nucleus sampling, introduced by Holtzman et al. in their 2020 paper "The Curious Case of Neural Text Degeneration", solves this problem by shifting the question entirely. Instead of asking "how many tokens should I consider?", it asks "how much probability mass should I keep?". The algorithm dynamically selects the smallest set of tokens whose cumulative probability exceeds a threshold $p$, a value you set once and reuse across all generation steps. This adaptive approach captures the "nucleus" of the probability mass at each step: all tokens the model considers plausible, while excluding the unreliable tail of unlikely candidates.

Think of the nucleus as the model's zone of reasonable options. In high-confidence situations, that zone is small. In open-ended situations, it is large. Nucleus sampling finds that zone automatically without you having to specify its size in advance. This lets it adapt better than top-k when generating over long sequences with varying levels of model certainty.

The paper itself coined the phrase "nucleus" to describe this idea: the high-density core of the probability distribution, surrounded by a diffuse tail of low-probability tokens. Just as an atomic nucleus contains most of the atom's mass in a tiny volume, the probability nucleus contains most of the probability mass in a small number of tokens. Everything outside the nucleus is statistically noise, and sampling from it produces incoherent text.

Understanding nucleus sampling deeply requires understanding why the tail is harmful. A language model trained on vast amounts of text assigns non-zero probability to millions of possible continuations, many of which are syntactically or semantically wrong in the current context. These wrong tokens sit in the tail. When you include the tail in your sampling pool, you occasionally pick one of these tokens, causing the generated text to veer off course. Once one wrong token is chosen, subsequent predictions are made with a corrupted context, compounding the error. Nucleus sampling prevents this degradation by excluding the tail before sampling even begins.

In this chapter, we'll build up nucleus sampling from first principles, trace through a concrete worked example, implement it in PyTorch, and explore how it interacts with temperature scaling. We'll also look at its limitations and the practical guidance you need to use it effectively in real systems.