Top-k Sampling

GPT-2

Medium

Temperature scaling and top-k sampling

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Top-(k) sampling restricts the next-token draw to the (k) highest logits, then samples from the renormalized softmax of those (k) values. The GPT-2 paper uses it as the default generation method in the appendix and in the summarization experiment. Section 3.6 generates 100 tokens with top-(k) random sampling (Fan et al., 2018) at (k=2) after a TL;DR: hint, arguing that this “reduces repetition and encourages more abstractive summaries than greedy decoding.” Figure 5 and Tables 7–13 use (k=40) for qualitative WebText completions.

Temperature is not given a numbered equation in the paper, but the official sampler always divides logits by a temperature before truncation. The two knobs are therefore a single pipeline: sharpen or flatten the distribution, keep the top (k) tokens, mask the rest, sample. The released sample.py later also implements nucleus (top_p) sampling; that is a repo addition, not a GPT-2 paper method.

How it works

Let (\ell\in\mathbb{R}^{V}) be the last-position logits (logits[:, -1, :] after the optional [:n_vocab] clip). Temperature (T>0) defines

\tilde\ell=\frac{\ell}{T}.

(T<1) sharpens; (T>1) flattens; (T=1) is the raw model distribution. The official code uses logits / tf.to_float(temperature) before top-(k), so the ranking is the same as ranking (\ell) whenever (T>0), but the probabilities among the kept tokens change.

Let (S_k) be the index set of the (k) largest entries of (\tilde\ell) (ties broken as in tf.nn.top_k). Top-(k) truncation is

\ell'_v=\begin{cases} \tilde\ell_v & v\in S_k,\\ -\infty & v\notin S_k, \end{cases} \qquad \hat y\sim\mathrm{Categorical}\!\left(\mathrm{softmax}(\ell')\right).

top_k_logits implements (-\infty) as -1e10. It takes the (k)th-largest value values[:, -1] and replaces every logit strictly below that threshold. k=0 is a special case meaning no truncation (full softmax). k=1 leaves a single finite logit, so sampling collapses to greedy.

After optional top_p_logits, tf.multinomial draws one id. That id is concatenated onto the context; the next model() call uses the KV present as past. Unconditional generation seeds <|endoftext|> and drops that first token from the printed sample. Conditional generation encodes a user string and continues for length steps.

Why (k=2) versus (k=40): small (k) stays near the mode (summaries that still vary a little), large (k) restores most of the nucleus of a high-entropy language model. Figure 5 reports that (k=40) samples overlap WebText 8-grams less than held-out articles, i.e. this is not a copy-out-the-training-set procedure by default. The paper does not prescribe a temperature for those tables; the code default is temperature=1.

Top-(k) is a hard count cutoff. It does not adapt to a flat or peaky distribution the way nucleus sampling (top_p in the same file) does. Combining both, as sample_sequence allows, first drops everyone outside the top (k), then drops the tail of the remaining mass — but again, only top_k is discussed in the paper.

Official code

src/sample.py: top_k_logits and the body of sample_sequence (temperature divide, then top-(k), then top-(p), then multinomial). CLI defaults and the “40 is generally a good value” note are in src/generate_unconditional_samples.py and src/interactive_conditional_samples.py.

Watch-outs

• Applying top-(k) before dividing by (T) does not change the kept set for (T>0), but applying it after a negative or zero temperature is undefined. The graph assumes (T>0).
• k=0 means “off,” not “empty support.” Treating 0 like “keep zero tokens” empties the distribution.
• The threshold is the (k)th score, inclusive. Using a strict top-(k) that drops ties at the boundary can disagree with tf.nn.top_k.
• Truncating in probability space (keep p > τ) is nucleus sampling, not top-(k). Section 3.6 cites Fan et al., which is count-based.

Sources

• Paper: Language Models are Unsupervised Multitask Learners (§3.6, Fig. 5, sample tables)
• Code: openai/gpt-2 (src/sample.py)