Top-k Sampling › Worked Example: Step-by-Step Numerical Trace

To make the algorithm fully concrete, let us trace through a complete top-k sampling step with small, explicit numbers. This example uses a toy vocabulary of 8 tokens so we can inspect every calculation by hand.

Suppose we have just processed the context "The sky is" and our language model outputs the following raw logits over 8 vocabulary tokens:

Step 1: Convert logits to probabilities using softmax. For each logit $z_i$, compute $e^{z_i}$, then divide by the sum of all exponentials.

The exponentials are:

• blue: $e^{3.2} = 24.53$
• clear: $e^{2.8} = 16.44$
• dark: $e^{1.5} = 4.48$
• gray: $e^{1.1} = 3.00$
• red: $e^{0.4} = 1.49$
• green: $e^{-0.3} = 0.74$
• loud: $e^{-1.2} = 0.30$
• empty: $e^{-2.5} = 0.08$

The sum of all exponentials is $24.53 + 16.44 + 4.48 + 3.00 + 1.49 + 0.74 + 0.30 + 0.08 = 51.06$.

Dividing each exponential by 51.06 gives the probability distribution:

Step 2: Select the top $k=4$ tokens. Ranking by probability, the top 4 are: blue (0.480), clear (0.322), dark (0.088), gray (0.059). The remaining tokens (red, green, loud, empty) are excluded.

Step 3: Compute the normalization constant. We sum the probabilities of the kept tokens:

$$Z_4 = 0.480 + 0.322 + 0.088 + 0.059 = 0.949$$

Step 4: Renormalize. Divide each kept token's probability by $Z_4 = 0.949$:

Step 5: Sample. We now draw one token from this renormalized distribution. With probability 0.506 we select "blue"; with probability 0.339, "clear"; with probability 0.093, "dark"; and with probability 0.062, "gray". The tokens "red", "green", "loud", and "empty" can never be selected.

Verification of ratio preservation. Let us check that the ratio of "blue" to "clear" is preserved. Originally: $0.480 / 0.322 = 1.491$. After renormalization: $0.506 / 0.339 = 1.493$. The tiny difference is rounding error; the ratio is preserved because both are divided by the same $Z_4$. The model's belief that "blue" is about 1.5 times more likely than "clear" is fully honored.

This numerical trace shows every operation explicitly: softmax converts logits to probabilities, top-$k$ selection identifies the candidates, summation computes the normalization constant, division produces a valid probability distribution, and multinomial sampling draws the final token.