Top-k Sampling

Overview

Top-k sampling is a decoding strategy that restricts the model’s token selection to only the k most probable next tokens at each generation step. Instead of sampling from the full vocabulary of 30,000–100,000+ tokens, the model considers only the top k candidates, sets all other token probabilities to zero, and re-normalizes the remaining distribution before sampling.

This technique was introduced to solve a fundamental problem with unrestricted sampling: even with moderate temperature settings, there’s always a non-zero chance of selecting extremely low-probability tokens that produce nonsensical output. A token with 0.001% probability might be “xylophone” in the middle of a sentence about database design. Top-k eliminates these tail tokens entirely, creating a hard boundary on the pool of candidate tokens.

For AI engineers, top-k sampling is one of the three primary “knobs” for controlling text generation quality, alongside temperature and top-p (nucleus sampling). Understanding how these parameters interact — and when top-k is the right choice versus alternatives — is critical for tuning production LLM applications.

How It Works

The Algorithm Step by Step

Top-k sampling modifies the standard sampling process with a simple but effective filter:

1. Generate logits: The model produces raw logit scores for every token in the vocabulary.
2. Apply temperature: Logits are divided by the temperature parameter T (this step happens first in most implementations).
3. Compute softmax: The scaled logits are converted to probabilities via the softmax function.
4. Rank and truncate: Sort all tokens by probability in descending order. Keep only the top k tokens; set all others to probability 0.
5. Re-normalize: Scale the remaining k probabilities so they sum to 1.0, creating a valid probability distribution.
6. Sample: Randomly select one token from this truncated distribution, weighted by the re-normalized probabilities.

A Concrete Example

Consider generating the next token after “The capital of France is”. The model produces probabilities for its full vocabulary. Here’s what top-k=5 does:

Full vocabulary (simplified to top 8):
Token       | Probability
------------|------------
"Paris"     | 0.72
"the"       | 0.08
"a"         | 0.05
"known"     | 0.04
"called"    | 0.03
"located"   | 0.02
"beautiful" | 0.01
"xylophone" | 0.0001
... (99,992 more tokens with tiny probabilities)

After top-k=5 filtering and re-normalization:
Token       | Original | Re-normalized
------------|----------|-------------
"Paris"     | 0.72     | 0.783
"the"       | 0.08     | 0.087
"a"         | 0.05     | 0.054
"known"     | 0.04     | 0.043
"called"    | 0.03     | 0.033

Tokens like “beautiful,” “located,” and “xylophone” are completely eliminated. The model will sample from only these 5 tokens, weighted by their re-normalized probabilities.

The Fixed-k Problem

The fundamental limitation of top-k sampling is that k is a fixed number, but probability distributions vary dramatically between generation steps. Consider two scenarios:

Scenario A — Confident prediction: After “The Eiffel Tower is located in”, the model is 95% confident the next token is “Paris.” With k=50, you’re including 49 tokens that collectively account for only 5% of the probability mass — unnecessary noise.

Scenario B — Uncertain prediction: After “I really enjoy eating”, many completions are plausible: “pizza,” “sushi,” “pasta,” “fruit,” “chocolate,” etc. The probability might be spread across 200+ reasonable tokens. With k=50, you’re cutting off many perfectly valid completions.

This inflexibility is why top-p (nucleus) sampling was developed as a dynamic alternative — it adapts the number of considered tokens based on how the probability mass is distributed.

Implementation in Code

Here’s how top-k sampling works in practice with the Hugging Face Transformers library:

# Using Hugging Face Transformers
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

input_ids = tokenizer("The future of AI is", return_tensors="pt").input_ids

# Generate with top-k sampling
output = model.generate(
    input_ids,
    max_new_tokens=50,
    do_sample=True,       # Enable sampling (vs greedy)
    top_k=40,             # Consider only top 40 tokens
    temperature=0.8,      # Moderate temperature
)

print(tokenizer.decode(output[0]))

And with the OpenAI API (note: OpenAI doesn’t expose top-k directly but some providers do):

// Google Vertex AI / Gemini API supports top-k directly
const response = await model.generateContent({
  contents: [{ role: "user", parts: [{ text: prompt }] }],
  generationConfig: {
    topK: 40,
    temperature: 0.8,
  },
});

Interaction with Temperature and Top-p

In the sampling pipeline, these parameters are typically applied in this order:

1. Temperature — reshapes the full probability distribution
2. Top-k — hard cutoff to the k most probable tokens
3. Top-p — dynamic cutoff based on cumulative probability
4. Sample — randomly select from the surviving tokens

When both top-k and top-p are applied, top-k acts as an upper bound on the number of tokens, and top-p further narrows within that set. For example, with k=100 and p=0.9, you’ll never consider more than 100 tokens, but if the top 20 tokens already cover 90% of the probability mass, only those 20 will be sampled from.

Practical Applications

Choosing the Right k Value

• k = 1: Equivalent to greedy decoding — always picks the most probable token. Fully deterministic.
• k = 10–20: Tight filtering, suitable for factual tasks where you want some variation but not wild departures.
• k = 40–50: The most commonly used range. GPT-2’s original implementation defaulted to k=40. Provides a good balance for general-purpose text generation.
• k = 100–200: Loose filtering, allows more creative token choices. Useful for brainstorming and creative writing.
• k = 0 (disabled): Most APIs interpret k=0 as “no top-k filtering.” The full vocabulary is considered (after temperature and top-p).

When to Use Top-k vs. Top-p

In modern practice, many engineers prefer top-p over top-k because it adapts to the distribution shape. However, top-k remains valuable in specific scenarios:

• Strict vocabulary control: When you need a hard guarantee on the maximum number of candidate tokens (e.g., for safety-critical applications where you want to limit the possibility space).
• Performance optimization: Top-k is computationally simpler than top-p since it only requires a partial sort rather than cumulative probability calculation.
• Combined with top-p: Using both gives you the best of both worlds — top-k sets an absolute ceiling, top-p handles the dynamic narrowing.

Comparing Decoding Strategies

Production Considerations

When deploying top-k sampling in production, keep these points in mind:

• Not all APIs expose top-k directly. OpenAI’s API does not have a top-k parameter; they rely on top-p and temperature. Google’s Vertex AI and Anthropic’s API do support it.
• The default k value varies by provider: some default to 40, others to 0 (disabled). Always set it explicitly.
• For reproducibility in testing, combine k with temperature=0 to get deterministic output regardless of k value.
• Monitor the distribution of selected token ranks — if your chosen k is too low, you’ll see frequent cases where the model was “forced” into suboptimal tokens.

Key Takeaways

• Top-k sampling restricts token selection to the k most probable tokens, eliminates all others, and re-normalizes the remaining probabilities before sampling.
• It provides a hard upper bound on the number of candidate tokens, making the generation process more predictable and preventing selection of extremely low-probability tokens.
• The fixed nature of k is both its strength (simplicity) and weakness (doesn’t adapt to varying distribution shapes). Top-p addresses this limitation.
• Common k values: 1 (greedy), 10–20 (conservative), 40–50 (balanced, most common), 100+ (creative).
• In the sampling pipeline, temperature is applied first, then top-k, then top-p. Using top-k and top-p together provides the most robust control.
• Not all LLM APIs expose top-k directly — check your provider’s documentation. Always set sampling parameters explicitly rather than relying on defaults.
• For production systems, the combination of moderate temperature (0.5–0.8) with top-k (40) and top-p (0.9) is a reliable starting point for most tasks.