Top-k & Top-p (Nucleus) Filtering

The problem

Even a good model gives a little probability to thousands of unlikely tokens, the tail of the distribution. Each one is rare, but together they add up, so sampling from the full distribution picks one of them now and then, and one bad token can derail the rest of the text. The two filters below cut the tail off before sampling.

Both work on the logits: a removed token gets the logit -np.inf. Since $e^{-\infty} = 0$, a softmax afterwards gives it probability 0 and renormalizes over the tokens that are left. So the sampling code of the previous exercise works unchanged after any filter.

Top-k: keep the $k$ most likely tokens

Find the $k$-th largest logit, keep every logit greater than or equal to it, and set the others to -inf.

Example. Logits [3, 1, 2, 0], k = 2. Sorted from the largest: 3, 2, 1, 0, so the 2nd largest is 2. Keep the logits $\geq 2$: the result is [3, -inf, 2, -inf].

With ties at the boundary this keeps more than $k$ tokens: for logits [3, 2, 2, 0] and k = 2, both 2s are kept. Hugging Face's TopKLogitsWarper does the same. If $k$ is at least the number of tokens, nothing is removed.

Top-p: keep the most likely tokens up to a total probability $p$

Top-k keeps the same number of tokens whatever the distribution looks like. Top-p (nucleus) sampling (Holtzman et al., 2020) adapts: it keeps the smallest group of most likely tokens whose total probability reaches $p$. When the model is confident that is a handful of tokens, when it is unsure, many.

1. Turn the logits into probabilities with a stable softmax.
2. Sort the tokens by probability, largest first.
3. Go down the sorted list adding up the probabilities, and stop at the first token where the total reaches $p$. Keep every token up to and including that one; remove the rest.

Example. Logits [3, 1, 2, 0], p = 0.9.

1. Probabilities: about [0.644, 0.087, 0.237, 0.032].
2. Sorted: token 0 (0.644), token 2 (0.237), token 1 (0.087), token 3 (0.032).
3. Running total: $0.644$ (below 0.9), $0.881$ (below), $0.968$ (reached at token 1). Keep tokens 0, 2 and 1.

The result is [3, 1, 2, -inf]. The most likely token is always kept, even if its probability alone is above $p$.

A convenient way to write step 3: a token is kept when the total probability of the tokens ranked above it is still below $p$. In the example these totals are 0 (token 0), 0.644 (token 2), 0.881 (token 1): all below 0.9, kept. For token 3 it is 0.968: removed.

What to implement

For a 1-D array of logits:

• top_k_filter(logits, k): the logits with everything below the $k$-th largest set to -inf.
• top_p_filter(logits, p): the logits with everything outside the nucleus set to -inf.

Requirements

• Return new arrays: do not modify logits.
• Compute the probabilities for top-p with a stable softmax.
• No Python loops over the vocabulary.

Hint

For top-p, order = np.argsort(-logits, kind="stable") lists the token indices from the most to the least likely. Compute the "total above" for the sorted probabilities with np.cumsum, decide which sorted positions to keep, then write that decision back to the original positions with keep[order] = ....