Top-p: let the candidate count follow uncertainty

Top-p, or nucleus sampling, keeps the smallest highest-probability token set whose cumulative probability reaches a chosen threshold. It then renormalizes that set and samples from it.

Unlike top-k, the candidate count changes with the distribution. A concentrated prediction needs few candidates; a flatter one needs more to cover the same probability mass.

Keep the token that crosses the boundary

This NumPy implementation sorts probabilities and retains the threshold-crossing token, not just those entirely below the cutoff.

import numpy as np

def nucleus(probabilities, threshold):
    order = np.argsort(-probabilities)
    cumulative = np.cumsum(probabilities[order])
    count = min(len(order), np.searchsorted(cumulative, threshold) + 1)
    keep = order[:count]
    filtered = np.zeros_like(probabilities)
    filtered[keep] = probabilities[keep]
    return filtered / filtered.sum()

peaked = nucleus(np.array([.75, .15, .05, .03, .02]), .8)
flat = nucleus(np.array([.2, .2, .2, .2, .2]), .8)

At p=0.8, the peaked distribution keeps two tokens while the flat distribution keeps four. The cutoff applies to cumulative probability, not to each token individually.

Temperature can change the probabilities before this filter, so the same p need not retain the same tokens at different temperatures. This controls diversity in token selection; it does not set the probability that the resulting answer is true.