What is top-k sampling and how does it work?

What is top-k sampling? It is a decoding strategy that restricts the next-token prediction to the k most probable tokens at each step, discarding the long tail of the distribution before sampling. By truncating low-probability candidates, top-k reduces the chance of incoherent or hallucinated output while preserving stochasticity.

Unlike greedy decoding (which always picks argmax) or pure multinomial sampling (which considers the full vocabulary), top-k imposes a hard cap on the candidate set. The value of k directly controls the diversity–coherence trade-off: smaller k yields safer, more deterministic text; larger k allows more creative but riskier continuations.

How top-k sampling works

At each decoding step, the model produces a probability distribution P over the vocabulary V (typically 32k–256k tokens). Top-k sampling applies three operations in sequence:

1. Sort the vocabulary by descending probability P(token | context).
2. Truncate to the top k tokens. All tokens ranked k+1 through |V| receive zero probability mass.
3. Renormalize the remaining k probabilities so they sum to 1, then sample from this truncated distribution.

Formally, let V_k be the set of k tokens with highest probability. The sampling distribution becomes:

P_top-k(t | context) = P(t | context) / Σ_{t' ∈ V_k} P(t' | context)   if t ∈ V_k
                       0                                                  otherwise

The temperature parameter T is applied before truncation (standard practice in most implementations):

