Prefix sums and binary search to top-p sampling

The "nucleus" in nucleus sampling is not a new idea - it is a running total and a binary search, two of the first tools you learned, pointed at a probability distribution.

A language model's final act is to hand you a probability for every token in its vocabulary. Top-p (nucleus) sampling keeps only the smallest set of the most likely tokens whose probabilities add up past a threshold $p$, then draws from that set. Strip away the ML vocabulary and the recipe is pure data structures: sort the probabilities, take a running total, and binary-search for the cutoff. This lesson makes that equivalence something you can grab with the mouse.

From bars to a rising staircase

Start with the next-token probabilities and sort them descending - biggest first. Now walk left to right keeping a running total. That running total is a prefix sum, and the array of prefix sums over a probability list has a name in statistics: the cumulative distribution function, or CDF. Because probabilities are non-negative, the CDF only ever climbs - it is a monotonically increasing staircase from $0$ up to $1$. Each step's height is exactly one token's probability, so the tallest bar is the tallest step. That single fact - the CDF is sorted - is what lets a binary search do the rest.

Drag the threshold, watch the nucleus

Below are ten sorted next-token probabilities for the prompt "The weather today is ___". The teal bars are the probabilities; the amber staircase is their CDF - notice each riser is the same height as the bar beneath it, because a prefix sum just stacks the bars. Drag the dashed $p$ line up and down (grab anywhere on the plot, or use the slider). The cutoff marker jumps to the first bar where the staircase crosses your line - that is the binary search landing. Then trace the search step by step, or throw a dart to sample.

One dart: inverse-CDF sampling

Finding the nucleus is only half the job - we still have to pick a token from it. The trick reuses the very same staircase. Draw a random height $u$ uniformly along the vertical axis (inside the nucleus's mass), then find the first token whose CDF rises above $u$. That is inverse-transform sampling, and it is again a binary search into the CDF. Because a wider bar owns a taller slice of the staircase, it catches more darts - so tokens are drawn in exact proportion to their probability. Press Throw a dart to watch $u$ drop and slide into a token's band.

import numpy as np

def top_p_sample(probs, p):
    order = np.argsort(probs)[::-1]     # 1. sort tokens high -> low   O(V log V)
    probs = probs[order]
    cdf   = np.cumsum(probs)            # 2. prefix sums = the CDF      O(V)
    m     = np.searchsorted(cdf, p)      # 3. binary search the cutoff   O(log V)
    keep  = order[: m + 1]              #    the nucleus (smallest prefix)
    u     = np.random.random() * cdf[m]  # 4. inverse-CDF: a random height
    j     = np.searchsorted(cdf, u)      #    first token whose CDF >= u
    return order[j]                     #    the sampled token

Top-p versus top-k. Top-k (last lesson) keeps a fixed count of tokens - a job for a heap or partial sort. Top-p keeps a fixed mass, so its count adapts: many tokens when the model is unsure, few when it is confident. Both cut the tail; top-p just lets the tail's own shape decide where the knife falls. In practice people stack them - a small top-k as a safety cap, then top-p inside it.

Check yourself