Gen AI

Top-K Sampling

Instead of risking a random draw from the thousands of terrible words in the long tail of the distribution, forcefully truncate the list to only the K best options before rolling the dice.

Why Does This Exist?

When an LLM generates text using standard sampling, it rolls a weighted die across its entire vocabulary (often 100,000+ words). While the top 5 or 10 words might hold 99% of the probability mass, there is still a tiny, lingering 1% chance distributed across 99,990 terrible, grammatically incorrect, or nonsensical words.

Normally, the odds of drawing one of those terrible words is low. But if you generate a 2,000-word essay, you are rolling that die 2,000 times. Eventually, the model will get unlucky, draw a word with a 0.001% probability, and the sentence will derail into gibberish. Top-K sampling exists to prevent this bad luck. It forcefully zeroes out the long tail of bad options, ensuring the model is physically incapable of making a catastrophic choice.

Think of It Like This

Hiring from a stack of resumes

Imagine you post a job and receive 1,000 resumes. You rank them all from best (Rank 1) to worst (Rank 1,000).

If you use pure randomness to pick someone to interview (weighted by how good their resume is), you will usually pick one of the top candidates. But there is always a tiny, mathematical chance that you accidentally draw the resume at Rank 999.

Top-K sampling is a hiring policy that says: "Take the top 40 resumes (K=40). Throw the other 960 in the trash. Now, roll your weighted die to pick one person from the 40 remaining." You keep the benefit of randomness and variety, but you entirely eliminate the risk of hiring a disaster.

How It Actually Works

The Truncation Step

Top-K sampling is applied after temperature has reshaped the logits, but before the final token is selected. The algorithm is straightforward:

Sort all the tokens in the vocabulary descending by their probability (or logit score).
Keep the top $K$ tokens (where $K$ is a user-defined integer, like 40 or 50).
Set the probability of all other tokens (from rank $K+1$ down to the bottom) to exactly 0.
Re-normalize the remaining $K$ tokens so their probabilities sum back up to 1.0.
Sample randomly from this truncated, safe list.

The flaw: A fixed cutoff is rigid

While Top-K is incredibly effective at preventing gibberish, its reliance on a fixed integer $K$ is a structural weakness.

Imagine a situation where the model is highly confident. The probabilities for the next word are: 90% for "apple", 9% for "banana", and <0.1% for everything else. If $K=40$ , Top-K keeps the top 40 tokens. It forces the model to consider 38 terrible, low-probability words just to fill the quota of 40, keeping risk in the pool.

Conversely, imagine a situation where the model is completely unsure. There are 100 perfectly valid, grammatically correct synonyms, each with a 1% probability. If $K=40$ , Top-K violently chops off 60 perfectly good words, artificially suppressing the model's creativity.

Because $K$ cannot adapt to the shape of the distribution, Top-K has largely been superseded by its dynamic cousin, Top-P (nucleus) sampling (upcoming).

Show Me the Code

Here is how you would implement a strict Top-K cutoff on a set of logits before feeding them to the softmax function.

1import numpy as np2
3def softmax(z: np.ndarray) -> np.ndarray:4    e = np.exp(z - np.max(z))5    return e / np.sum(e)6
7def top_k_logits(logits: np.ndarray, k: int) -> np.ndarray:8    # If K is larger than the vocabulary, do nothing9    if k >= len(logits):10        return logits11        12    # Find the threshold logit value at rank K13    # np.partition is a fast way to find the Kth largest element14    kth_largest_value = np.partition(logits, -k)[-k]15    16    # Create a copy so we don't mutate the original17    truncated_logits = logits.copy()18    19    # Set everything strictly less than the Kth value to -infinity.20    # When -infinity goes through softmax, it becomes exactly 0 probability.21    truncated_logits[truncated_logits < kth_largest_value] = -np.inf22    23    return truncated_logits24
25# A vocabulary of 6 words26logits = np.array([5.0, 4.0, 3.0, 2.0, 1.0, 0.0])27
28print(f"Original Probs: {softmax(logits).round(3)}")29# -> [0.636 0.234 0.086 0.032 0.012 0.004] (All 6 have a chance)30
31# Apply Top-K (K=3)32safe_logits = top_k_logits(logits, k=3)33print(f"Top-3 Probs:    {softmax(safe_logits).round(3)}")34# -> [0.665 0.245 0.09  0.    0.    0.   ] (Bottom 3 are forced to 0%)

By setting the rejected logits to -infinity, the subsequent softmax function automatically crushes their probability to exactly zero, while scaling the remaining three tokens up slightly so they sum to 1.0.

Watch Out For

Using Top-K and Top-P together haphazardly

Most API providers allow you to set both Top-K and Top-P simultaneously. When both are active, they are applied sequentially (usually Top-K first, then Top-P). If you set $K=10$ and $P=0.9$ , the $K=10$ will aggressively chop the distribution first, often making the $P=0.9$ threshold entirely irrelevant. It is generally recommended to leave $K$ at a very high default (or disabled) and rely entirely on $P$ to dynamically handle the truncation.

The Quick Version

Standard sampling leaves a tiny mathematical risk that the model draws a completely nonsensical word from the tail end of the vocabulary.
Top-K sampling eliminates this risk by keeping only the $K$ highest-probability tokens and deleting the rest.
The remaining $K$ tokens are re-normalized to sum to 100%, and the model samples from this safe, truncated pool.
Because $K$ is a fixed number, it often keeps too many bad options when the model is confident, and cuts off too many good options when the model is unsure.