Skip to content
Module 3 ยท Training and Generation
Chapter 16 of 17
Run it in your browser

Chapter 16: Temperature and Top-k

You are here in the big picture Figure 16.1: Where we are: we are tuning the text generation process.

Plain sampling gives us varied output, but it can be unpredictable. Sometimes the model chooses a rare character that breaks the grammar or creates a nonsense word. We need controls to balance creativity with coherence. In this chapter you will:

  • Use temperature to adjust how risky the model's choices are.
  • Use top-k to cut the list of choices short, so the model cannot pick truly bad characters.
  • Combine these controls for optimal text generation.

Words to Know

  • Temperature: a number the scores (logits) are divided by before they become probabilities. A low temperature plays it safe, and a high one takes risks.
  • Top-k: a limit that lets the model sample only from the k most likely next tokens, where k is a number you choose (40 here).

Theory

Two controls sit between the raw scores the model produces and the character it finally picks. Temperature reshapes the odds, while top-k decides which characters are allowed to compete at all. They are independent, and in practice you set both.

The two generation controls side by side Figure 16.2: Temperature changes the shape of the odds; top-k changes how many characters stay in the running.

Temperature

Temperature is a simple math trick applied to the raw scores (logits) before the softmax step. We divide every score by one number, the temperature, and that changes how the chances are spread across the characters.

Terminal
$ python src/examples/ch16_temperature.py
Saved chart: ch16_temperature.png

Bars of next-character probabilities at three temperatures Figure 16.3: A plotted chart showing probabilities: low temperature sharpens the choice, high temperature flattens it.

  • Low temperature (e.g., 0.5): Dividing by a number smaller than 1 pushes the scores further apart, so the top choice becomes overwhelmingly favored and the model plays it safe.
  • High temperature (e.g., 2.0): Dividing by a large number squishes the scores together. The probabilities flatten out, meaning unusual choices become more likely, and the model takes risks.
  • T = 1.0: This is standard sampling, where the model uses its learned probabilities exactly as they are.

Top-k

Even at a safe temperature, there is a tiny mathematical chance (say, 0.001%) that the model might pick an absurd character. If you generate many thousands of characters, those tiny chances add up, and sooner or later a mistake appears.

Top-k sampling solves this by putting a hard limit on the choices. If top_k = 40, we look at the 65 possible characters, keep the 40 with the highest scores, and discard the bottom 25 by setting their probabilities to zero. This cuts off the "long tail" of bad choices, guaranteeing that the model only samples from the most reasonable options.

Ranked scores with the low-scoring tail greyed out Figure 16.4: Top-k keeps the highest-scoring characters and zeroes the rest. Drawn here with ten characters and k = 4; the book's code uses 65 and k = 40.

These controls adjust the Next-Token Scores on the "where we are" map, right before generation loops back to produce the New Text.

In Business

For the house-style assistant, you might use a low temperature (0.5) when generating compliance documentation to ensure it stays close to the safest boilerplate text. When brainstorming marketing slogans, a higher temperature (0.9) with top-k sampling (40) will produce creative, surprising slogans that still make grammatical sense.

Code

We apply temperature and top-k right before the softmax function in the generation loop.

src/ch15_generate_sampling.py (excerpt)
        # Temperature controls randomness: lower is less random
        logits = logits / temperature

        # Top-k sampling limits the choices to the k most likely tokens
        if top_k is not None:
            threshold = logits.topk(top_k).values[:, -1, None]
            logits = logits.masked_fill(logits < threshold, float("-inf"))

        probs   = F.softmax(logits, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)

Line 2 divides the scores by the temperature to adjust randomness. Line 7 discards any score below the top-k threshold by setting it to negative infinity.

Code flow: logits divided by temperature, filtered by top-k, then softmaxed Figure 16.5: The code flow applies temperature first, then top-k, and finally softmax.

In PyTorch, we use masked_fill to apply top-k. We find the score of the 40th best token (threshold), and any token with a score lower than that gets its value changed to negative infinity (-inf). When softmax calculates probabilities, anything with a score of -inf becomes exactly 0.

The run below shows how different settings affect the generated text.

Terminal
$ python src/ch15_generate_sampling.py
--- temp=0.5, top_k=40 ---
JULIET:
My must grace prove a son the come on my lord.

