Temperature, Top-k & Top-p Sampling

Temperature changes the shape of the probability distribution used to select an output token. Dividing logits by a lower positive temperature concentrates probability; raising temperature makes the distribution flatter.

Temperature, truncation, and sampling are separate operations. Top-k or top-p chooses which candidates remain eligible, then a token is drawn from the renormalized distribution. A temperature of zero is usually handled as a special greedy-decoding setting.

Temperature scales the logits before softmax. Low temperature ($T \to 0$) sharpens the distribution toward the most likely token (approaching greedy); high temperature ($T > 1$) flattens it, giving lower-probability tokens more chance. Diversity and factual accuracy depend on the model and task, not temperature alone. Top-k and top-p then truncate: top-k keeps the k highest-probability tokens; top-p keeps the smallest set whose probabilities sum to at least p, adapting how many candidates survive to how confident the model is.

Given logits $z$, temperature $T$ rescales before softmax:

$T = 1$ is the raw distribution; $T \to 0^+$ concentrates mass on the maximum-logit tokens (a unique maximum gives the greedy limit); $T > 1$ flattens toward uniform.

Top-k keeps the set $V_k$ of the $k$ highest-probability tokens, zeros the rest, and renormalizes. Top-p (nucleus) keeps the smallest set $V_p$ such that

then renormalizes and samples from it. The key difference: top-k uses a fixed candidate count regardless of confidence; top-p uses a dynamic count: few candidates when the model is sure, many when it's uncertain. Beam search isn't sampling at all: it keeps the $b$ highest-scoring candidates among expansions of the current beam, often with length normalization. This is an approximate search, not a guarantee of the globally most probable sequence. It is used in directed generation, but can yield dull or repetitive open-ended text.

Key takeaway: Decoding controls how much you trust the distribution's peak versus its tail: temperature reshapes it, top-k/top-p truncate it, beam search abandons sampling entirely for high-probability sequences.

Greedy and beam search are deterministic selection rules given identical logits and tie-breaking. Full serving systems may still vary because of numerical kernels, batching, or model changes. Neither method guarantees a correct answer.

For one finite logit vector, apply temperature, then top-k, then nucleus truncation. The nucleus must include the token that first reaches or crosses the requested probability mass.

1import torch
2import torch.nn.functional as F
3
4def sampling_probs(logits: torch.Tensor, temperature=1.0, top_k=0, top_p=1.0):
5    if logits.ndim != 1 or not torch.isfinite(logits).all():
6        raise ValueError("Expected one finite logit vector")
7    if temperature <= 0 or not 0 < top_p <= 1 or not 0 <= top_k <= logits.numel():
8        raise ValueError("Use positive temperature, 0 < top_p <= 1, and a valid top_k")
9    scaled = logits / temperature
10    if top_k > 0:
11        keep = torch.topk(scaled, top_k).indices
12        filtered = torch.full_like(scaled, float("-inf"))
13        filtered[keep] = scaled[keep]
14        scaled = filtered
15    probs = F.softmax(scaled, dim=-1)
16    if top_p < 1.0:
17        sorted_probs, idx = torch.sort(probs, descending=True)
18        # Remove a token only if the mass BEFORE it already reaches p.
19        mass_before = torch.cat([sorted_probs.new_zeros(1), sorted_probs.cumsum(0)[:-1]])
20        sorted_probs[mass_before >= top_p] = 0.0
21        probs = torch.zeros_like(probs).scatter(0, idx, sorted_probs)
22        probs = probs / probs.sum()
23    return probs
24
25
26def sample(logits: torch.Tensor, temperature=1.0, top_k=0, top_p=1.0) -> int:
27    if temperature == 0:                         # explicit greedy convention
28        return logits.argmax().item()
29    probs = sampling_probs(logits, temperature, top_k, top_p)
30    return torch.multinomial(probs, 1).item()
31
32logits = torch.randn(50000)
33token = sample(logits, temperature=0.8, top_p=0.9)

1. Conceptual: What does temperature do to the next-token distribution? (Scales logits before softmax: low T sharpens toward greedy, high T flattens toward uniform/more random.)
2. Implementation: What's the difference between top-k and top-p sampling? (Top-k keeps a fixed number of candidates; top-p keeps a dynamic number whose cumulative probability reaches p: adapting to the model's confidence.)
3. Applied: When would you prefer greedy or beam search over sampling? (When selecting high-scoring outputs is more useful than sampling diverse ones, as in some directed generation tasks. Translation and summarization have many valid answers; greedy and beam search do not guarantee the best sequence.)
4. Systems-level: What does temperature = 0 give you, and why is it useful? (Usually a greedy-decoding convention. The selection rule is deterministic for fixed logits and tie-breaking, but an entire serving system need not be bit-for-bit reproducible.)
5. Failure modes: Why can beam search produce dull or repetitive text in open-ended generation? (It optimizes for high total probability, which favors safe, generic continuations and degenerate repetition over natural diversity.)

From memory: write the temperature-scaled softmax, explain top-k vs. top-p in one sentence each, and name a task where greedy/beam beats sampling. Check against Stage 3.

Sources and further reading

• Holtzman et al., The Curious Case of Neural Text Degeneration: nucleus sampling and comparisons with maximization-based decoding.