Decoding Strategies: Temperature & Top-k · Build an LLM From Scratch

The trained model has one more problem: with greedy decoding (ch4’s generate_text_simple), every request returns the identical highest- probability continuation — monotonous output, and a magnifier for the memorization we just diagnosed. This section re-implements generation with two knobs that trade determinism for controlled randomness: temperature scaling and top-k sampling.

The toy setup: a 9-token vocabulary

To see every number, the book shrinks the world to nine tokens and one next-token prediction (context: “every effort moves you”):

</> A small vocabulary and one set of next-token logits

vocab = {
  "closer": 0, "every": 1, "effort": 2, "forward": 3, "inches": 4,
  "moves": 5, "pizza": 6, "toward": 7, "you": 8,
}
inverse_vocab = {v: k for k, v in vocab.items()}

next_token_logits = torch.tensor(
  [4.51, 0.89, -1.90, 6.75, 1.63, -1.62, -1.89, 6.28, 1.79]
)

next_token_logits (1×9)

4.51	0.89	-1.90	6.75	1.63	-1.62	-1.89	6.28	1.79
closer	every	effort	forward	inches	moves	pizza	toward	you

Three tokens dominate: «forward» (6.75), «toward» (6.28), «closer» (4.51). The rest — including «pizza» — trail far behind.

Greedy vs. probabilistic sampling

</> Greedy decoding — deterministic

probas = torch.softmax(next_token_logits, dim=0)
next_token_id = torch.argmax(probas).item()
print(inverse_vocab[next_token_id])
# "forward"  — always, every single time #A

#A The maximum logit 6.75 stays maximal after softmax, so index 3 ("forward") wins deterministically — identical results for repeated requests.

</> Probabilistic sampling — multinomial replaces argmax

torch.manual_seed(123)
next_token_id = torch.multinomial(probas, num_samples=1).item()  #A
print(inverse_vocab[next_token_id])

#A Draw the token at random, in proportion to its probability: regardless of ranking, every token has a chance to be selected — «forward» often, «toward» sometimes, even «pizza» once in a long while.

Probabilistic sampling injects variety — but on its own it occasionally draws a truly unsuitable token. The two knobs below shape which randomness survives.

Knob 1 — temperature scaling

$p_i = \mathrm{softmax}\!\left(\frac{z_i}{T}\right)$

symbol	meaning
$z_i$	the raw next-token logits
$T$	temperature — $T>1$ flattens the distribution (more diverse), $T<1$ sharpens it (more confident), $T \to 0$ approaches greedy
$p_i$	the sampling distribution multinomial draws from

</> softmax_with_temperature

def softmax_with_temperature(logits, temperature):
  scaled_logits = logits / temperature
  return torch.softmax(scaled_logits, dim=0)

Drag the slider across the real toy logits. At T = 1: «forward» ≈ 0.56. Push T up and watch «pizza» come alive; pull it toward 0.1 and the distribution collapses onto greedy:

temperature T = 1.0

next-token logitssum = 16.44 (arbitrary)

closer

4.5100

every

0.8900

effort

-1.9000

forward

6.7500

inches

1.6300

moves

-1.6200

pizza

-1.8900

toward

6.2800

you

1.7900

softmax(x / 1.0)sum = 1, all positive

closer

0.0609

every

0.0016

effort

0.0001

forward

0.5721

inches

0.0034

moves

0.0001

pizza

0.0001

toward

0.3576

you

0.0040

Temperature reshapes the sampling distribution without touching the model: T>1 → more uniform (creative but riskier), T<1 → sharper (safer but repetitive). The same softmax-sharpness physics as ch3's √d_k scaling, now used as a control dial.

Knob 2 — top-k sampling

Temperature alone can still surface nonsense: a very low-probability token sometimes gets drawn, derailing the sentence. Top-k sampling first restricts the pool to the $k$ most likely tokens, masking everything else to $-\infty$ — the same $e^{-\infty} = 0$ trick as chapter 3’s causal mask, applied to the vocabulary axis:

</> Top-k masking (k = 3)

top_k = 3
top_logits, top_pos = torch.topk(next_token_logits, top_k)  #A
new_logits = torch.where(
  condition=next_token_logits < top_logits[-1],  #B
  input=torch.tensor(float("-inf")),
  other=next_token_logits
)
topk_probas = torch.softmax(new_logits, dim=0)  #C

