Top-k sampling is a stochastic text-generation technique that limits each next-token choice to the k highest-probability candidates. It samples from this reduced set instead of the model’s full vocabulary, preserving randomness while favoring coherent continuations. section 4.1.3
At each generation step, the model produces a probability distribution, the method retains the top k tokens, and a token is randomly selected from that subset. Applied after softmax, it modifies probabilities rather than logits. section 4.1.3
The k parameter controls the main trade-off:
- A smaller k narrows the choice to the most likely continuations, emphasizing exploitation and coherence. section 4.1.3
- A larger k admits more high-probability alternatives, increasing sampling variety while retaining a probability-based filter. section 4.1.3
- Unlike greedy search, top-k sampling does not always select the single most likely token; unlike beam search, it maintains one stochastic generation path. section 4.1 section 4.1.3
Hugging Face Transformers configuration
Enable sampling with do_sample=True and set the candidate limit with top_k. For example: model.generate(input_ids=input_ids, max_new_tokens=128, do_sample=True, top_k=50, temperature=1.0, num_return_sequences=3, pad_token_id=tokenizer.eos_token_id). section 4.1.3
Figure 4.2 If we assume we want to model the probability of the next word to choose for our text generation, we would do P(w |mid "The" ), which will then be followed by the probability P(w |mid "The," "movie") and so on. The choice of the model will be limited to the two k-probable words, so only movie or pie will be considered for the first probability and was or is for the next choice, respectively.
Visualizes successive next-token choices, showing how k=2 restricts candidates such as movie or pie, then was or is. section 4.1.3