...
--- temp=0.8, top_k=40 ---
JULIET:
That will, my call thee for thee king win a be parder
The his in suppy ascented a not a say?
...
--- temp=1.0, top_k=40 ---
JULIET:
Look, do indo witned me devise thy grohed:
A give gring brothe whith ighers paison:
...
--- temp=1.5, top_k=40 ---
JULIET:
fattlyals give niet, thou stain, viful-wid thel,
Brob dayse mernion shing you fea,
...
--- temp=1.0, no top_k ---
JULIET:
Than mide smear of thou I am go vious, you me
now you scalf them frate ell, to 'till will she
...
--- Sweet spot ---
temperature=0.8 to 1.0 and top_k=40 usually gives the best results.

What just happened:

  • At temp=0.5, the text is coherent but relies heavily on common, safe words.
  • At temp=0.8 to 1.0, the text feels more natural and varied.
  • At temp=1.5, the text quickly devolves into chaos and made-up words ("fattlyals").
  • The sweet spot is typically a temperature between 0.8 and 1.0, combined with top_k=40.

Shape Check

Table 16.1 lists the shapes of tensors used during top-k filtering.

Table 16.1: Tensors modified during the top-k sampling process.

Tensor Shape What it means
threshold [1, 1] The cutoff score (the 40th highest value).
logits [1, 65] 65 scores, where 25 of them have been set to -inf.

Try It

Try It

Open the existing script src/examples/ch16_explore_temp.py and run it. It uses a very low temperature (0.1), and you will notice that it behaves almost exactly like greedy decoding, picking the safe top character every time.

Terminal
$ python src/examples/ch16_explore_temp.py
--- Prompt: 'JULIET:\n' ---
Generating with T=0.1...
JULIET:
I will the shall the shall be the son th...

Key Takeaways

  • Temperature adjusts the spread of the probabilities. Low temperature makes the model conservative, while high temperature makes it creative and risky.
  • Top-k restricts the model from ever picking the worst-scoring tokens, preventing bizarre errors.
  • To combine them, first apply temperature, then apply top-k to filter out the bad options, and finally convert to probabilities with softmax.
  • A combination of temperature=0.8 and top_k=40 is a widely used default for high-quality text generation.

Check Your Understanding

  1. If you set the temperature to 0.1, what happens to the gap between the highest and lowest scores?
  2. Why do we replace the eliminated scores in top-k with negative infinity (-inf) instead of 0?
  3. Which step must happen first in the code: applying top-k or calculating softmax probabilities?

Assignment

Assignment 16: Turn up the heat

About 15 minutes. You change one number and write no new code. Work in this chapter's Colab notebook or on your own laptop.

  1. Run it. Run src/examples/ch16_explore_temp.py as it is. Copy the text the model writes after JULIET:.
  2. Change one number. Near the bottom of the file, find temperature=0.1. Temperature is the number the scores are divided by before the model draws a character. Change 0.1 to 2.0 and run the script again. The line Generating with T=0.1... is a fixed label, so it will still say 0.1.
  3. Compare. Which text keeps repeating the same words? Which one is full of made-up words? Which one is closer to real English?
  4. Explain it. In two or three sentences a manager could follow, say why neither setting would suit an assistant that writes to customers, and which way you would turn the dial from each one.

Hand in: both texts, your answers to step 3, and your sentences.

Further Reading

Where top-k sampling comes from. This paper is about writing stories from a prompt, and along the way it introduced the sampling rule you use in Chapter 16: keep only the k most likely next tokens and draw from those. It keeps the text varied without letting the model pick something absurd from the long tail.

Why always picking the likeliest word goes wrong. Always taking the highest-scoring next token, which Chapter 15 calls greedy decoding, produces flat and repetitive text, and this paper shows why: human writing is not made of the most predictable word at every turn. Their alternative keeps the smallest set of tokens whose probabilities add up to a chosen share and samples from that. How you choose the next token matters as much as how well the model was trained.

Fan, A., Lewis, M., & Dauphin, Y. (2018). Hierarchical neural story generation (arXiv:1805.04833). arXiv. https://doi.org/10.48550/arXiv.1805.04833

Holtzman, A., Buys, J., Du, L., Forbes, M., & Choi, Y. (2019). The curious case of neural text degeneration (arXiv:1904.09751). arXiv. https://doi.org/10.48550/arXiv.1904.09751

Next up: Chapter 17: Putting It All Together