Nucleus (Top-p) Sampling - LLM Generation & Decoding

Nucleus (Top-p) Sampling

What you'll build

Sample from the nucleus: the smallest set of most-likely tokens whose total probability reaches p. Return the sampled token and the filtered, renormalized distribution.

Background

Top-k keeps a fixed number of tokens. Nucleus sampling (Holtzman et al., 2019) keeps a fixed amount of probability mass instead, so the cutoff adapts: after "The capital of France is", one token holds 95% of the mass and the nucleus is tiny; after "My favourite food is", the distribution is flat and the nucleus is large. Most chat APIs expose top_p (often 0.9 to 0.95) alongside temperature.

Algorithm

1. probs = softmax(logits / temperature) (stable).
2. Sort the probabilities in descending order and take the cumulative sum c.
3. Keep the sorted tokens up to and including the first one where c >= p (the token that crosses the threshold is in the nucleus).
4. Zero out every other token and renormalize the kept ones to sum to 1.
5. Sample one token from the result.

Inputs

• logits: float array (vocab_size,)
• p: float in (0, 1]
• temperature: float > 0, default 1.0 (applied before the nucleus is computed)

Expected Return Format

A tuple (token, filtered_probs):

• token: Python int, an index inside the nucleus
• filtered_probs: float array (vocab_size,), non-zero exactly on the nucleus, sums to 1

Edge cases and conventions

• The most likely token is always kept, even if its probability alone exceeds p.
• p = 1.0 keeps every token with non-zero probability.
• Keep the np.random.seed(42) line and sample with np.random.choice(len(probs), p=probs).
• Equivalent rule: keep a sorted token if the mass before it, c - prob, is still < p.

Constraints

NumPy only: np.argsort, np.cumsum, np.searchsorted (or a boolean mask).

Follow-up

• In practice, frameworks apply temperature, then top-k, then top-p. Does the order of top-k and top-p matter?
• What goes wrong with top-p when the model is over-confident and wrong? (See min-p sampling for one fix.)

Input: top_p_sampling(np.array([2.0, 1.0, 0.5, 0.1, -1.0]), p=0.8)
Output: filtered_probs = [0.6285, 0.2312, 0.1402, 0.0, 0.0]

Input: top_p_sampling(np.array([5.0, 3.0, 2.0, 1.0, 0.1]), p=0.5)
Output: (0, [1.0, 0.0, 0.0, 0.0, 0.0])

import numpy as np

def top_p_sampling(logits, p, temperature=1.0):
    """
    Nucleus (top-p) sampling from logits.

    Args:
        logits: Raw logits (vocab_size,)
        p: Cumulative probability threshold (0 < p <= 1)
        temperature: Temperature, applied first

    Returns:
        token: Sampled token index (int)
        filtered_probs: (vocab_size,), non-zero on the nucleus only, sums to 1
    """
    np.random.seed(42)
    # Your code here
    pass