Top-K and Top-P Sampling

Your program will build the full sampling distribution used by real chat APIs: temperature first, then top-k, then top-p, producing the exact probabilities a token would be drawn from. The draw itself is the one you wrote in the last lesson.

Why clip the tail

Greedy decoding is repetitive. Sampling from the full softmax fixes that but has a different problem: a vocabulary of 50,000 tokens has a long tail of nonsense, each token with a tiny probability, and together the tail can hold 5 or 10 percent of the mass. Over a few hundred tokens you will certainly draw from it. Top-k and top-p cut the tail before sampling.

The three filters, in order

For one set of logits z:

1. Temperature. Divide every logit by T. Smaller T sharpens, larger T flattens (lesson 12).
2. Top-k. Keep the k largest logits and drop the rest (their probability becomes 0). k = 0 means no limit. Ties for the last place go to the lower index.
3. Softmax over the survivors, so their probabilities sum to 1.
4. Top-p (nucleus). Sort the survivors from most to least probable and add them up from the top. Stop at the first token where the running total reaches p, and keep that token too. Drop everything else and divide the kept probabilities by their total so they sum to 1 again. p = 1 means no limit.

The order matters: top-p looks at the probabilities left after top-k has renormalised them.

Worked example

Logits [4, 3, 2, 1, 0], temperature 1:

softmax          [0.6364, 0.2341, 0.0861, 0.0317, 0.0117]

k = 3:           keep the three largest: softmax over [4,3,2]
                 -> [0.6652, 0.2447, 0.0900, 0, 0]

p = 0.8:         running totals 0.6364, 0.8705
                 0.8705 >= 0.8 at the second token, so keep two
                 -> divide by 0.8705: [0.7311, 0.2689, 0, 0, 0]

Top-k always keeps 3 here. Top-p kept 2 because the first token already holds most of the mass; with logits [8, 0, 0, 0, 0] and p = 0.7 it would keep exactly one, and with a flat distribution it would keep many. That adaptivity is why top-p tends to read more naturally. Typical settings: T = 0.2 and p = 0.95 for code, T = 0.9 and p = 0.95 for creative writing, k in the range 40 to 100 when it is used.

Details that bite

• Ties. When two logits are equal, which one survives top-k must not depend on the sort routine. "Lower index first" gives one answer everywhere.
• The boundary token is kept. With probabilities [0.5, 0.5] and p = 0.5, the first token alone reaches the target, so the nucleus is one token. A nucleus that stops one token early would be smaller than the mass you asked for; one that stops one late is larger.
• Renormalise. After dropping tokens the probabilities no longer sum to 1; sampling from them without dividing by the total biases the draw.
• k larger than the vocabulary is the same as no limit.

Cost is small: sorting is O(V log V) per token for a vocabulary V, a rounding error next to the forward pass.

Common mistakes

• Applying top-p to the original softmax instead of to the distribution left after top-k.
• Setting dropped logits to 0 instead of excluding them; a logit of 0 still gets probability.
• Stopping the nucleus before the token that crosses p, or after one more.
• Forgetting that k = 0 and p = 1 mean "off", and so filtering with k = 0 to nothing.

Your exercise

Implement the top-k selection and the top-p filter. The starter divides by the temperature, runs the masked softmax between your two steps, and prints the result.