Top-p (Nucleus) Filtering

Turn a vector of next-token logits into the probability distribution that nucleus (top-p) sampling draws from.

Sampling from the full softmax lets rare, low-quality tokens through. Top-p sampling keeps only the smallest set of the most likely tokens whose total probability reaches p, and renormalises over that set. When the model is confident, the set is small. When it is unsure, the set grows.

1. Apply a temperature: probs = softmax(logits / temperature).
2. Sort the tokens by probability from high to low. Break ties by the lower token index.
3. Keep tokens from the top until the cumulative probability is >= p. The token that crosses p is kept.
4. Set every other probability to 0 and divide by the sum of the kept ones.

Function signature

def top_p_filter(logits, p, temperature=1.0):

Example

Input

logits = np.log(np.array([0.5, 0.3, 0.15, 0.05]))
top_p_filter(logits, p=0.75)

Output

array([0.625, 0.375, 0.   , 0.   ])

Reasoning

The softmax gives back [0.5, 0.3, 0.15, 0.05]. The top token alone has 0.5, which is below 0.75. Adding the second gives 0.8, which reaches it, so two tokens are kept. They are renormalised: 0.5 / 0.8 = 0.625 and 0.3 / 0.8 = 0.375.

Constraints and notes

• logits has shape (V,), and 0 < p <= 1.
• At least one token is always kept.
• A lower temperature makes the distribution sharper before the cut.
• Compare the cumulative sum with a small tolerance, such as 1e-12, so that a sum that should equal p exactly is not lost to rounding.