# More efficient variant:
# new_logits = torch.full_like(next_token_logits, -torch.inf)
# new_logits[top_pos] = next_token_logits[top_pos]

#A topk returns the k highest logits and their positions: here [6.75, 6.28, 4.51] at [3, 7, 0].
#B Everything below the k-th highest value (top_logits[-1] = 4.51) becomes −∞.
#C Softmax turns the −∞ entries into exact zeros and renormalizes over the 3 survivors — sampling can now only pick «forward», «toward» or «closer».

Combine both knobs — restrict the pool with top-k, then shape it with temperature:

temperature T = 1.0top-k = 3

next-token logitssum = 16.44 (arbitrary)

closer

4.5100

every

0.8900

effort

-1.9000

forward

6.7500

inches

1.6300

moves

-1.6200

pizza

-1.8900

toward

6.2800

you

1.7900

top-3 → softmax(x / 1.0)sum = 1 over the 3 kept tokens

closer

0.0615

every

0.0000

effort

0.0000

forward

0.5775

inches

0.0000

moves

0.0000

pizza

0.0000

toward

0.3610

you

0.0000

Top-k zeroes the tail (dropped tokens get exactly 0 probability), temperature reshapes what remains. k = 9 disables the filter; k = 1 IS greedy decoding regardless of temperature.

Top-p (nucleus) sampling — the adaptive cousin

Top-p sampling sorts tokens by probability and keeps the smallest set whose cumulative probability reaches p — e.g. top_p=0.9 samples from the smallest group covering ≥ 90% of the mass. Unlike top-k’s fixed count, the pool size adapts to the distribution’s shape. When both are set, top-k filters first, then top-p filters within the survivors — so a token inside the top k can still be excluded by the cumulative-probability cut.

The upgraded generate() function

</> generate — temperature + top-k + early stopping in one loop

def generate(model, idx, max_new_tokens, context_size,
           temperature=0.0, top_k=None, eos_id=None):
  for _ in range(max_new_tokens):
      idx_cond = idx[:, -context_size:]  #A
      with torch.no_grad():
          logits = model(idx_cond)
      logits = logits[:, -1, :]

      # New: filter logits with top_k sampling
      if top_k is not None:
          top_logits, _ = torch.topk(logits, top_k)
          min_val = top_logits[:, -1]
          logits = torch.where(logits < min_val,
              torch.tensor(float("-inf")).to(logits.device), logits)  #B

      # New: apply temperature scaling
      if temperature > 0.0:  #C
          logits = logits / temperature
          probs = torch.softmax(logits, dim=-1)
          idx_next = torch.multinomial(probs, num_samples=1)
      else:  # greedy fallback
          idx_next = torch.argmax(logits, dim=-1, keepdim=True)

      if idx_next == eos_id:  #D
          break

      idx = torch.cat((idx, idx_next), dim=1)
  return idx

#A The loop skeleton is generate_text_simple from ch4, unchanged: crop context, forward pass, take the last position.
#B Top-k first: mask the tail to −∞ (device-aware) so softmax will zero it.
#C temperature > 0 → scale, softmax, multinomial sample. temperature = 0.0 → the deterministic argmax of ch4. One function now covers the whole greedy↔creative spectrum.
#D Optional early exit when the model emits a designated end-of-sequence token.

Key idea — decoding is a dial, not a switch

The model’s logits are fixed; how you turn them into a token is a free choice made at sampling time. Greedy (T=0) is one end; high temperature with no top-k is the other. In practice: moderate temperature (≈0.7–1.0) + top-k (or top-p) keeps text coherent and varied — and mitigates, to a degree, the training-data memorization diagnosed in §5.2. Next: saving, loading, and borrowing weights.

📝 Check yourself: decoding strategies

0 / 5

1.What is the practical problem with pure greedy decoding?
2.What does torch.multinomial(probas, num_samples=1) do differently from argmax?
3.Predict: with the toy logits, temperature T = 5 is applied before softmax. What happens to the distribution?
4.How does top-k sampling (k=3) change the sampling pool?
5.In the combined generate() function, what does temperature = 0.0 mean?