Skip to content

← JAX Stochasticity step 10 of 25

Hard Primitives

Nucleus (Top-p) Masking

Jump to the task

Why this matters

Top-p masking chooses candidates by probability mass rather than a fixed count. Sort tokens from most to least likely, then keep enough to reach the target probability p. A concentrated distribution may need only one token; a flatter distribution may need many.

This function prepares logits for sampling. It does not draw a token. The choice of p changes the candidates available to a sampler, but it does not guarantee coherent text or make top-p best for every task.

Worked mini-example

For logits [1, 2, 3, 4], the two largest scores account for about 88.1% of the probability. That is below p = 0.9, so keep the third as well. Their combined probability is about 96.8%.

import jax
import jax.numpy as jnp

logits = jnp.array([1.0, 2.0, 3.0, 4.0])
p = 0.9
sorted_logits = jnp.sort(logits)[::-1]  # [4, 3, 2, 1]
probs = jax.nn.softmax(sorted_logits)
# Approximately [0.644, 0.237, 0.087, 0.032].

cumulative = jnp.cumsum(probs)
before = jnp.concatenate([jnp.zeros_like(probs[:1]), cumulative[:-1]])
# Approximately [0, 0.644, 0.881, 0.968].
keep = (before < p) | (p == 1.0)
# [True, True, True, False]

threshold = jnp.where(keep, sorted_logits, jnp.inf).min()  # 2.0
masked = jnp.where(logits >= threshold, logits, -1e9)
# [-1e9, 2, 3, 4], in the original token order.

Problem

Implement nucleus_top_p_mask(logits, p).

  • logits is a nonempty, finite float32 array of shape (K,).
  • p is a scalar probability with 0 < p <= 1.
  • For p < 1, find the shortest descending prefix whose softmax probabilities sum to at least p. Always include the first token, and include the token that reaches or crosses the target.
  • Retain every original logit at least as large as the last score in that prefix. This includes all ties at the cutoff, so the result may contain more tokens than the shortest prefix. Equal logits all survive.
  • For p == 1, retain every logit, including those whose computed softmax probability underflows to zero.

Return a float32 array of shape (K,) in the original order. Replace rejected entries with -1e9 and leave retained entries unchanged.

Common pitfalls

  • Compare p with the probability mass before each token. Comparing with the mass after it would drop the token needed to reach the target.
  • Build the preceding cumulative sums by shifting the cumulative array and inserting zero. Subtracting each probability from its cumulative sum can introduce additional rounding error near a cutoff.
  • Handle p == 1 explicitly: floating-point cumulative sums can reach one before the final token. Cutoffs very close to a cumulative sum can otherwise depend on floating-point rounding.
  • -1e9 is a finite output marker, not negative infinity. It is convenient for these JSON fixtures but does not guarantee zero sampling probability. If retained logits are even smaller, a rejected entry can receive more probability than a retained one. A downstream sampler needs a true exclusion mask or negative infinity to guarantee exclusion.