Temperature Scaling for Text Generation

When a text generator produces the next word, it does not simply output one answer. It scores every token in its vocabulary and then draws one token from the scores. Top-k sampling is one way of controlling that draw, and its name suggests something specific: the k largest scores are involved. Many learners read that name and conclude that top-k always returns the single best-scoring token. That reading is incorrect, and the rest of this article is about why.

The source makes a sharp distinction between narrowing and choosing. Top-k sampling does narrow the set of options, because it removes every token outside the k most likely ones. But narrowing is not the same as choosing. The k survivors are a pool of allowed candidates. The sampler still has to pick one of them, and it does so according to the probabilities that the candidates receive.

The contrast with greedy decoding makes this clear. Greedy decoding, which is the argmax choice, always returns the single highest-scoring token. Given the same scores, it would return the same token every time. Top-k sampling keeps a small group of plausible tokens open, so a lower-ranked survivor can sometimes be drawn. That is the source of the extra variety that top-k is designed to allow.

Top-k sampling works on the logits, which are the raw scores the model produces for each vocabulary token before any probabilities are computed. The procedure has three stages. First, find the k largest logits. Second, mask every other logit so that it cannot compete. Third, apply softmax to the result to convert the scores into probabilities. The output is a probability distribution over the surviving candidates, and the next token is drawn from it.

Read the pipeline 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 source states that the purpose of this arrangement is that the next token is always sampled from a top-k position.

The source uses a nine-token vocabulary: forward, closer, effort, inches, toward, every, moves, pizza, and you. With k set to 3, only three of those nine positions survive. The source does not list numeric logit values for this vocabulary, so the numbers below are illustrative. They are chosen only to show the mechanics.

Negative infinity is a deliberate choice. Softmax converts each logit into a probability by exponentiating it and dividing by the total. A masked position therefore contributes nothing once its value is negative infinity, because its exponential is zero. The remaining probabilities still sum to 1, since the normalization runs only over the unmasked positions.

After the other tokens are zeroed out, the k remaining probabilities must be rescaled so they add up to 1. Softmax handles this automatically. The exponentials of the three survivors are approximately 66.69 for every, 22.20 for inches, and 13.46 for you. Their total is about 102.35. Dividing each by that total gives the probabilities 0.65, 0.22, and 0.13. The rescaling depends only on the survivors, because the masked positions contributed zero to the total.

The source is clear that top-k sampling is not meant to stand alone. Top-k sets the boundary of the candidate set. Temperature scaling adjusts how evenly the probabilities are spread within that boundary. Multinomial sampling then makes the final draw. A higher temperature spreads the next-token probabilities more evenly, so the model is less likely to keep choosing the single most probable token, which opens the door to less obvious but potentially more creative continuations.

The same flexibility has a cost. Because low-probability tokens still receive some chance of being selected, the output can drift into text that is grammatically broken or plainly nonsensical. The source illustrates this with the phrase every effort moves you pizza. 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 masking is the boundary that keeps those tokens out.

The source says the next step is to modify the text generation function so that temperature scaling and the multinomial function operate on the non-zero probabilities. That modification is beyond the scope of this article. The point to carry forward is the ordering: the top-k masking happens before the probabilities are finalized.