Sampling Random Tokens with torch.multinomial

Earlier in this module, the text generation loop produced each new token by looking only at the last time step of the model's output and choosing the single most likely token. That approach is simple and predictable, but it always makes the same choice when given the same input. A generator that always picks the top candidate tends to produce repetitive and narrow text. This article covers the modified version of that loop, which gives the generator two controls over its choice: temperature scaling, which changes how random the choice is, and top-k sampling, which restricts the set of tokens the choice can come from.

The source modifies the earlier function, called generate_text_simple, into a new function named generate. Both controls are added inside the same loop, so the whole mechanism can be read in one function. The plan for this article is to run the modified function first, then trace the mechanism step by step.

The loop body starts the same way as before. It crops the context to context_size, runs the model, and keeps only the logits for the last time step. Everything after that line is new. On each pass, the data moves through a single pipeline: top-k filtering runs first, then temperature scaling runs, then softmax converts the logits into probabilities, and finally torch.multinomial samples one token from those probabilities.

The temperature step is conditional, so it does not always run. The branch that decides between the two selection paths depends on the temperature value. When temperature is greater than 0.0, the function divides the logits by the temperature and samples a token from the resulting probabilities. When temperature is 0.0 or lower, the code skips the scaling and sampling lines and calls torch.argmax instead, which selects the token with the highest logit. The earlier generator always made that greedy choice, so this branch preserves the earlier behavior as a special case.

After a token is selected, the function checks whether it equals eos_id, the end-of-sequence token ID. The source states that generation can stop early when that token appears. The break statement exits the loop and the function returns the sequence built so far. If no end-of-sequence token appears, the loop runs for max_new_tokens iterations, and each pass appends one token with torch.cat. In the source's example, the limit is 15 newly generated tokens.

Top-k sampling restricts candidates to the k largest logits. The function finds the cutoff, which is the smallest of the top k logits. Any logit below that cutoff is replaced with negative infinity. Because the probabilities are produced by softmax later in the function, a token whose logit is negative infinity receives no probability and cannot be selected by sampling. The function narrows the candidate set to the top k tokens without removing the remaining tokens from the vocabulary arrays.

Temperature scaling divides the logits by a temperature value before they reach softmax. In the formula, $z$ stands for the logits for the candidate tokens, $T$ stands for the temperature, and $z'$ is the scaled logits that the function passes to softmax.

In the code, this corresponds to the line logits = logits / temperature. The source's example uses temperature=1.4, a value greater than 0.0, so this division runs and the sampling branch is used. The temperature value is the divisor, and the scaled logits are what softmax turns into probabilities for sampling.

The purpose of this division is to change how sharply the probabilities differ after softmax. The source does not show the probability values, so the exact effect depends on the logits at each step. The direction of the effect follows from the arithmetic: dividing by a value below 1 makes the scaled logits further apart, while dividing by a value above 1 pulls them closer together. Closer logits produce a flatter distribution, which gives lower-ranked tokens a better chance of being sampled.

The two selection paths are both controlled by the temperature value. With temperature greater than 0.0, the function samples a token from the probabilities using torch.multinomial. Sampling picks one token at random according to the probabilities, so a lower-ranked token can be chosen. With temperature 0.0 or lower, the function calls torch.argmax and always takes the token with the highest logit. That is the greedy path, and it matches the behavior of the earlier generator.

Whether a run is deterministic depends on the branch it takes. On the greedy path, torch.argmax selects the highest logit, so the same input gives the same token. On the sampling path, torch.multinomial draws from the probabilities, and the source notes that torch.manual_seed fixes the random number generator so the sampling step behaves reproducibly on that run. The divergence point is therefore the sampling call: two runs with different seeds, or with the seed unset, can choose different tokens from the same probabilities.