Decoding and sampling

Problem

The decoder-only model hands back a score (a logit) for each of its tens of thousands of vocabulary entries. To write a sentence we must choose one token, feed it back, and choose again. The choice is not obvious. Always taking the top token is repetitive and can talk itself into a loop. Sampling straight from the probabilities sometimes picks a very unlikely word, and one odd word can derail everything after it. Different tasks want different behaviour: a code completion should be near-deterministic, a story should surprise you.

Intuition

Think of the model's output as a weighted die with one face per word. Decoding is how you roll it. Greedy always shows the heaviest face. Temperature reshapes the die: low temperature makes heavy faces heavier, high temperature evens the faces out. Top-k and top-p cut away the long tail of silly faces before rolling. Beam search does not roll at all; it keeps a few promising partial sentences and extends them in parallel, because the best first word does not always lead to the best sentence.

Real libraries apply these steps as a chain of logit processors; the order shown is the common one, not a law.

Mechanics

Greedy and beam search

Greedy decoding is argmax of the logits at the last position. It is fast and deterministic, but it optimises one step at a time, while what we want is the sequence with the highest total probability, the product of its step probabilities (the sum of their logs). Beam search keeps the B best partial sequences by total log-probability. At each step it extends every beam with every token, scores the results, and keeps the best B. With B=1 it is greedy. Beam search is a heuristic: it is not guaranteed to find the best sequence either, only a better one.

Temperature

Divide the logits by a temperature T before the softmax: p_i = exp(z_i/T) / Σ_j exp(z_j/T). As T→0 all the mass goes to the largest logit (greedy); T=1 is the model's own distribution; as T→∞ it approaches uniform. The name comes from the Boltzmann distribution in physics, where the same formula has the temperature as a divisor.

Top-k and top-p (nucleus)

Top-k keeps only the k highest-probability tokens and renormalises. Top-p keeps the smallest set of top tokens whose probabilities add up to at least p, so the number kept adapts: a confident step keeps a few tokens, an uncertain one keeps many. Both finish by dividing the survivors by their total so they sum to 1.

Repetition penalties

Models are prone to repeating themselves. A repetition penalty lowers the logit of tokens already generated. One common form (the idea comes from the CTRL model's "penalized sampling") divides a positive logit by a penalty θ>1 and multiplies a negative one by θ, so that the token gets less likely in both cases. This page uses that form with θ=1.5.

Constrained decoding

When the output must follow a format (valid JSON, a grammar, one of a few labels), set the logit of every token that would break the format to −∞ before the softmax. Those tokens get probability exactly 0, and the remaining ones are renormalised. Each step needs the set of allowed tokens, which a grammar or state machine supplies from what has been written so far.

Speculative decoding

Generation is slow because each token needs a pass through the big model. Speculative decoding has a small, cheap draft model propose several tokens, then lets the big target model score all of them in one pass. For a drafted token x with draft probability q(x) and target probability p(x), accept it with probability min(1, p(x)/q(x)). On rejection, draw instead from the residual distribution max(0, p−q), renormalised. The result: the final tokens follow the target's distribution exactly, so quality is unchanged and only speed varies. The check on this page confirms it for one step.

Worked example

Implementation

import torch
import torch.nn.functional as F

def sample_next(logits, temperature=1.0, top_k=None, top_p=None):
    # logits: [B, V] for the last position
    if temperature == 0:
        return logits.argmax(-1, keepdim=True)            # greedy
    logits = logits / temperature
    if top_k is not None:
        kth = torch.topk(logits, top_k).values[..., -1, None]
        logits = logits.masked_fill(logits < kth, float("-inf"))
    if top_p is not None:
        sorted_logits, order = torch.sort(logits, descending=True)
        probs = F.softmax(sorted_logits, dim=-1)
        before = probs.cumsum(-1) - probs                  # mass strictly ahead
        sorted_logits = sorted_logits.masked_fill(before >= top_p, float("-inf"))
        logits = torch.full_like(logits, float("-inf")).scatter(-1, order, sorted_logits)
    probs = F.softmax(logits, dim=-1)
    return torch.multinomial(probs, num_samples=1)        # [B, 1]

# repetition penalty on tokens already in `ids` ([B, T]):
def penalize(logits, ids, theta=1.5):
    seen = logits.gather(-1, ids)
    seen = torch.where(seen > 0, seen / theta, seen * theta)
    return logits.scatter(-1, ids, seen)

The mask for top-p keeps a token while the mass strictly ahead of it is still below p, so the token that crosses the threshold is kept, matching "smallest set whose mass reaches p". Ties in topk are not stable, so top-k can keep slightly more than k tokens here when logits tie at the cut. torch.multinomial accepts unnormalised weights, but softmax output is already normalised.

Tradeoffs

• Greedy and beam suit tasks with one right answer (translation, extraction). They are repetitive for open-ended text, and wide beams can favour short or generic outputs, so implementations usually add length normalisation.
• Temperature trades diversity for reliability. It is not calibration: dividing logits does not make the model more correct, only more or less committed.
• Top-k versus top-p. A fixed k keeps too many tokens when the model is sure and too few when it is not. Top-p adapts to that, which is why it is the usual default. Both are heuristics with no guarantee of better text.
• Penalties stop loops but also penalise words that should repeat (names, code identifiers). Keep θ small.
• Constrained decoding guarantees format but can force the model into low-probability continuations if the constraint disagrees with what it wanted to write.
• Speculative decoding is lossless in distribution but only helps when the draft model is much cheaper and agrees with the target often (high α); a poor draft wastes compute. The speed-up also depends on how verification of γ tokens costs compared with one, which is a hardware question (see the KV cache).

Check yourself

Connections

References

• Holtzman, Buys, Du, Forbes and Choi. The Curious Case of Neural Text Degeneration (nucleus sampling). ICLR, 2020 (arXiv 1904.09751).
• Fan, Lewis and Dauphin. Hierarchical Neural Story Generation (top-k sampling). ACL, 2018 (arXiv 1805.04833).
• Leviathan, Kalman and Matias. Fast Inference from Transformers via Speculative Decoding. ICML, 2023 (arXiv 2211.17192).
• Chen et al. Accelerating Large Language Model Decoding with Speculative Sampling. 2023 (arXiv 2302.01318).
• Keskar, McCann, Varshney, Xiong and Socher. CTRL: A Conditional Transformer Language Model for Controllable Generation. 2019 (arXiv 1909.05858).
• Willard and Louf. Efficient Guided Generation for Large Language Models. 2023 (arXiv 2307.09702).
• PyTorch v2.15.0: categorical.py, multinomial, topk, softmax.