Skip to content

← Modern Techniques step 9 of 10

Medium Research

Nucleus (Top-P) Filtering

Keep the highest-probability tokens until their combined probability reaches a threshold, then rescale the retained probabilities to sum to one. This is the filtering stage of nucleus sampling; do not draw a token in this exercise.

Your task

Implement nucleus_filter(probs, p).

  • probs is a nonempty floating-point vector of finite, nonnegative probabilities summing to one, with one entry per vocabulary token.
  • p is a threshold with 0 < p <= 1.
  • Return a vector with the same shape, dtype and device, leaving probs unchanged. Keep vocabulary positions in their original order.

Rank tokens by descending probability, breaking ties by lower token ID. Keep the shortest prefix of that ranking whose cumulative probability is at least p. Include the token that reaches or crosses p, but no later token is needed once the threshold is reached. Zero all excluded positions and divide the retained probabilities by their total.

For probs [0.5, 0.3, 0.2] and p=0.8, keep the first two tokens. Their total is 0.8, so return [0.625, 0.375, 0]. Do not divide by p in general: for p=0.7 the same two tokens survive and their actual total is still 0.8.

For probs [0.1, 0.1, 0.8] and p=0.9, rank IDs as [2, 0, 1]. Keep IDs 2 and 0, then return approximately [0.1111111, 0, 0.8888889] in original vocabulary order. Sorting is for choosing the set, not the output layout.

Build the cutoff mask

Use a stable descending sort so equal probabilities keep their original ID order. Compute cumulative sums of the sorted probabilities. A token is retained when the cumulative mass before it is less than p.

Build that preceding-mass vector by placing zero before all but the last cumulative sum. Avoid subtracting each probability from its cumulative sum: rounding during that subtraction can move a value across the cutoff. Threshold comparisons use the input dtype, so values extremely close to the boundary are subject to floating-point rounding.

If the largest probability already reaches p, only that token is retained. With p=1, retain the full probability mass. Excluded entries must stay exactly zero, and the retained entries keep their relative proportions.

Unlike top-k, the number of retained tokens depends on how concentrated the probabilities are. This filtering rule does not itself guarantee better text; a generation loop would draw from the returned distribution as a separate step.