Top-p sampling, also known as nucleus sampling, is a stochastic decoding strategy used to generate sequences from autoregressive probabilistic models, such as large language models. It was originally proposed by Ari Holtzman, Yejin Choi, and colleagues in 2019 for natural language generation to address the issue of repetitive and nonsensical text produced by other common decoding methods like beam search. The technique has since been applied in other scientific fields, including protein engineering and geophysics.
In top-p sampling, a probability threshold p is set, and the next item in a sequence is sampled only from the smallest possible set of high-probability candidates whose cumulative probability exceeds p. This method adapts the size of the candidate pool based on the model's certainty, making it more flexible than top-k sampling, which samples from a fixed number of candidates. Due to its effectiveness, top-p sampling is widely used in many large language model applications.
Technique
At each step of the text generation process, a language model calculates a probability distribution over its entire vocabulary for the next token. While simply picking the token with the highest probability (greedy search) or a limited set of high-probability sequences (beam search) is possible, these deterministic methods often produce text that is dull, repetitive, or nonsensical. Top-p sampling introduces randomness to avoid these issues while maintaining quality.
The core idea is to sample from a smaller, more credible set of tokens at each step, called the nucleus. This nucleus contains the most likely next tokens whose combined, or cumulative probability, just exceeds the threshold p. By sampling only from this dynamically sized group, the model can adapt to different situations. When the model is confident about the next token (e.g., one token has a very high probability), the nucleus will be small. When the model is uncertain (the probabilities are more evenly distributed), the nucleus will be larger, allowing for more diversity.
The process at each step is as follows:
- The model calculates the probabilities for all possible next tokens.
- The tokens are sorted by their probability in descending order.
- The nucleus is formed by selecting tokens from the top of the list until their cumulative probability exceeds the predefined threshold, p.
- The probabilities of tokens within this nucleus are then rescaled so that they sum to 1. All tokens outside the nucleus are discarded (given a probability of 0).
- The final next token is randomly sampled from this new, smaller distribution.
Formally, the nucleus, V^(p) ⊆ V, is defined as the smallest set of tokens satisfying:
∑_{x ∈ V^(p)} P(x | x_1, …, x_{t-1}) ≥ p
In this formula, P(x | x_1, …, x_{t-1}) represents the probability of a token x given the preceding tokens x_1, …, x_{t-1}.
Example
Imagine at a certain step, a language model has a vocabulary of five words: [the, a, cat, dog, eats] and produces the following probabilities:
- the: 0.5
- a: 0.2
- cat: 0.1
- dog: 0.1
- eats: 0.1
If we set p = 0.8:
- The tokens are sorted by probability: [the, a, cat, dog, eats].
- The cumulative probability is calculated:
- the: 0.5
- the + a: 0.5 + 0.2 = 0.7
- the + a + cat: 0.7 + 0.1 = 0.8
- The nucleus is the smallest set with cumulative probability ≥ 0.8, which is V^(0.8) = {the, a, cat}.
- The probabilities for this set are rescaled to sum to 1:
- P(the) = 0.5 / 0.8 = 0.625
- P(a) = 0.2 / 0.8 = 0.25
- P(cat) = 0.1 / 0.8 = 0.125
- The next token is then sampled from this new distribution, meaning dog and eats have a 0% chance of being chosen.
Top-k sampling
Top-k sampling is a similar technique where the pool of candidate tokens is restricted to the k most likely tokens. The main advantage of top-p is its adaptability. When the model is very certain about the next token (a peaked distribution), the nucleus V^(p) can be very small. When the model is uncertain (a flat distribution), the nucleus can be much larger, allowing for more diversity. In contrast, top-k always samples from a fixed number of tokens, which may be too restrictive or too broad depending on the context.
Applications
While top-p sampling is most famously used as a decoding strategy for large language models, the technique has also been adapted for use in other scientific domains that involve generating or analyzing sequential data from probabilistic models.
Natural language generation
In its original domain of natural language generation, top-p sampling is valued for its ability to produce more diverse and coherent text compared to deterministic methods. It has been shown to be beneficial in tasks like automatic question generation, where sample diversity is important for creating effective training data for question answering models.
Drug and protein design
Top-p sampling is used in computational biology to generate novel molecular and protein sequences from specialized language models. In de novo drug design, chemical language models trained on molecular structures use nucleus sampling to generate focused libraries of new, valid drug candidates. Similarly, in protein engineering, top-p sampling helps explore sequence space while maintaining the likelihood of producing functional proteins.
Geophysics
In geophysics, top-p sampling has been applied to generate seismic waveforms or other sequential data from probabilistic models, aiding in tasks such as earthquake simulation or subsurface imaging. The adaptive nature of the nucleus allows for realistic variability in generated data, which is crucial for modeling complex natural phenomena.
Relationship to Other Decoding Strategies
Top-p sampling is one of several stochastic decoding methods used in generative models. It is often combined with temperature scaling to further control the randomness of the output. While temperature scaling adjusts the sharpness of the probability distribution before sampling, top-p sampling truncates the distribution to a subset of tokens. These techniques can be used together to achieve a balance between diversity and coherence.
Compared to deterministic methods like greedy search or beam search, top-p sampling introduces stochasticity, which can prevent repetitive loops and produce more varied outputs. However, this randomness can also lead to occasional incoherence, and the choice of p is critical. A low p value (e.g., 0.5) makes the output more focused and deterministic, while a high p value (e.g., 0.95) increases diversity but may reduce quality.
Implementation Considerations
In practice, top-p sampling is implemented in most modern deep learning frameworks and is a standard parameter in APIs for large language models, such as those from OpenAI, Anthropic, and Google DeepMind. The threshold p is typically set between 0.9 and 0.95 for general text generation tasks, but the optimal value depends on the specific application and desired trade-off between creativity and accuracy.
One challenge in implementation is ensuring that the cumulative probability calculation is efficient, especially for large vocabularies. However, since the nucleus is usually small, the computational overhead is minimal. Additionally, top-p sampling can be combined with other techniques like RLHF (reinforcement learning from human feedback) to align generated text with human preferences.
Limitations and Extensions
Despite its advantages, top-p sampling has limitations. The threshold p is static and does not adapt to the context, which may lead to suboptimal performance in some scenarios. Researchers have explored extensions such as dynamic p values based on the entropy of the distribution, but these are not yet widely adopted. Furthermore, top-p sampling does not guarantee global coherence, as it only considers local probabilities at each step.
Another limitation is that top-p sampling can still produce repetitive text if the model's distribution is heavily peaked on a few tokens. In such cases, combining top-p with other strategies like top-k or temperature adjustments may be necessary. Despite these challenges, top-p sampling remains a fundamental tool in the generative AI toolbox, widely used in both research and production systems.