LM5-T6 — Top-k, Top-p and Stopping Criteria

Learning objectives

After this lesson, you should be able to:

• Apply top-k and top-p to a distribution by hand and in PyTorch, with renormalisation.
• Explain why top-p adapts to the shape of the distribution and top-k does not.
• List the common stopping criteria and their failure modes.
• Combine temperature, truncation and a repetition penalty in a generation step.

The idea: sample only from the plausible head

Suppose the model is about to continue “the robot picks up the red”. A few tokens (box, bin, cone) get almost all the probability. The remaining 49,000 tokens have a total of, say, 4% between them. Over a 200-token text, the chance that at least one of those is drawn is high, and once an odd token is in the context the model has to continue from it. Truncation sets the tail to zero and renormalises the rest before sampling.

Top-k

Top-k keeps the k most likely tokens and discards the rest (Fan et al., 2018). It is simple and cheap. Its weakness is the fixed size. If the model is almost certain (one token has 99%), a k of 50 still lets in 49 junk candidates. If the model is genuinely unsure (a hundred plausible continuations of an open question), k = 50 cuts good options away.

Top-p (nucleus sampling)

Top-p sorts the tokens by probability and keeps the smallest set whose probabilities add up to at least p (for example 0.9), the nucleus (Holtzman et al., 2020). The size of the set adapts: a confident model gives a nucleus of one or two tokens, an unsure one of dozens. This matches what we want: trust the model’s confidence when it has it, and allow variety when it does not.

Other tools

• Min-p keeps tokens whose probability is at least a fraction of the top token’s probability; another adaptive rule.
• Repetition and frequency penalties lower the logits of tokens that already occurred, to reduce loops. They are blunt: they can also penalise words that legitimately repeat.
• Beam search keeps several candidate continuations and picks the highest-probability sequence. It suits translation and short outputs but tends to produce bland, repetitive text for open-ended generation.

When to stop

Generation must end. The common criteria are: (1) the model emits the end-of-text token (Lesson LM2-T4), the preferred natural stop; (2) a maximum number of new tokens, a safety limit that always applies; (3) a stop string, such as a newline or a role marker, checked on the decoded text; (4) a limit on time or cost. A model that never emits the end token runs to the limit; a stop string that is also a normal part of the answer truncates it. Always return the reason generation stopped, so that a cut-off answer can be recognised.

The mathematics: truncation and renormalisation

Let p₁ ≥ p₂ ≥ … ≥ p_V be the probabilities sorted in descending order (ties broken by token index). For top-k with integer k and top-p with threshold p, the kept set S is

k

number of tokens kept by top-k.

p

cumulative probability threshold of top-p (0 < p ≤ 1); p = 1 keeps everything.

m

the size of the nucleus: the smallest number of top tokens whose total probability is at least p.

The new distribution is renormalised over the kept set:

If both rules are used, apply them in sequence (temperature first, then top-k, then top-p) and renormalise at the end; the order of the two truncations does not change the final set when both are computed on the same sorted list.

Example. Probabilities [0.50, 0.30, 0.10, 0.06, 0.04]. Top-k with k = 2 keeps {0.50, 0.30}, total 0.80; renormalised [0.625, 0.375, 0, 0, 0]. Top-p with p = 0.85: cumulative sums 0.50, 0.80, 0.90; 0.80 < 0.85 but 0.90 ≥ 0.85, so m = 3; renormalised [0.556, 0.333, 0.111, 0, 0]. With p = 0.8 the nucleus would be two tokens, since 0.80 ≥ 0.8.

Adaptivity. For a peaked distribution [0.97, 0.02, 0.007, 0.003], top-p with p = 0.9 keeps one token while top-k with k = 3 keeps three, two of them junk. For a flat distribution over 10 tokens of 0.1 each, top-p with p = 0.9 keeps nine tokens while top-k with k = 3 keeps only three.

Repetition penalty. A common form divides the logit of a token that already occurred by a factor θ > 1 if the logit is positive and multiplies it by θ if negative, so its probability always falls. The value θ = 1.1 to 1.3 is typical.

A worked example: top-k, top-p and a stopping loop

The functions below operate on one logit vector. They are the PyTorch versions of the Python you write in Lab LM5-L2.

import torch

def top_k_filter(logits, k):
    kth = torch.topk(logits, k).values[-1]
    return logits.masked_fill(logits < kth, float("-inf"))

def top_p_filter(logits, p):
    probs = torch.softmax(logits, -1)
    sorted_p, order = torch.sort(probs, descending=True)
    cum = torch.cumsum(sorted_p, -1)
    drop = (cum - sorted_p) >= p                      # drop tokens once the mass BEFORE them already reaches p
    drop_orig = torch.zeros_like(drop).scatter(0, order, drop)
    return logits.masked_fill(drop_orig, float("-inf"))

probs = torch.tensor([0.50, 0.30, 0.10, 0.06, 0.04])
logits = probs.log()                                  # logits whose softmax is exactly probs