logits' = logits / T
probs = softmax(logits')

Then top-k filtering runs on probs. This order matters: temperature sharpens or flattens the original distribution, changing which tokens make the top-k cut.

Minimal implementation

import torch
import torch.nn.functional as F

def top_k_logits(logits: torch.Tensor, k: int) -> torch.Tensor:
    """Mask logits so only top-k remain. Returns modified logits."""
    if k <= 0:
        return logits
    # Get the k-th largest logit value per batch item
    values, _ = torch.topk(logits, k, dim=-1)
    # values[:, -1:] is the threshold (smallest of the top-k)
    threshold = values[:, -1:]
    # Mask everything below threshold to -inf
    return torch.where(logits < threshold, torch.full_like(logits, float('-inf')), logits)

def sample_top_k(logits: torch.Tensor, k: int, temperature: float = 1.0) -> torch.Tensor:
    logits = logits / temperature
    logits = top_k_logits(logits, k)
    probs = F.softmax(logits, dim=-1)
    return torch.multinomial(probs, num_samples=1)

This pattern — mask logits to -inf before softmax — is numerically stable and matches what you’ll find in Hugging Face generate, vLLM, and TensorRT-LLM.

Why top-k matters for production systems

Controlling tail risk

LLM vocabularies contain thousands of tokens that are grammatical but semantically nonsense in context (e.g., “the”, “a”, “banana” when completing a SQL query). The aggregate probability mass of this “long tail” can exceed 10–20% even when the top token has 40%+ probability. Pure multinomial sampling will occasionally draw from this tail, producing non-sequiturs or hallucinations. Top-k eliminates that risk by construction.

Predictable latency and memory

Top-k sampling avoids the overhead of nucleus (top-p) sorting when k is small and fixed. The topk kernel is highly optimized on GPU; it runs in O(k log k) or better with specialized implementations. For batch inference with fixed k (e.g., k=50), the per-step cost is deterministic — useful for latency SLAs.

Composition with other parameters

Top-k composes cleanly with temperature, repetition penalty, and logit bias. A common production recipe:

Applying top-p after top-k (the default in most libraries) means the nucleus operates on an already-truncated distribution. This is usually what you want: top-k removes the absolute garbage, top-p removes the marginal garbage.

Concrete example: completing a code snippet

Prompt:

def fibonacci(n: int) -> int:
    if n <= 1:
        return n
    return fibonacci(n - 1) + fibonacci(n - 2)

Assume the model’s next-token distribution (simplified):

With top_k=5, temperature=0.7:

1. Temperature 0.7 sharpens the distribution (higher peaks, lower valleys).
2. Top-5 tokens kept: \n, , #, def, if.
3. Remaining 93% of probability mass (including banana and 30k others) zeroed out.
4. Renormalized probabilities sum to 1.0 over the 5 candidates.
5. Sample — most likely \n or , producing valid Python continuation.

With top_k=200, temperature=1.0:

1. No sharpening.
2. Top-200 includes else, return, and many low-prob tokens.
3. banana (rank ~500) still excluded, but weirder continuations possible.
4. Higher variance in output — sometimes creative, sometimes broken.

With top_k=0 (disabled) + top_p=0.95:

1. Nucleus finds smallest set covering 95% cumulative probability.
2. If the head is sharp, this might be 3 tokens; if flat, could be 500+.
3. Variable candidate set size per step — less predictable latency.

Common misconceptions

“Top-k and top-p do the same thing”

False. Top-k is a fixed-cardinality filter; top-p is a fixed-mass filter. On sharp distributions (code, JSON), top-p may keep only 2–3 tokens while top-k=50 keeps 50. On flat distributions (creative writing), top-p=0.95 may keep 200+ tokens while top-k=50 keeps 50. They answer different questions:

• Top-k: “Never consider more than k options.”
• Top-p: “Never consider options that collectively exceed (1-p) probability.”

“Top-k=1 is the same as greedy decoding”

Mostly true, but not identical. With top_k=1 and temperature > 0, you still run softmax + multinomial over a single token — which always returns that token. The difference is numerical: greedy uses argmax directly on logits; top-k=1 goes through softmax. Results match unless logits contain inf/nan or you’re using float8 quantization where softmax behaves differently. For practical purposes, treat them as equivalent.

“Higher top-k always means more creative”

Only up to a point. Once k exceeds the “effective vocabulary” — the number of tokens with non-negligible probability in context — increasing k further adds only noise tokens. For typical English contexts at temperature 0.7, the effective vocabulary is often 30–80 tokens. Setting top_k=500 vs top_k=100 rarely changes output quality but does increase the chance of sampling a hallucination token that happened to sneak into the top 500.

“Top-k prevents repetition”

It does not. Repetition arises from the model assigning high probability to recently generated tokens (a known transformer pathology). Top-k merely limits which high-probability tokens are eligible. If the model puts 0.6 probability on “the” and 0.3 on “the” again, top-k=10 keeps both. Use repetition_penalty, frequency_penalty, or presence_penalty for actual repetition control.

“You must choose either top-k or top-p”

Most production stacks apply both sequentially: top-k first (hard cap), then top-p (adaptive mass cap). This combination — sometimes called “top-k + top-p” or “truncated nucleus” — gives you a safety ceiling (k) and a quality floor (p). The Hugging Face generate default (top_k=50, top_p=1.0) effectively disables top-p; setting top_p=0.95 enables the combo.

Practical tuning guidelines

Tune in this order:

1. Set temperature for desired randomness level.
2. Set top_k to eliminate the long tail (start at 50).
3. Add top_p=0.95 if you want adaptive truncation on top of the hard cap.
4. Add penalties only if you observe repetition loops.

Interaction with provider defaults

Different inference endpoints ship different defaults. OpenAI’s chat.completions does not expose top-k directly (only top_p). Anthropic exposes top_k. Open-source servers (vLLM, TGI, TensorRT-LLM) expose both. If you’re routing traffic across multiple providers — as you might through a gateway that normalizes the OpenAI-compatible interface — verify each provider’s default top_k and whether they apply top-p before or after. The order changes the effective distribution.

Summary

Top-k sampling is a hard cardinality constraint on the next-token candidate set. It removes the long tail of low-probability tokens that cause hallucinations and non-sequiturs, while preserving stochasticity within the retained head. The parameter k is a direct knob on the diversity–coherence trade-off: lower for code and structured output, higher for creative tasks. Combine with temperature (applied first) and optionally top-p (applied after) for a complete decoding strategy. Avoid the misconception that top-k alone controls repetition — it doesn’t — and don’t assume higher k always improves quality past the effective vocabulary size.