Probabilistic Sampling for Next-Token Selection

When a language model is asked for the next token, it does not hand you a word directly. It produces a raw score for every token in its vocabulary, and those scores must be turned into a choice. The simplest rule would be to always take the token with the highest score. Top-k sampling works differently, and understanding why is the goal of this article.

Keeping the randomness in the final step has a real effect on the text that comes out. Because the sampler can choose less obvious tokens, the generated text can become more varied and sometimes more creative. The same freedom has a cost: lower-probability tokens still have some chance of being selected, and that can push the output toward text that is grammatically broken or nonsensical. The problem is not that the sampler is random. The problem is that the random choice is allowed to reach tokens that should never have been candidates at that position. Top-k sampling exists to draw that boundary.

The process starts with logits. A logit is the raw score the model produces for each vocabulary token before any probabilities are computed. The masking step described later works directly on these logits, so it helps to picture them as a list with one score per vocabulary entry.

The vocabulary used in this lesson has nine tokens: forward, closer, effort, inches, toward, every, moves, pizza, and you. Each one sits at a fixed position in the list. The scores below are illustrative values made up for this lesson, chosen so that the example is easy to follow. They are not outputs of any particular model.

Notice that these numbers are not yet probabilities. They do not add up to 1, and some of them are negative. Their only job at this stage is to be compared with one another. Higher logits mean the model prefers that token more. The next question is how to decide which entries remain eligible.

Top-k sampling removes tokens from the selection process by masking their scores. The procedure has three stages that always run in the same order. First, find the k largest logits. Second, mask every other logit. Third, apply softmax to convert the result into probabilities over the tokens that survived.

Read the flow from left to right. The logits are the input. The top-k step decides which positions are eligible. The masking step removes the rest by assigning them negative infinity. The softmax step turns the surviving scores into probabilities. The purpose of this arrangement is that the next token is always sampled from one of the top-k positions.

The masking rule compares each logit with the smallest logit inside the top-k group. Call that smallest value the threshold. Any logit below the threshold is replaced with negative infinity. Logits at or above the threshold are kept exactly as they were.

Negative infinity is a deliberate choice because of how softmax works. Softmax turns each logit into a probability by exponentiating it and dividing by the total of all exponentiated values. The softmax formula is shown below, where z_i is the logit at position i after masking and p_i is the resulting probability.

When a logit z_j is set to negative infinity, the term e^{z_j} becomes zero. A masked position therefore adds nothing to the sum and ends up with probability 0. The remaining probabilities still add up to 1, because the normalization is carried out only over the unmasked positions. A reader who has met masking in the causal attention material of chapter 3 will recognize the same idea.

Notice what the example does not do. It does not simply output you. Token you is the most likely single choice, with a probability of about 0.678, but the other two candidates still have a real chance of being drawn. The sampling step, not the masking step, makes the final decision.

Greedy selection always takes the single highest-scoring token, so its candidate pool has exactly one member. Top-k sampling keeps a pool of k tokens and then chooses among them probabilistically. The two approaches agree on the highest-scoring token being the most likely outcome, but they differ in what they allow to happen next.

Top-k is designed to work alongside probabilistic sampling and temperature scaling. Each part has a separate job. Top-k sets the boundary of the candidate set. Temperature scaling adjusts how evenly the probabilities are spread within that boundary, and a higher temperature spreads them more evenly. Multinomial sampling then makes the final draw from the resulting distribution.