print([round(v, 3) for v in torch.softmax(top_k_filter(logits, 2), -1).tolist()])      # [0.625, 0.375, 0.0, 0.0, 0.0]
print([round(v, 3) for v in torch.softmax(top_p_filter(logits, 0.85), -1).tolist()])   # [0.556, 0.333, 0.111, 0.0, 0.0]
print([round(v, 3) for v in torch.softmax(top_p_filter(logits, 0.75), -1).tolist()])   # [0.625, 0.375, 0.0, 0.0, 0.0]

peaked = torch.tensor([0.97, 0.02, 0.007, 0.003]).log()
print((top_p_filter(peaked, 0.9) > float("-inf")).sum().item(), (top_k_filter(peaked, 3) > float("-inf")).sum().item())   # 1 3

# a stopping loop: end token, maximum length and a stop string, with the reason returned
END = 0
def generate(next_logits, prompt, max_new=20, stop="."):
    ids = list(prompt)
    for _ in range(max_new):
        nxt = int(torch.argmax(next_logits(ids)))
        if nxt == END:
            return ids, "end token"
        ids.append(nxt)
        if stop is not None and nxt == 9:             # in this toy, token 9 stands for the stop string
            return ids, "stop string"
    return ids, "max length"

scripted = lambda ids: torch.nn.functional.one_hot(torch.tensor((len(ids) * 3) % 10), 10).float()
print(generate(scripted, [1]))                         # ([1, 3, 6, 9], 'stop string')
print(generate(scripted, [1], stop=None, max_new=4))   # ([1, 3, 6, 9, 2], 'max length')

The line drop = (cum - sorted_p) >= p expresses “keep a token if the mass before it is still below p”, which keeps the first token that crosses the threshold, exactly the nucleus definition. Using cum >= p directly would drop that crossing token.

Engineering implications and limitations

Order of operations. A conventional chain is: penalties on the logits, temperature, top-k, top-p, softmax, sample. State your order in documentation; different libraries differ in small ways.

Always keep at least one token. The first sorted token must always survive; implementations that drop on cum >= p without the shift can remove everything when the top token alone exceeds p.

Typical values. top-p 0.9 to 0.95 with temperature 0.7 to 1.0; top-k 40 to 50 as an additional guard. For deterministic tasks use greedy or a very low temperature, with no truncation.

Always set a maximum length. A model that does not emit its end token will otherwise run until memory or time is exhausted. Log how often the limit is hit; if often, the prompt or the model needs attention.

Stop strings and tokenization. A stop string may be split across tokens; check the decoded text, not single token IDs, and trim the stop string from the output if it should not be shown.

Repetition penalties are a last resort. Loops usually point to a too-low temperature, a too-narrow truncation or a model that needs more training. Penalties can damage code and structured text where repetition is correct.

Evaluate decoding on your task. The best settings differ by task. Evaluate with a fixed set of prompts and a metric or human review (Module 9), and keep the settings with the model's version.

For robots. For a model that outputs an action or a plan, restrict decoding to valid tokens (constrained decoding, Module 11) rather than relying on truncation to avoid invalid ones.

Common misconceptions

Engineering takeaways

1. Top-k keeps the k most likely tokens; top-p keeps the smallest set with cumulative probability at least p; both renormalise afterwards.
2. For [0.50, 0.30, 0.10, 0.06, 0.04]: k = 2 gives [0.625, 0.375]; p = 0.85 keeps three tokens: [0.556, 0.333, 0.111].
3. Top-p adapts to the model’s confidence; top-k does not.
4. Typical chain: penalties, temperature, top-k, top-p, softmax, sample; always keep the first sorted token.
5. Stop on the end token, a maximum length or a stop string, and report why generation stopped.

References and further reading

See Chapter 5 of the book (decoding strategies to control randomness: top-k sampling, and modifying the text generation function).

• Holtzman, A. et al. (2020). The Curious Case of Neural Text Degeneration. arXiv:1904.09751.
• Fan, A., Lewis, M. and Dauphin, Y. (2018). Hierarchical Neural Story Generation. arXiv:1805.04833.
• Holtzman et al. (2020), Neural Text Degeneration (research paper)
• Fan, Lewis and Dauphin (2018), Hierarchical Neural Story Generation (research paper)
• Build a Large Language Model (From Scratch), Chapter 5: Pretraining on Unlabeled Data (primary book reference)
• Build a Large Language Model (From Scratch), Appendix D: Adding Bells and Whistles to the Training Loop (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

That completes Module 5. Next, Module 6 adapts a pretrained GPT to a classification task.

Related content

Nothing below is required to finish this lesson.

• Before
  LM5-T5: Temperature and Sampling

  Turn logits into probabilities, scale them by temperature and sample.

• Lab
  LM5-L2: Compare Sampling Strategies

  Implement temperature, top-k and top-p sampling and compare their diversity and quality.