Gen AI

Top-P (Nucleus) Sampling

Instead of keeping a fixed number of tokens, Top-P dynamically keeps as many tokens as needed until their combined probability hits a target percentage, adapting beautifully to how confident the model is.

Why Does This Exist?

Top-K sampling was a great first attempt at stopping LLMs from generating absolute gibberish. By chopping off the "long tail" of the vocabulary, it prevented the model from accidentally picking a terrible word with a 0.001% probability.

But Top-K is rigid. If you set $K=40$ , the model always keeps exactly 40 words. If the model is extremely confident and the first 2 words make up 99% of the probability, Top-K forces it to keep 38 terrible words in the pool just to hit the quota. If the model is completely unsure and there are 100 equally great synonyms, Top-K violently chops off 60 of them, destroying the model's creativity.

Top-P (often called nucleus sampling) fixes this by adapting dynamically. Instead of a fixed count, it targets a fixed probability mass.

Think of It Like This

Packing a suitcase to a weight limit

Top-K is like packing for a flight with a rule that says: "You must pack exactly 10 items, no matter what." If you pack 10 heavy coats, your bag is overweight. If you pack 10 pairs of socks, you freeze on your trip.

Top-P is like packing with a normal airline weight limit (e.g., 50 pounds). If you pack heavy coats, you might hit 50 pounds after just 3 items, so you stop. If you pack socks, you might fit 100 items before you hit the limit. The number of items changes dynamically, but the total "weight" (probability mass) remains constant.

How It Actually Works

The Cumulative Cutoff

Like Top-K, Top-P is applied after temperature but before the final token is chosen. The algorithm works by accumulating probability from the top down:

Sort all the tokens in the vocabulary descending by probability.
Start adding the probabilities together, one by one, from largest to smallest.
Stop exactly when the running total (cumulative probability) crosses the user-defined threshold $P$ (e.g., $P = 0.90$ ).
Keep the tokens that contributed to that 90%. Throw everything else away (set their probability to 0).
Re-normalize the surviving tokens so they sum to 1.0, and sample.

Adapting to Confidence

The genius of Top-P is how it reacts to the shape of the distribution.

If the model is confident, the distribution is sharp. Token 1 might be 85% likely, and Token 2 is 10%. The running total hits 95% after just two tokens. Top-P (with $P=0.90$ ) immediately stops and throws away the other 99,998 words. It acted just like $K=2$ .

If the model is unsure, the distribution is flat. The top 50 tokens might each have a 1% probability. The running total slowly creeps up, and it takes 90 tokens to finally reach the 90% threshold. Top-P keeps all 90 tokens, preserving the model's creative options. It acted just like $K=90$ .

Show Me the Code

Implementing Top-P requires a cumulative sum (cumsum) over the sorted probabilities.

1import numpy as np2
3def top_p_sampling_mask(probs: np.ndarray, p: float) -> np.ndarray:4    # 1. Sort the probabilities descending, keeping track of original indices5    sorted_indices = np.argsort(probs)[::-1]6    sorted_probs = probs[sorted_indices]7    8    # 2. Calculate the cumulative sum9    cumulative_probs = np.cumsum(sorted_probs)10    11    # 3. Find where the cumulative sum exceeds our threshold P12    # We shift the mask by 1 so we always include the token that crosses the line13    cutoff_mask = cumulative_probs > p14    cutoff_mask[1:] = cutoff_mask[:-1].copy()15    cutoff_mask[0] = False16    17    # 4. Map the sorted mask back to the original array layout18    original_layout_mask = np.zeros_like(cutoff_mask)19    original_layout_mask[sorted_indices] = cutoff_mask20    21    # 5. Set everything beyond the nucleus to 022    safe_probs = probs.copy()23    safe_probs[original_layout_mask] = 0.024    25    # 6. Re-normalize26    return safe_probs / np.sum(safe_probs)27
28# Example: Confident distribution29confident_probs = np.array([0.80, 0.15, 0.02, 0.01, 0.01, 0.01])30print(top_p_sampling_mask(confident_probs, p=0.90).round(3))31# -> [0.842 0.158 0. 0. 0. 0.] (Kept only 2 tokens!)32
33# Example: Unsure (Flat) distribution34unsure_probs = np.array([0.25, 0.25, 0.20, 0.20, 0.05, 0.05])35print(top_p_sampling_mask(unsure_probs, p=0.90).round(3))36# -> [0.278 0.278 0.222 0.222 0. 0.] (Kept 4 tokens to reach the 90% threshold)

Notice how the exact same algorithm dynamically kept 2 tokens in the first scenario and 4 tokens in the second scenario, perfectly matching the model's confidence level.

Watch Out For

Setting Top-P too low

If you set $P$ to a very low number (e.g., $P = 0.10$ ), the cumulative sum will almost always trigger on the very first token. The algorithm will instantly throw away everything else, and the model will behave exactly like greedy decoding, losing all of its creativity. Standard values for Top-P are usually between 0.85 and 0.95.

The Quick Version

Top-P (nucleus) sampling is a dynamic alternative to the rigid Top-K algorithm.
It sums the probabilities of the most likely tokens until the total reaches a user-defined threshold $P$ (e.g., 90%).
If the model is confident, it hits 90% quickly, keeping only 1 or 2 safe tokens.
If the model is unsure, it takes many tokens to hit 90%, keeping a wide variety of creative options.
This dynamic adaptation is why Top-P is the industry standard safety net used by almost all commercial LLM APIs today.