Autoregressive Sampling: Greedy vs Temperature Scaling

The Big Idea

Now that our GPT model is trained from Lesson 4.1, how do you make it generate an original story? You provide a seed prompt (e.g. 'ROMEO:'), feed it through the network, and inspect the predicted scores for the next character.

The simplest strategy is 'greedy decoding': always pick the single token with the highest logit (argmax). While this sounds logical, greedy decoding is notoriously flawed: it gets trapped in sterile, repetitive infinite loops ('The King said the King said the King said...').

To introduce creativity, we divide raw logits by a positive number called Temperature (T) before applying softmax, as detailed in temperature-scaled Softmax. A low temperature (T = 0.2) sharpens differences, making the top choice 98% likely (conservative and factual). A high temperature (T = 1.5) flattens probabilities, giving surprising tokens a fighting chance (creative and poetic). We sample a token using torch.multinomial, append it to our sequence, crop to our context window, and repeat!

Architecture Blueprint & Tensor Mechanics

The autoregressive generation loop crops incoming context to block_size, divides final-step logits by temperature, and samples via torch.multinomial.

import torch
import torch.nn.functional as F

@torch.no_grad()
def generate(model, idx, max_new_tokens, block_size, temperature=0.8):
    model.eval()
    for _ in range(max_new_tokens):
        # 1. Crop current context to the maximum block_size
        idx_cond = idx if idx.size(1) <= block_size else idx[:, -block_size:]
        
        # 2. Get forward predictions
        logits, _ = model(idx_cond)
        
        # 3. Focus solely on the final time step and apply temperature
        logits = logits[:, -1, :] / max(temperature, 1e-5)
        
        # 4. Convert logits to probabilities
        probs = F.softmax(logits, dim=-1)
        
        # 5. Sample from the probability distribution
        idx_next = torch.multinomial(probs, num_samples=1)
        
        # 6. Append sampled token to the running sequence
        idx = torch.cat((idx, idx_next), dim=1)
        
    return idx

# Demonstration with dummy model
# prompt = torch.tensor([[10, 24, 32]], dtype=torch.long) # Encoded prompt
# out_tokens = generate(model, prompt, max_new_tokens=100, block_size=256, temperature=0.8)
# print("Generated text:", decode(out_tokens[0].tolist()))