Softmax and Probabilities for Token Selection

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 it is given the same input. A generator that always picks the top candidate tends to produce repetitive and narrow text.

The modified generate function addresses this by adding two controls to the same loop. Temperature scaling changes how random the choice is, and top-k sampling restricts the set of tokens the choice can come from. Both controls live inside one function, so the whole mechanism can be read in one place.

Each pass through the loop starts the same way as the earlier generator. The function 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. The lines that follow the logits selection are where the added diversity comes from.

On each pass, the data moves through one fixed sequence of operations. Top-k filtering runs first, then temperature scaling, then softmax converts the logits into probabilities, and finally torch.multinomial samples one token from those probabilities. Whether the temperature step runs at all depends on the temperature value, which is why the flowchart below shows a branch.

Read the diagram from the top. The top-k step always comes before the branch. The temperature decision then sends the flow either to the sampling lines, which divide, normalize, and draw one token at random, or to torch.argmax, which takes the highest-scoring token. Both paths lead to the same end-of-sequence check, which decides whether the loop continues.

Top-k sampling restricts the 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. Every other logit is kept as it is.

In this generated example, the two highest logits, 3.2 and 2.7, stay in place. The other two are set to negative infinity. The vocabulary arrays are not shortened. The remaining positions still exist, but they carry the negative infinity value instead of their original scores.

Temperature scaling is applied to the logits before softmax. In the formula, z stands for the logits of the candidate tokens, T stands for the temperature, and z' is the scaled result that the function passes to softmax. In the code, this is the line logits = logits / temperature.

The purpose of this division is to change how sharply the probabilities differ after softmax. The source does not show the probability values directly, so the exact effect depends on the logits at each step. The direction of the mechanism is what matters: the temperature is the divisor, and the scaled logits are what softmax turns into probabilities for sampling. Reasoning from the formula, dividing by a larger T shrinks the gaps between the logits, which pushes the probabilities closer together. Dividing by a smaller T enlarges the gaps.

In this generated example, the temperature is 2.0. The gap between the two logits drops from 2.0 to 1.0 after division. Because the gap is smaller, softmax assigns the two tokens more similar probabilities, so the sampling step has more variety to draw from.

The temperature value decides which selection path runs. With temperature greater than 0.0, the function samples a token from the probabilities. With temperature 0.0 or lower, it selects the token with the highest logit. The earlier generator always made the greedy choice, so this branch keeps that behavior as a special case.

The seed matters because the source's example calls torch.manual_seed(123) before generating. That call fixes the random number generator, so the sampling step behaves reproducibly on that run. Without the seed, the same sampling settings are not guaranteed to produce the same tokens each time.

The source runs the modified function with a fixed seed, top_k set to 25, and temperature set to 1.4. Both controls are active, so the loop takes the sampling path on every pass. The output is the prompt followed by up to 15 newly generated tokens. The source notes that the text reads as plausible words that drift in meaning, which is the variety the function is designed to allow.

To trace a single pass, follow the same order each time. Filter to the top 25 logits, divide the surviving logits by 1.4, apply softmax, and let torch.multinomial draw one token. Append that token, then check whether it equals eos_id. If it does, the loop breaks early and returns the sequence so far. If not, the loop continues until max_new_tokens iterations have run.