Top-k Sampling › How Top-k Sampling Works

Now that we understand the problem, let us build up the solution piece by piece. The core insight behind top-k sampling is deceptively simple: if most of the probability mass concentrates in a handful of tokens, why not just sample from those and ignore the rest?

The elegance of this solution is that it makes almost no assumptions about the shape of the distribution. It does not try to model why probability falls off or where the "natural" cutoff should be. It simply imposes a hard count-based boundary and accepts the consequences. This makes top-k both easy to understand and easy to implement, which explains why it became so widely adopted despite its simplicity.

Think of top-k sampling as a casting director who says: "I will only consider actors in the top $k$ of my ranking. Everyone else is not being called back, regardless of how close they were to the boundary." The director does not re-evaluate the top actors based on who was excluded; they simply rescale their assessments of the remaining candidates. Someone ranked 3rd out of 50,000 is still ranked 3rd out of the top $k$, and their relative standing compared to other top-$k$ candidates is unchanged.

This analogy also highlights an important subtlety: the excluded tokens receive zero probability, not some small residual probability. This is a hard cut, not a soft fade. The benefit is that we guarantee no sampling from the tail. The cost is that tokens just outside the boundary (ranked $k+1$ and $k+2$) receive zero probability even if they were only slightly less probable than the token at rank $k$. Whether this boundary effect matters in practice depends on how steeply the probability drops off near rank $k$.

Think about how you complete sentences. Given "The capital of France is", you do not mentally consider every possible word in English. You immediately focus on a small set of plausible continuations: "Paris", maybe "a", "the", or "known". Your mental probability distribution is not uniform across all words, and neither is the model's. Top-k sampling formalizes this intuition by explicitly restricting attention to the most likely candidates.

The approach works in five steps:

1. Compute probabilities: The model outputs a distribution over all vocabulary tokens
2. Rank by likelihood: Sort tokens from most to least probable
3. Keep the top k: Select only the $k$ highest-probability tokens
4. Zero out the rest: Set all other token probabilities to exactly zero
5. Renormalize and sample: Rescale the remaining probabilities so they sum to 1, then sample

Let us formalize this process. At generation step $i$, the model has seen context $x_{<i}$ (all previous tokens) and outputs a probability $P(x_i | x_{<i})$ for each possible next token $x_i$. This distribution spans the entire vocabulary, but we only want to sample from the top performers.

First, we need to identify which tokens make the cut. We define $V_k$ as the set containing the $k$ tokens with highest probability. Formally, $V_k = \{v_1, v_2, \ldots, v_k\}$ where $P(v_1 | x_{<i}) \geq P(v_2 | x_{<i}) \geq \ldots \geq P(v_k | x_{<i})$ and every token outside $V_k$ has probability no greater than the smallest element of $V_k$. If the vocabulary has 50,000 tokens and $k=50$, then $V_k$ contains exactly the 50 most likely tokens for this specific context.

Now we need to handle a subtle but important issue: after removing tokens from consideration, the probabilities of the remaining tokens no longer sum to 1. A probability distribution must sum to 1 for sampling to work correctly. If the top-$k$ tokens collectively hold 85% of the probability mass, the removed tokens held the other 15%, and our truncated distribution sums to only 0.85. The solution is renormalization: we divide each remaining probability by the total mass of the kept tokens, scaling everything up so the distribution is valid again.

To compute the renormalization constant, we sum the probabilities of all top-$k$ tokens:

$$Z_k = \sum_{x \in V_k} P(x | x_{<i})$$

where:

• $V_k$: the set of $k$ tokens with highest probability
• $P(x | x_{<i})$: the model's original probability for token $x$ given context $x_{<i}$
• $Z_k$: the total probability mass held by the top-$k$ tokens, which acts as a normalizing constant

This gives us the formal definition of top-k sampling. Given the original distribution and the normalization constant, the truncated distribution assigns probability:

$$P_k(x_i | x_{<i}) = \begin{cases} \frac{P(x_i | x_{<i})}{Z_k} & \text{if } x_i \in V_k \\ 0 & \text{otherwise} \end{cases}$$

where:

• $x_i$: a candidate token at position $i$
• $x_{<i}$: the context (all tokens before position $i$)
• $P(x_i | x_{<i})$: the original probability the model assigns to token $x_i$ given the context
• $V_k$: the set of $k$ tokens with highest probability under $P(\cdot | x_{<i})$
• $Z_k = \sum_{x \in V_k} P(x | x_{<i})$: the normalization constant, which equals the sum of probabilities of the top-$k$ tokens
• $P_k(x_i | x_{<i})$: the renormalized probability used for sampling

Why does this formula make sense? Notice that for tokens inside $V_k$, dividing by $Z_k$ (which is less than 1) makes each probability larger than the original. This is correct: we are redistributing the probability that was held by the excluded tokens and giving it proportionally to the included tokens. For tokens outside $V_k$, the formula assigns zero. The resulting distribution sums to exactly 1 because $\sum_{x \in V_k} P_k(x | x_{<i}) = \sum_{x \in V_k} P(x | x_{<i}) / Z_k = Z_k / Z_k = 1$.

A necessary property of this construction is that renormalization preserves the relative likelihood of tokens within the top-$k$. This matters because the model's ranking captures important information about which tokens are more semantically appropriate than others, and we want to honor that information.

Consider two tokens: "Paris" with original probability 0.30 and "the" with probability 0.15. In the original distribution, "Paris" is exactly twice as likely as "the".

After truncation, suppose the top-$k$ tokens have cumulative probability $Z_k = 0.90$. The renormalized probabilities become:

• "Paris": $0.30 / 0.90 = 0.333$
• "the": $0.15 / 0.90 = 0.167$

Notice that $0.333 / 0.167 = 2$. The ratio is preserved! This happens because we divide both probabilities by the same constant $Z_k$. Renormalization scales all probabilities equally, maintaining their relative ordering and ratios.

This property matters because it means top-k sampling respects the model's preferences among plausible tokens. We are not arbitrarily reweighting tokens; we are simply restricting which tokens can be sampled while honoring the model's ranking within that restricted set. If you trust the model's learned distribution over the top candidates, renormalized top-k sampling samples from that distribution faithfully.

The key insight is that top-k sampling changes the support of the distribution (which tokens have non-zero probability) without changing the shape of the distribution over the supported tokens. The relative probabilities among kept tokens are identical to what they would be under the original distribution, up to the constant scaling factor $1/Z_k$.

The following visualization shows this process concretely. On the left, we see the original distribution over the top 20 tokens. On the right, we see the truncated distribution after keeping only $k=10$ tokens and renormalizing.