Module 08: Generation

Introduction

After training, the next step is generating text. The model predicts probabilities for the next token - but how do we choose which one? This module explores decoding strategies that shape output quality, diversity, and coherence.

Text generation produces text from a trained model. The model predicts probability distributions over the vocabulary, and we must decide how to select the next token from these probabilities.

Why does the decoding strategy matter?

  • Different strategies, different outputs: Greedy decoding gives deterministic results, sampling gives variety
  • Control creativity vs coherence: Temperature and filtering parameters let us tune this tradeoff
  • Application-specific needs: Code generation wants precision, creative writing wants diversity
  • Ensure quality outputs: Proper temperature and filtering prevent repetition, incoherence, and gibberish

Understanding generation is essential.

What You’ll Learn

After this module, you can:

  • Implement the autoregressive generation loop from scratch
  • Apply and combine decoding strategies (greedy, temperature, top-k, top-p, min-p)
  • Use repetition penalties to prevent degenerate outputs
  • Constrain decoding to a grammar so output is valid JSON/dates/enums by construction
  • Understand KV-caching for efficient generation
  • Choose appropriate generation parameters for different use cases

Prerequisites

This module requires familiarity with:

The Generation Loop

Text generation is autoregressive - we generate one token at a time, feeding previous tokens back to the model:

Each step:

  1. Feed current tokens to the model
  2. Get probability distribution over vocabulary for next position
  3. Apply decoding strategy to select next token
  4. Append selected token to sequence
  5. Repeat until stopping criterion met

From Scratch: The Generation Loop

Generation is just repeated next-token prediction with sampling. Let’s build it step by step.

import numpy as np

def softmax(x: np.ndarray) -> np.ndarray:
    """Stable softmax: subtract max to prevent overflow."""
    x_max = x.max(axis=-1, keepdims=True)
    exp_x = np.exp(x - x_max)
    return exp_x / exp_x.sum(axis=-1, keepdims=True)

def generate_scratch(get_logits, context: np.ndarray, max_new_tokens: int = 10, temperature: float = 1.0) -> np.ndarray:
    """
    Generate tokens autoregressively from scratch.

    Args:
        get_logits: Function that takes context (1, seq_len) and returns logits (1, vocab_size)
        context: Starting token ids, shape (1, seq_len)
        max_new_tokens: How many tokens to generate
        temperature: Sampling temperature (higher = more random)

    Returns:
        Extended context with generated tokens
    """
    ctx = context.copy()

    for _ in range(max_new_tokens):
        # 1. Get logits for last position
        logits = get_logits(ctx)  # (1, vocab_size)

        # 2. Apply temperature (scale before softmax)
        logits = logits / temperature

        # 3. Convert to probabilities
        probs = softmax(logits)[0]  # (vocab_size,)

        # 4. Sample next token
        next_token = np.random.choice(len(probs), p=probs)

        # 5. Append to context
        ctx = np.concatenate([ctx, [[next_token]]], axis=1)

    return ctx

Key insight: Generation is surprisingly simple. The model predicts, we sample, and we feed the result back in. That’s it.

From Logits to Tokens

The logits-to-token pipeline is the heart of generation:

# Step-by-step: logits -> probabilities -> token

# Simulate model output (logits for 8-token vocabulary)
logits = np.array([[2.0, 1.5, 0.5, 0.0, -0.5, -1.0, -1.5, -2.0]])
token_names = ["the", "cat", "sat", "on", "mat", "dog", "ran", "fast"]

print("Step 1: Raw logits from model")
for i, (name, logit) in enumerate(zip(token_names, logits[0])):
    print(f"  {name:>4}: {logit:+.1f}")

print("\nStep 2: Apply softmax to get probabilities")
probs = softmax(logits)[0]
for name, prob in zip(token_names, probs):
    bar = "█" * int(prob * 40)
    print(f"  {name:>4}: {prob:.3f} {bar}")

print("\nStep 3: Sample from the distribution")
np.random.seed(42)
sampled_idx = np.random.choice(len(probs), p=probs)
print(f"  Sampled token: '{token_names[sampled_idx]}' (index {sampled_idx})")
Step 1: Raw logits from model
   the: +2.0
   cat: +1.5
   sat: +0.5
    on: +0.0
   mat: -0.5
   dog: -1.0
   ran: -1.5
  fast: -2.0

Step 2: Apply softmax to get probabilities
   the: 0.466 ██████████████████
   cat: 0.283 ███████████
   sat: 0.104 ████
    on: 0.063 ██
   mat: 0.038 █
   dog: 0.023 
   ran: 0.014 
  fast: 0.009 

Step 3: Sample from the distribution
  Sampled token: 'the' (index 0)

Temperature controls the randomness by scaling logits before softmax:

print("Effect of temperature on the same logits:\n")

for temp in [0.5, 1.0, 2.0]:
    scaled_logits = logits / temp
    probs = softmax(scaled_logits)[0]

    print(f"Temperature = {temp}:")
    for name, prob in zip(token_names[:4], probs[:4]):  # Show top 4
        bar = "█" * int(prob * 30)
        print(f"  {name:>4}: {prob:.3f} {bar}")
    print()

print("Lower temp = sharper (more deterministic)")
print("Higher temp = flatter (more random)")
Effect of temperature on the same logits:

Temperature = 0.5:
   the: 0.691 ████████████████████
   cat: 0.254 ███████
   sat: 0.034 █
    on: 0.013 

Temperature = 1.0:
   the: 0.466 █████████████
   cat: 0.283 ████████
   sat: 0.104 ███
    on: 0.063 █

Temperature = 2.0:
   the: 0.291 ████████
   cat: 0.227 ██████
   sat: 0.137 ████
    on: 0.107 ███

Lower temp = sharper (more deterministic)
Higher temp = flatter (more random)

Now let’s see it in action with a real model.

Code Walkthrough

Let’s explore generation interactively:

import torch
import torch.nn.functional as F
import numpy as np

print(f"PyTorch version: {torch.__version__}")
PyTorch version: 2.12.1+cu130

Setting Up

import sys
sys.path.insert(0, '..')

from generation import (
    top_k_filtering,
    top_p_filtering,
    apply_repetition_penalty,
    generate,
    generate_greedy,
    generate_sample,
    get_token_probabilities,
    get_top_tokens,
)
from m06_transformer.transformer import create_gpt_tiny

# Create a small model for demonstration
vocab_size = 50
model = create_gpt_tiny(vocab_size=vocab_size)

# Create a sample prompt
prompt = torch.randint(0, vocab_size, (1, 5))
print(f"Prompt tokens: {prompt[0].tolist()}")
Prompt tokens: [9, 33, 35, 29, 24]

Understanding Model Output

A language model outputs logits (unnormalized scores) that become probabilities after softmax:

# Get probability distribution for next token
probs = get_token_probabilities(model, prompt)

print(f"Probability distribution shape: {probs.shape}")
print(f"Sum of probabilities: {probs.sum().item():.4f}")

# Show top tokens
top = get_top_tokens(probs, k=5)
print("\nTop 5 most likely next tokens:")
for token_id, prob in top:
    print(f"  Token {token_id}: {prob*100:.2f}%")
Probability distribution shape: torch.Size([1, 50])
Sum of probabilities: 1.0000

Top 5 most likely next tokens:
  Token 5: 3.62%
  Token 19: 3.58%
  Token 6: 3.00%
  Token 27: 2.61%
  Token 17: 2.59%

Probability Distribution Explorer

Explore how different sampling parameters affect the probability distribution:

Decoding Strategies

1. Greedy Decoding

Always pick the token with the highest probability - simple but often repetitive.

Pros: Deterministic, coherent output Cons: Boring, repetitive, can get stuck in loops

# Generate with greedy decoding
output_greedy = generate_greedy(model, prompt, max_new_tokens=15)

print(f"Prompt: {prompt[0].tolist()}")
print(f"Generated: {output_greedy[0, 5:].tolist()}")
Prompt: [9, 33, 35, 29, 24]
Generated: [5, 19, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 21, 5]
# Greedy is deterministic - same output every time
print("Multiple greedy generations (should all be identical):")
for i in range(3):
    out = generate_greedy(model, prompt, max_new_tokens=10)
    print(f"  Run {i+1}: {out[0, 5:].tolist()}")
Multiple greedy generations (should all be identical):
  Run 1: [5, 19, 5, 5, 5, 5, 5, 5, 5, 5]
  Run 2: [5, 19, 5, 5, 5, 5, 5, 5, 5, 5]
  Run 3: [5, 19, 5, 5, 5, 5, 5, 5, 5, 5]

2. Temperature Sampling

Temperature controls the “sharpness” of the probability distribution before sampling:

\[P_{\text{new}} = \text{softmax}(\text{logits} / T)\]

  • Temperature < 1.0: Sharper distribution (more like greedy)
  • Temperature = 1.0: Original distribution
  • Temperature > 1.0: Flatter distribution (more random)

Note: Temperature = 0 would cause division by zero. In practice, very low temperatures (e.g., 0.01) approximate greedy decoding, and many implementations treat temperature = 0 as an alias for greedy mode.

TipTry It!

Use the Probability Distribution Explorer above to see how temperature affects the distribution. Try T=0.3 (sharper, more confident) vs T=2.0 (flatter, more random).

# Generate with different temperatures
print("Generating with different temperatures:\n")

for temp in [0.3, 0.7, 1.0, 1.5]:
    print(f"Temperature = {temp}:")
    for i in range(3):
        torch.manual_seed(42 + i)
        out = generate(model, prompt, max_new_tokens=10, temperature=temp, do_sample=True)
        print(f"  Sample {i+1}: {out[0, 5:].tolist()}")
    print()
Generating with different temperatures:

Temperature = 0.3:
  Sample 1: [24, 25, 12, 27, 11, 47, 30, 30, 14, 30]
  Sample 2: [18, 40, 25, 21, 22, 23, 27, 23, 45, 0]
  Sample 3: [2, 5, 15, 38, 47, 5, 24, 31, 23, 5]

Temperature = 0.7:
  Sample 1: [24, 25, 29, 27, 11, 47, 30, 30, 14, 30]
  Sample 2: [18, 40, 25, 21, 18, 23, 27, 23, 45, 0]
  Sample 3: [2, 8, 15, 38, 47, 10, 24, 31, 23, 38]

Temperature = 1.0:
  Sample 1: [24, 25, 29, 27, 11, 47, 10, 30, 14, 30]
  Sample 2: [18, 40, 25, 21, 18, 23, 27, 23, 21, 0]
  Sample 3: [2, 8, 15, 38, 47, 10, 24, 31, 23, 38]

Temperature = 1.5:
  Sample 1: [24, 25, 29, 27, 11, 47, 10, 30, 14, 30]
  Sample 2: [18, 10, 25, 21, 18, 23, 27, 23, 21, 0]
  Sample 3: [2, 8, 15, 38, 47, 10, 24, 31, 23, 38]

3. Top-k Sampling

Only sample from the k most likely tokens - filters out unlikely tokens:

# Demonstrate top-k filtering
logits = torch.tensor([[1.0, 3.0, 0.5, 2.5, 0.0, 2.0, -1.0, 1.5]])
original_probs = F.softmax(logits, dim=-1)

print("Original probabilities:")
for i, p in enumerate(original_probs[0]):
    print(f"  Token {i}: {p.item():.3f}")

# Apply top-k filtering
for k in [3, 5]:
    filtered = top_k_filtering(logits.clone(), k)
    filtered_probs = F.softmax(filtered, dim=-1)

    print(f"\nAfter top-k = {k}:")
    for i, p in enumerate(filtered_probs[0]):
        if p > 0:
            print(f"  Token {i}: {p.item():.3f}")
Original probabilities:
  Token 0: 0.055
  Token 1: 0.403
  Token 2: 0.033
  Token 3: 0.244
  Token 4: 0.020
  Token 5: 0.148
  Token 6: 0.007
  Token 7: 0.090

After top-k = 3:
  Token 1: 0.506
  Token 3: 0.307
  Token 5: 0.186

After top-k = 5:
  Token 0: 0.058
  Token 1: 0.429
  Token 3: 0.260
  Token 5: 0.158
  Token 7: 0.096
TipTry It!

In the Probability Distribution Explorer, enable Top-k and adjust the slider. Watch how k=3 keeps only 3 tokens while k=20 keeps most.

4. Top-p (Nucleus) Sampling

Keep the smallest set of tokens whose cumulative probability exceeds p. This adapts to the distribution - keeps more tokens when uncertain, fewer when confident.

Key advantage: Top-p adapts to the distribution shape: - Peaked (confident): Keeps fewer tokens - Flat (uncertain): Keeps more tokens

# Demonstrate top-p filtering
logits = torch.tensor([[3.0, 2.0, 1.5, 1.0, 0.5, 0.0, -0.5, -1.0]])
probs = F.softmax(logits, dim=-1)[0]

# Sort and show cumulative probabilities
sorted_probs, sorted_idx = torch.sort(probs, descending=True)
cumulative = torch.cumsum(sorted_probs, dim=0)

print("Tokens sorted by probability:")
print(f"{'Token':<8} {'Prob':<10} {'Cumulative':<10}")
print("-" * 28)
for i, (idx, p, c) in enumerate(zip(sorted_idx, sorted_probs, cumulative)):
    marker = " <- cutoff (p=0.9)" if c.item() > 0.9 and (i == 0 or cumulative[i-1].item() <= 0.9) else ""
    print(f"{idx.item():<8} {p.item():<10.3f} {c.item():<10.3f}{marker}")
Tokens sorted by probability:
Token    Prob       Cumulative
----------------------------
0        0.524      0.524     
1        0.193      0.717     
2        0.117      0.834     
3        0.071      0.905      <- cutoff (p=0.9)
4        0.043      0.948     
5        0.026      0.975     
6        0.016      0.990     
7        0.010      1.000
TipTry It!

Compare how top-p behaves on “Peaked” vs “Flat” distributions in the explorer. Notice how it keeps fewer tokens when confident (peaked) and more when uncertain (flat).

5. Min-p Sampling

Top-p budgets a fixed slice of cumulative probability. Min-p (Nguyen et al., 2024) takes a different tack: it sets a relative floor on individual token probabilities, scaled by how confident the model is. Keep every token whose probability is at least min_p times the probability of the single most likely token:

\[ \text{threshold} = \texttt{min\_p} \cdot p_{\max}, \qquad \text{keep token } i \iff p_i \ge \text{threshold} \]

where \(p_{\max} = \max_i p_i\). Because the floor rides on \(p_{\max}\), the cutoff adapts to the model’s confidence for free:

  • Confident step — one token dominates, so \(p_{\max}\) is large, the floor is high, and only a handful of strong candidates survive. Coherent, low-risk.
  • Uncertain step — the distribution is flat, \(p_{\max}\) is small, the floor drops, and many plausible tokens stay in play. Room to be creative.

The top token always clears its own threshold (since \(\texttt{min\_p} \le 1\)), so the candidate set is never empty — which is exactly why min-p stays well-behaved at high temperature, where top-p tends to either admit a long tail of junk or clamp down too hard. Typical values are min_p = 0.05–0.1.

# min_p_filtering lives in generation.py alongside top_k / top_p
from generation import min_p_filtering

def survivors(logits, **kw):
    """How many tokens survive a filter (non -inf entries)."""
    if "min_p" in kw:
        f = min_p_filtering(logits.clone(), kw["min_p"])
    else:
        f = top_p_filtering(logits.clone(), kw["top_p"])
    return int((f > float("-inf")).sum().item())

# A CONFIDENT step (one token dominates) vs an UNCERTAIN step (flat)
confident = torch.tensor([[6.0, 2.0, 1.0, 0.5, 0.0, -0.5, -1.0, -1.5]])
uncertain = torch.tensor([[0.4, 0.3, 0.2, 0.1, 0.0, -0.1, -0.2, -0.3]])

for name, logits in [("confident", confident), ("uncertain", uncertain)]:
    p = F.softmax(logits, dim=-1)[0]
    thresh = 0.1 * p.max().item()          # min_p = 0.1
    print(f"{name:<10} p_max={p.max():.3f}  floor(min_p=0.1)={thresh:.3f}  "
          f"-> keep {survivors(logits, min_p=0.1)} tokens")
confident  p_max=0.967  floor(min_p=0.1)=0.097  -> keep 1 tokens
uncertain  p_max=0.173  floor(min_p=0.1)=0.017  -> keep 8 tokens

The same min_p=0.1 keeps only the top few tokens when the model is sure, but opens up when it is not — a single knob that tracks confidence. Watch the floor move as you reshape the distribution:

NoteKey Insight

Top-k fixes the count, top-p fixes the cumulative mass, and min-p fixes the relative height — a floor that automatically rises when the model is confident and falls when it is not. That single adaptive rule is why min-p keeps generations coherent at temperatures where top-p unravels.

TipTry It!
  1. Slide confidence up. As the distribution peaks, p_max grows, the dashed floor lifts, and the kept set shrinks to the few strong tokens — all at a fixed min_p.
  2. Slide it back down. On a flat distribution the floor drops and more tokens survive. min-p never has to be re-tuned per step.
  3. Push min_p toward 0.5. A high floor is aggressive — near-greedy on confident steps. Around 0.05–0.1 is the sweet spot most tools default to.

Combining Strategies

Each strategy has trade-offs: temperature affects the overall distribution shape, top-k provides a hard cutoff, and top-p adapts to model confidence. In practice, combining them often works better than any single approach:

# The typical generation pipeline
logits_example = torch.randn(1, vocab_size) * 2

# Step 1: Apply temperature
temperature = 0.7
logits_temp = logits_example / temperature

# Step 2: Apply top-k filtering
logits_topk = top_k_filtering(logits_temp.clone(), top_k=20)

# Step 3: Apply top-p filtering
logits_topp = top_p_filtering(logits_topk.clone(), top_p=0.9)

# Step 4: Sample from the distribution
probs = F.softmax(logits_topp, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)

print("Combined filtering pipeline:")
print(f"  Original vocab: {vocab_size} tokens")
print(f"  After top-k=20: {(F.softmax(logits_topk, dim=-1) > 0).sum().item()} tokens")
print(f"  After top-p=0.9: {(probs > 0).sum().item()} tokens")
print(f"  Sampled token: {next_token.item()}")
Combined filtering pipeline:
  Original vocab: 50 tokens
  After top-k=20: 20 tokens
  After top-p=0.9: 4 tokens
  Sampled token: 5
# Compare different strategy combinations
strategies = [
    ("Greedy", {"do_sample": False}),
    ("Temperature=0.5", {"temperature": 0.5, "do_sample": True}),
    ("Temperature=1.0", {"temperature": 1.0, "do_sample": True}),
    ("Top-k=5", {"top_k": 5, "do_sample": True}),
    ("Top-p=0.9", {"top_p": 0.9, "do_sample": True}),
    ("Min-p=0.1", {"min_p": 0.1, "do_sample": True}),
    ("Combined (T=0.7, k=20, p=0.9)", {"temperature": 0.7, "top_k": 20, "top_p": 0.9, "do_sample": True}),
]

print("Comparing strategies (3 samples each):\n")

for name, kwargs in strategies:
    print(f"{name}:")
    for i in range(3):
        torch.manual_seed(100 + i)
        out = generate(model, prompt, max_new_tokens=10, **kwargs)
        tokens = out[0, 5:].tolist()
        print(f"  {tokens}")
    print()
Comparing strategies (3 samples each):

Greedy:
  [5, 19, 5, 5, 5, 5, 5, 5, 5, 5]
  [5, 19, 5, 5, 5, 5, 5, 5, 5, 5]
  [5, 19, 5, 5, 5, 5, 5, 5, 5, 5]

Temperature=0.5:
  [17, 34, 14, 30, 27, 30, 24, 28, 9, 5]
  [17, 4, 33, 13, 30, 11, 46, 2, 23, 3]
  [0, 47, 13, 32, 49, 4, 40, 46, 44, 27]

Temperature=1.0:
  [3, 34, 14, 27, 27, 8, 24, 28, 9, 48]
  [17, 4, 33, 13, 30, 11, 46, 2, 23, 3]
  [0, 47, 13, 32, 49, 4, 28, 46, 15, 27]

Top-k=5:
  [17, 27, 33, 27, 27, 30, 40, 27, 40, 5]
  [17, 6, 19, 6, 30, 5, 23, 19, 6, 40]
  [6, 27, 30, 6, 27, 40, 40, 19, 23, 27]

Top-p=0.9:
  [3, 34, 14, 27, 27, 8, 24, 28, 9, 48]
  [17, 4, 33, 13, 30, 11, 46, 2, 23, 3]
  [0, 47, 13, 32, 49, 4, 28, 46, 15, 27]

Min-p=0.1:
  [3, 34, 14, 27, 27, 8, 24, 28, 9, 48]
  [17, 4, 33, 13, 30, 11, 46, 2, 23, 3]
  [0, 47, 13, 32, 49, 4, 28, 46, 15, 27]

Combined (T=0.7, k=20, p=0.9):
  [17, 27, 42, 27, 27, 30, 24, 27, 9, 48]
  [17, 40, 33, 33, 30, 21, 15, 33, 23, 3]
  [9, 35, 30, 17, 24, 40, 40, 46, 44, 27]

Measuring Output Diversity

Let’s quantify how different strategies affect output diversity:

def measure_diversity(model, prompt, num_samples=20, **kwargs):
    """Measure how diverse the generated outputs are."""
    outputs = []
    for i in range(num_samples):
        torch.manual_seed(i)
        out = generate(model, prompt, max_new_tokens=15, **kwargs)
        outputs.append(tuple(out[0].tolist()))

    unique = len(set(outputs))
    return unique / num_samples

# Compare diversity across settings
settings = [
    ("Greedy", {"do_sample": False}),
    ("Temp=0.3", {"temperature": 0.3, "do_sample": True}),
    ("Temp=0.7", {"temperature": 0.7, "do_sample": True}),
    ("Temp=1.0", {"temperature": 1.0, "do_sample": True}),
    ("Temp=1.5", {"temperature": 1.5, "do_sample": True}),
]

diversities = []
for name, kwargs in settings:
    div = measure_diversity(model, prompt, num_samples=20, **kwargs)
    diversities.append((name, div))
    print(f"{name}: {div*100:.0f}% unique outputs")
Greedy: 5% unique outputs
Temp=0.3: 100% unique outputs
Temp=0.7: 100% unique outputs
Temp=1.0: 100% unique outputs
Temp=1.5: 100% unique outputs

Choosing Parameters

Recommended settings for different use cases:

Goal Temperature Top-k Top-p
Code generation 0.2-0.4 10-20 0.8-0.9
Factual/deterministic 0.3-0.5 5-10 0.5-0.7
Coherent responses 0.7-0.9 20-50 0.85-0.92
Creative writing 0.8-1.2 40-100 0.9-0.95 
# Example settings for different applications
use_cases = {
    "Code generation": {"temperature": 0.2, "top_p": 0.9, "do_sample": True},
    "Balanced chat": {"temperature": 0.7, "top_p": 0.9, "do_sample": True},
    "Creative writing": {"temperature": 1.0, "top_p": 0.95, "do_sample": True},
    "Brainstorming": {"temperature": 1.5, "top_p": 0.95, "do_sample": True},
}

print("Sample outputs for different use cases:\n")

for name, kwargs in use_cases.items():
    print(f"{name}:")
    for i in range(2):
        torch.manual_seed(42 + i)
        out = generate(model, prompt, max_new_tokens=12, **kwargs)
        print(f"  {out[0, 5:].tolist()}")
    print()
Sample outputs for different use cases:

Code generation:
  [24, 25, 40, 27, 11, 9, 30, 30, 5, 30, 17, 23]
  [6, 40, 5, 21, 44, 23, 27, 23, 21, 0, 39, 27]

Balanced chat:
  [24, 25, 12, 27, 11, 9, 30, 30, 14, 30, 17, 14]
  [32, 40, 25, 21, 18, 23, 27, 23, 21, 0, 39, 27]

Creative writing:
  [24, 25, 12, 27, 11, 47, 30, 30, 14, 30, 17, 14]
  [18, 40, 25, 21, 18, 23, 27, 23, 21, 0, 39, 27]

Brainstorming:
  [24, 25, 12, 27, 11, 47, 30, 30, 14, 30, 17, 14]
  [18, 40, 25, 21, 18, 23, 27, 23, 21, 0, 39, 27]

Repetition Penalty

A common problem with text generation is repetition - the model gets stuck repeating the same tokens or phrases. Repetition penalties address this by reducing the probability of tokens that have already appeared.

The repetition penalty works as follows:

  • For tokens that have appeared before:
    • If the logit is positive, divide by the penalty (reduces probability)
    • If the logit is negative, multiply by the penalty (makes it more negative)
  • Penalty = 1.0 means no change
  • Penalty > 1.0 discourages repetition (common values: 1.1 - 1.5)
# Demonstrate repetition penalty
logits = torch.tensor([[2.0, 1.5, 1.0, 0.5, -0.5, -1.0]])
previous_tokens = torch.tensor([[0, 1, 4]])  # Tokens 0, 1, and 4 appeared

print("Original logits:")
for i, l in enumerate(logits[0]):
    marker = " (appeared)" if i in [0, 1, 4] else ""
    print(f"  Token {i}: {l.item():.2f}{marker}")

# Apply penalty
penalized = apply_repetition_penalty(logits, previous_tokens, penalty=1.5)

print("\nAfter repetition penalty (1.5):")
for i, l in enumerate(penalized[0]):
    marker = " (appeared)" if i in [0, 1, 4] else ""
    print(f"  Token {i}: {l.item():.2f}{marker}")

# Compare probabilities
orig_probs = F.softmax(logits, dim=-1)
new_probs = F.softmax(penalized, dim=-1)

print("\nProbability changes:")
for i in [0, 1, 2]:
    print(f"  Token {i}: {orig_probs[0,i].item():.3f} -> {new_probs[0,i].item():.3f}")
Original logits:
  Token 0: 2.00 (appeared)
  Token 1: 1.50 (appeared)
  Token 2: 1.00
  Token 3: 0.50
  Token 4: -0.50 (appeared)
  Token 5: -1.00

After repetition penalty (1.5):
  Token 0: 1.33 (appeared)
  Token 1: 1.00 (appeared)
  Token 2: 1.00
  Token 3: 0.50
  Token 4: -0.75 (appeared)
  Token 5: -1.00

Probability changes:
  Token 0: 0.429 -> 0.324
  Token 1: 0.260 -> 0.232
  Token 2: 0.158 -> 0.232
# Generate with and without repetition penalty
print("Generation without repetition penalty:")
out = generate_greedy(model, prompt, max_new_tokens=30)
tokens = out[0].tolist()
from collections import Counter
counts = Counter(tokens)
print(f"  Tokens: {tokens[5:]}")
print(f"  Most common: {counts.most_common(3)}")

print("\nGeneration with repetition penalty (1.3):")
out = generate(model, prompt, max_new_tokens=30, do_sample=False, repetition_penalty=1.3)
tokens = out[0].tolist()
counts = Counter(tokens)
print(f"  Tokens: {tokens[5:]}")
print(f"  Most common: {counts.most_common(3)}")
Generation without repetition penalty:
  Tokens: [5, 19, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 21, 5, 5, 5, 5, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 39]
  Most common: [(5, 16), (21, 12), (9, 1)]

Generation with repetition penalty (1.3):
  Tokens: [5, 19, 5, 5, 6, 5, 5, 5, 21, 5, 5, 5, 12, 5, 21, 23, 5, 5, 21, 21, 21, 39, 21, 21, 21, 21, 21, 12, 12, 12]
  Most common: [(5, 12), (21, 10), (12, 4)]

When to use repetition penalty:

  • Always for open-ended generation (stories, chat)
  • Less critical for short, structured outputs (classification, extraction)
  • Typical values: 1.1 for mild effect, 1.3-1.5 for stronger effect
  • Too high (> 2.0) can make outputs incoherent

Stop Conditions

Generation requires clear stopping conditions. Two conditions stop generation:

  1. Maximum length (max_new_tokens) - Hard limit on generated tokens
  2. EOS token (eos_token_id) - Stop when a special end-of-sequence token is generated
# Demonstrate EOS stopping
# In real models, EOS is a special token. Here we use token 42 as our "EOS"
eos_id = 42

print(f"Generating with EOS token = {eos_id}")
print(f"(Generation stops early if token {eos_id} is produced)")

# Without EOS
out_no_eos = generate_greedy(model, prompt, max_new_tokens=20)
print(f"\nWithout EOS check: {len(out_no_eos[0]) - 5} new tokens generated")
print(f"  Tokens: {out_no_eos[0, 5:].tolist()}")

# With EOS (may stop early if 42 is generated)
out_with_eos = generate_greedy(model, prompt, max_new_tokens=20, eos_token_id=eos_id)
print(f"\nWith EOS check: {len(out_with_eos[0]) - 5} new tokens generated")
print(f"  Tokens: {out_with_eos[0, 5:].tolist()}")

if len(out_with_eos[0]) < len(out_no_eos[0]):
    print(f"  (Stopped early due to EOS token)")
Generating with EOS token = 42
(Generation stops early if token 42 is produced)

Without EOS check: 20 new tokens generated
  Tokens: [5, 19, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 21, 5, 5, 5, 5, 21, 21]

With EOS check: 20 new tokens generated
  Tokens: [5, 19, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 21, 5, 5, 5, 5, 21, 21]

Practical notes on stopping:

  • Always set a reasonable max_new_tokens to prevent runaway generation
  • EOS tokens are essential for chat/instruction models to indicate response completion
  • Batched generation continues until ALL sequences hit a stop condition
  • Some APIs support multiple stop sequences (not just EOS)

Constrained Decoding: Guaranteeing Valid Output

Every strategy so far — greedy, temperature, top-k, top-p, min-p, repetition penalty — is the same move: reshape the logits, then sample. Constrained decoding is one more logit transform, with a new rule: at each step, keep only the tokens that would leave the output valid, and mask the rest to \(-\infty\). Because a masked token has probability exactly zero, the model cannot emit it. The result is valid by construction, not by hoping and re-checking afterward.

Why it matters for LLMs:

  • Structured output. When you need a date, an enum, or a JSON object your program will parse, “usually valid” is not good enough — one stray token breaks the parser. Constrained decoding makes malformed output impossible.
  • Function calling & agents. A tool call the runtime can execute must match the tool’s schema exactly. This is how “JSON mode” and function-calling APIs guarantee parseable calls.
  • It’s orthogonal to the model. The constraint lives in a mask, not in the weights — no fine-tuning. The same untrained-or-trained generate() loop gains guaranteed structure for free.

Here is the problem, made concrete. Ask an unconstrained sampler for a date and one bad token derails it:

import sys
sys.path.insert(0, '..')
import torch
from constrained import (
    DIGITS, field_class, choice, template_fsm,
    date_fsm, json_answer_fsm,
    allowed_token_ids, build_token_index, apply_constraint,
    constrained_decode,
)

# A toy character vocabulary (id -> string). In a real model these are BPE
# tokens; here single characters keep the picture readable.
vocab = {i: s for i, s in enumerate(list("0123456789-/: ") + ["true", "false"])}
eos_id = len(vocab)
vocab_size = eos_id + 1

# The "model": random logits, standing in for an untrained network.
g = torch.Generator().manual_seed(0)
def random_logits(ids):
    return torch.randn(vocab_size, generator=g)

# Unconstrained greedy pick over 10 steps -> garbage that no parser accepts.
junk = "".join(
    vocab[int(torch.argmax(random_logits(None)[:len(vocab)]))] for _ in range(10)
)
print(f"Unconstrained: {junk!r}")
print(f"  valid date? {date_fsm().accepts(junk)}")
Unconstrained: '7-8:61: 01'
  valid date? False

Now add the constraint. The same random logits, masked to a date grammar, can only ever spell a date:

ids, text = constrained_decode(
    random_logits, date_fsm(), vocab,
    eos_token_id=eos_id, max_tokens=12,
)
print(f"Constrained:   {text!r}")
print(f"  valid date? {date_fsm().accepts(text)}")
Constrained:   '7977-78-52'
  valid date? True

The model’s preferences still choose which valid date — the constraint only removes the invalid options. This lives in constrained.py.

A Grammar Is a Finite-State Machine

What does “valid so far” mean? We encode it as a finite-state machine (FSM): a set of states, and for each state a table saying which next character is legal and where it leads. A character with no entry is forbidden — the machine “dies.” A schema or regex compiles into such a machine.

A YYYY-MM-DD date is a straight chain: four digit-steps, a dash, two digits, a dash, two digits. At every state the allowed characters are exactly the out-edges — digits at a number position, - at a separator. That set is what we will turn into a token mask.

fsm = date_fsm()
print("At the start, the only legal characters are:")
print(" ", sorted(fsm.allowed_chars(fsm.start)))            # digits
after_year = fsm.advance(fsm.start, "2024")
print("After '2024', the only legal character is:")
print(" ", sorted(fsm.allowed_chars(after_year)))           # ['-']
print("Is '2024-01-31' accepted?", fsm.accepts("2024-01-31"))
print("Is '2024/01/31' accepted?", fsm.accepts("2024/01/31"))
At the start, the only legal characters are:
  ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
After '2024', the only legal character is:
  ['-']
Is '2024-01-31' accepted? True
Is '2024/01/31' accepted? False

We build these machines from a tiny grammar toolkit — literal("..."), field_class(chars, min, max), and choice([...]) — that template_fsm concatenates left to right. A fixed JSON schema is still regular, so it is just an FSM too:

schema = json_answer_fsm()   # {"answer": <1-3 digits>, "correct": true|false}
print(schema.accepts('{"answer": 42, "correct": true}'))    # True
print(schema.accepts('{"answer": 1234, "correct": true}'))  # False: >3 digits
print(schema.accepts('{"answer": 42, "correct": maybe}'))   # False: bad enum
True
False
False

Step through the date automaton below: watch the active state advance as each character is consumed, and the allowed-character set shrink to just what’s legal next.

From Characters to Tokens

The model does not emit characters — it emits tokens over a fixed vocabulary. So we translate “allowed characters” into “allowed token ids”: a token is legal from state \(s\) if feeding its whole string through the FSM keeps it alive. Walking character by character is what makes token boundaries a non-issue — a multi-character token that straddles a boundary (say "}\n") is allowed only if every one of its characters is legal in order.

fsm = date_fsm()
allowed = allowed_token_ids(fsm, fsm.start, vocab, eos_token_id=eos_id)
print("Legal token strings at the start:", [vocab[t] for t in allowed])
# -> only the digits; '-', '/', ' ', 'true', 'false' are all masked out.
Legal token strings at the start: ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']

Because the legal set depends only on the grammar and vocabulary — never on the logits — we can precompute it for every state, once. This is the key idea from Willard & Louf’s Efficient Guided Generation (the Outlines library): build an index up front, then each decoding step is a dict lookup instead of a scan over the whole vocabulary.

index, goto = build_token_index(fsm, vocab)
print(f"Precomputed allowed-token sets for {len(index)} states.")
print("State 0 (start) allows ids:", index[fsm.start])
Precomputed allowed-token sets for 11 states.
State 0 (start) allows ids: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

Masking itself is exactly the -inf trick from top_k_filtering, so forbidden tokens get zero probability after softmax:

logits = torch.randn(vocab_size, generator=torch.Generator().manual_seed(1))
masked = apply_constraint(logits, allowed)
probs = torch.softmax(masked, dim=-1)
print("Total probability on legal (digit) tokens:", round(probs[allowed].sum().item(), 6))
print("Probability on the illegal 'true' token:  ", round(probs[vocab_size - 3].item(), 6))
Total probability on legal (digit) tokens: 1.0
Probability on the illegal 'true' token:   0.0

The decode loop then just threads these together — mask, sample, advance the FSM — and this is the whole of constrained_decode. Applied to the JSON schema, the output is guaranteed parseable:

import json
schema = json_answer_fsm()
jvocab = {i: s for i, s in enumerate(
    ['{', '}', '"', ':', ',', ' ', 'answer', 'correct', 'true', 'false']
    + list("0123456789"))}
jeos = len(jvocab)
jg = torch.Generator().manual_seed(3)

def jlogits(_ids):
    return torch.randn(jeos + 1, generator=jg)

_, out = constrained_decode(
    jlogits, schema, jvocab, eos_token_id=jeos,
    do_sample=True, generator=jg, max_tokens=40,
)
print("Generated:", out)
print("json.loads parses it:", json.loads(out))   # never raises
Generated: {"answer": 857, "correct": true}
json.loads parses it: {'answer': 857, 'correct': True}
NoteKey Insight

Constrained decoding does not make the model smarter — it makes whole regions of output space unreachable. The mask is applied at every single step, so there is no “mostly valid” failure mode: the output either matches the grammar or generation stops. The model’s logits still decide which valid string you get.

Interactive: Watch the Mask

Below is the real trace of a constrained run over the date grammar. Step through it: at each position the vocabulary is split into legal tokens (lit) and forbidden tokens (greyed, logit \(= -\infty\)), and the chosen token is ringed. Notice how the legal set flips between “all digits” and “just -” as the automaton advances — the model never even sees the invalid options.

TipTry This
  1. Read the flip. Drag the step slider and watch the legal set toggle between the ten digits and the lone -. That toggle is the grammar.
  2. Break the parser. In the JSON demo above, change field_class(DIGITS, 1, 3) to field_class(DIGITS, 1, 5) inside json_answer_fsm and confirm larger answers become legal — the schema is the code.
  3. Swap the constraint for a keyword. Build template_fsm([choice(["yes", "no", "maybe"])]) and decode with random logits; every run lands on one of the three words, never a typo.

Common Pitfalls

WarningConstraint is not correctness

A grammar guarantees the shape of the output, not its truth. {"answer": 7, "correct": true} is valid JSON even when the real answer is 42. Constrained decoding removes parser errors, not reasoning errors — pair it with the training and evaluation from earlier modules.

WarningWatch for an empty legal set

If your grammar and vocabulary disagree — a required character that no token can produce — every token gets masked and generation stalls. constrained_decode stops when the legal set is empty; in practice, make sure the tokenizer can spell every string the grammar allows (real systems split multi-character tokens when needed).

WarningFixed schemas are regular; nested ones are not

A JSON object with a fixed set of fields is a regular language, so a plain FSM suffices. Arbitrary nesting (objects inside arrays inside objects) needs a stack — a pushdown automaton or a full grammar (llama.cpp’s GBNF, Grammar-Constrained Decoding). The token-masking idea is identical; only the “what’s legal next” bookkeeping grows a stack.

KV-Cache Optimization

In generation, we process one new token at a time. Without optimization, we’d recompute attention for ALL previous tokens every step - wasting computation!

KV-Cache stores Key and Value projections from previous tokens:

  • Without cache: O(n^2) per token, O(n^3) total for n tokens
  • With cache: O(n) per token, O(n^2) total for n tokens

KV-caching is crucial for fast inference.

How KV-Cache Works

In attention, we compute: \[\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V\]

For autoregressive generation:

  1. First forward pass (prompt): Compute K, V for all prompt tokens and cache them
  2. Each new token: Only compute Q, K, V for the new token
  3. Attention: New Q attends to cached K, V plus new K, V
  4. Update cache: Append new K, V to the cache

From Scratch: Caching Keys and Values

An optimization is only trustworthy if it changes speed, not answers. So rather than describe the cache, let’s build one and prove it produces the exact same output. Two small pieces live in generation.py:

  • KVCache — a growable store for one layer’s keys and values.
  • CachedCausalSelfAttention — multi-head causal attention that, given a cache, projects K, V for the new token only and attends over the full history.
from generation import KVCache, CachedCausalSelfAttention, demonstrate_kv_cache

# One attention layer; eval() so there is no dropout randomness.
attn = CachedCausalSelfAttention(embed_dim=32, num_heads=4).eval()

# A short sequence of token embeddings (batch=1, seq=6, embed=32).
x = torch.randn(1, 6, 32)
print(f"Input sequence: {tuple(x.shape)}")
Input sequence: (1, 6, 32)

The cache is just remembered K and V. Watch it grow one token at a time — the same thing that happens during decoding:

cache = KVCache()
print(f"Start: {len(cache)} tokens cached")

for t in range(x.size(1)):
    # Feed ONLY the new token; the layer appends its K, V to the cache.
    attn(x[:, t:t + 1, :], cache=cache)
    print(f"  after token {t}: {len(cache)} tokens cached")
Start: 0 tokens cached
  after token 0: 1 tokens cached
  after token 1: 2 tokens cached
  after token 2: 3 tokens cached
  after token 3: 4 tokens cached
  after token 4: 5 tokens cached
  after token 5: 6 tokens cached

The key/value projection cost is the whole point. Naively, step t re-projects all t+1 tokens; with a cache it projects exactly one:

seq_len = x.size(1)
naive_kv = sum(range(1, seq_len + 1))   # 1 + 2 + ... + n
cached_kv = seq_len                      # one new token per step

print(f"K/V projections over {seq_len} steps:")
print(f"  Without cache: {naive_kv}")
print(f"  With cache:    {cached_kv}")
print(f"  Redundant work avoided: {naive_kv - cached_kv} projections")
K/V projections over 6 steps:
  Without cache: 21
  With cache:    6
  Redundant work avoided: 15 projections

Proving It’s the Same Computation

The reason caching is free correctness-wise: the keys and values of past tokens never depend on future tokens, so storing them changes nothing about the result. Here we run the identical layer two ways — a full forward pass over the whole sequence, and token-by-token through the cache — and compare:

# 1. Full forward: recompute attention over the entire sequence at once.
full = attn(x)

# 2. Incremental: one token at a time, reusing cached K, V.
cache = KVCache()
steps = [attn(x[:, t:t + 1, :], cache=cache) for t in range(x.size(1))]
incremental = torch.cat(steps, dim=1)

max_diff = (full - incremental).abs().max().item()
print(f"Full-forward shape:  {tuple(full.shape)}")
print(f"Incremental shape:   {tuple(incremental.shape)}")
print(f"Max |difference|:    {max_diff:.2e}")
print(f"Identical output?    {torch.allclose(full, incremental, atol=1e-5)}")
Full-forward shape:  (1, 6, 32)
Incremental shape:   (1, 6, 32)
Max |difference|:    1.79e-07
Identical output?    True

The difference is at the level of floating-point rounding (~1e-7), not algorithmic. demonstrate_kv_cache() wraps this same check with the op-count summary:

_ = demonstrate_kv_cache(seq_len=6, embed_dim=32, num_heads=4)
============================================================
KV-CACHE EQUIVALENCE
============================================================

Sequence length: 6, embed_dim: 32, heads: 4
Full-forward output shape:  (1, 6, 32)
Cached (incremental) shape: (1, 6, 32)

Max |difference|: 1.19e-07  (identical up to float rounding)

Key/Value projections computed while decoding:
  Without cache: 21 (recomputes the whole past each step)
  With cache:    6 (only the new token each step)
  Saved:         15 redundant projections
NoteKey Insight

A KV-cache is not an approximation — it’s memoization. Past keys and values are a pure function of past tokens, so caching them and only projecting the new token yields bit-for-bit the same output while turning O(n) redundant work per step into O(1). Speed changes; answers do not.

The trick that makes prefill (many prompt tokens at once) and decode (one new token) share one code path is an offset-aware causal mask: a query at position len(cache) + i may attend to key positions 0 … len(cache) + i. When the cache is empty this is the familiar lower-triangular mask; mid-generation it lets the single new query see the entire cached history.

WarningReset the cache between sequences

A KVCache holds the K, V of one specific sequence. Reuse it for a new prompt without calling cache.reset() and the new tokens will attend to stale history from the previous generation — silently wrong output, not a crash. One cache per sequence, per layer. Real models also cap the cache at the context length; past that, entries must be dropped or the positions re-scaled (see Module 04’s RoPE).

Memory Tradeoff

KV-cache trades memory for speed:

Aspect Without Cache With Cache
Computation per token O(n^2) O(n)
Memory O(1) extra O(n * layers * d)
Total time for n tokens O(n^3) O(n^2)

For a model with:

  • 32 layers, d_model = 4096, 8K context
  • KV cache size = 2 (K and V) × 32 × 4096 × 8192 × 2 bytes (float16) = ~4GB per sequence

Long-context models demand significant GPU memory for this reason.

Practical Considerations

  • Prompt processing: First pass processes entire prompt (batches efficiently)
  • Generation: Subsequent tokens are generated one at a time (memory-bound)
  • Batch size tradeoff: Larger batches amortize overhead but need more KV-cache memory
  • Context length: Longer contexts need more cache memory per sequence

Note: Our generate() function prioritizes clarity over efficiency; production implementations use KV-caching.

Exercises

Exercise 1: Temperature Exploration

Experiment with extreme temperatures and observe the output behavior:

# Try very low and very high temperatures
print("Extreme temperature exploration:\n")

for temp in [0.1, 0.5, 1.0, 2.0, 5.0]:
    print(f"Temperature = {temp}:")
    outputs = set()
    for i in range(5):
        torch.manual_seed(i)
        out = generate(model, prompt, max_new_tokens=8, temperature=temp, do_sample=True)
        outputs.add(tuple(out[0, 5:].tolist()))
    print(f"  {len(outputs)}/5 unique sequences")
    # Show one sample
    torch.manual_seed(42)
    sample = generate(model, prompt, max_new_tokens=8, temperature=temp, do_sample=True)
    print(f"  Sample: {sample[0, 5:].tolist()}")
    print()
Extreme temperature exploration:

Temperature = 0.1:
  5/5 unique sequences
  Sample: [19, 19, 5, 27, 11, 19, 30, 30]

Temperature = 0.5:
  5/5 unique sequences
  Sample: [24, 25, 12, 27, 11, 47, 30, 30]

Temperature = 1.0:
  5/5 unique sequences
  Sample: [24, 25, 29, 27, 11, 47, 10, 30]

Temperature = 2.0:
  5/5 unique sequences
  Sample: [24, 25, 29, 27, 11, 47, 10, 30]

Temperature = 5.0:
  5/5 unique sequences
  Sample: [24, 25, 29, 27, 11, 47, 10, 30]

Exercise 2: Top-k vs Top-p

Compare how top-k and top-p behave differently:

# Compare filtering approaches
print("Top-k vs Top-p filtering:\n")

# Create a bimodal distribution (two likely options)
bimodal_logits = torch.tensor([[3.0, 3.0, -1.0, -1.0, -2.0, -2.0, -3.0, -3.0, -4.0, -4.0]])

print("Original probabilities (bimodal - two equally likely tokens):")
bimodal_probs = F.softmax(bimodal_logits, dim=-1)
for i, p in enumerate(bimodal_probs[0][:5]):
    print(f"  Token {i}: {p.item():.3f}")

# Top-k=2 keeps exactly 2 tokens
topk_filtered = top_k_filtering(bimodal_logits.clone(), 2)
topk_probs = F.softmax(topk_filtered, dim=-1)

# Top-p=0.5 adapts to distribution
topp_filtered = top_p_filtering(bimodal_logits.clone(), 0.5)
topp_probs = F.softmax(topp_filtered, dim=-1)

print(f"\nTop-k=2 keeps: {(topk_probs > 0).sum().item()} tokens")
print(f"Top-p=0.5 keeps: {(topp_probs > 0).sum().item()} tokens")

print("\n** Key insight: Top-k always keeps exactly k tokens.")
print("   Top-p adapts: it may keep fewer tokens if one dominates.")
Top-k vs Top-p filtering:

Original probabilities (bimodal - two equally likely tokens):
  Token 0: 0.486
  Token 1: 0.486
  Token 2: 0.009
  Token 3: 0.009
  Token 4: 0.003

Top-k=2 keeps: 2 tokens
Top-p=0.5 keeps: 2 tokens

** Key insight: Top-k always keeps exactly k tokens.
   Top-p adapts: it may keep fewer tokens if one dominates.

Exercise 3: Repetition Penalty Effects

Explore how different repetition penalty values affect generation:

# Compare different repetition penalties
print("Repetition penalty comparison (greedy decoding, 40 tokens):\n")

for penalty in [1.0, 1.1, 1.3, 1.5, 2.0]:
    out = generate(model, prompt, max_new_tokens=40, do_sample=False, repetition_penalty=penalty)
    tokens = out[0, 5:].tolist()

    # Count unique tokens
    unique_ratio = len(set(tokens)) / len(tokens)

    print(f"Penalty = {penalty}:")
    print(f"  Unique tokens: {len(set(tokens))}/{len(tokens)} ({unique_ratio*100:.0f}%)")
    print(f"  First 15: {tokens[:15]}")
    print()
Repetition penalty comparison (greedy decoding, 40 tokens):

Penalty = 1.0:
  Unique tokens: 5/40 (12%)
  First 15: [5, 19, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 21, 5]

Penalty = 1.1:
  Unique tokens: 5/40 (12%)
  First 15: [5, 19, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 21, 5]

Penalty = 1.3:
  Unique tokens: 9/40 (22%)
  First 15: [5, 19, 5, 5, 6, 5, 5, 5, 21, 5, 5, 5, 12, 5, 21]

Penalty = 1.5:
  Unique tokens: 15/40 (38%)
  First 15: [5, 19, 5, 6, 27, 5, 5, 23, 3, 5, 5, 5, 30, 21, 5]

Penalty = 2.0:
  Unique tokens: 17/40 (42%)
  First 15: [5, 19, 6, 23, 30, 40, 5, 5, 21, 5, 5, 5, 46, 3, 5]

Exercise 4: Observing Repetition in Long Generation

Observe how greedy decoding can lead to repetition:

# Generate longer sequences to see repetition patterns
print("Long greedy generation (may show repetition):\n")

# Generate more tokens
long_output = generate_greedy(model, prompt, max_new_tokens=50)
tokens = long_output[0].tolist()

print(f"Generated sequence ({len(tokens)} tokens):")
print(tokens)

# Count token frequency
from collections import Counter
token_counts = Counter(tokens)
print(f"\nMost common tokens:")
for token, count in token_counts.most_common(5):
    print(f"  Token {token}: {count} times ({count/len(tokens)*100:.1f}%)")
Long greedy generation (may show repetition):

Generated sequence (55 tokens):
[9, 33, 35, 29, 24, 5, 19, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 21, 5, 5, 5, 5, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 39, 39, 39, 12, 12, 12, 12, 12, 12, 39, 39, 12, 12, 39, 39, 12, 39, 39, 39, 39, 39]

Most common tokens:
  Token 5: 16 times (29.1%)
  Token 21: 12 times (21.8%)
  Token 39: 12 times (21.8%)
  Token 12: 9 times (16.4%)
  Token 9: 1 times (1.8%)

The Complete Generation Function

Here’s the main generation function from our codebase:

# Display the signature and key parts
import inspect
from generation import generate

print("generate() function signature:")
print(inspect.signature(generate))
print()
print("Key parameters:")
print("  - model: The language model")
print("  - prompt_tokens: Starting sequence (batch, seq_len)")
print("  - max_new_tokens: How many tokens to generate")
print("  - temperature: Distribution sharpness (default 1.0)")
print("  - top_k: Filter to top k tokens (optional)")
print("  - top_p: Nucleus sampling threshold (optional)")
print("  - min_p: Relative probability floor, min_p * p_max (optional)")
print("  - do_sample: If False, use greedy decoding")
print("  - eos_token_id: Stop token (optional)")
print("  - repetition_penalty: Penalize repeated tokens (default 1.0)")
generate() function signature:
(model: torch.nn.modules.module.Module, prompt_tokens: torch.Tensor, max_new_tokens: int = 50, temperature: float = 1.0, top_k: Optional[int] = None, top_p: Optional[float] = None, min_p: Optional[float] = None, do_sample: bool = True, eos_token_id: Optional[int] = None, repetition_penalty: float = 1.0, verbose: bool = False) -> torch.Tensor

Key parameters:
  - model: The language model
  - prompt_tokens: Starting sequence (batch, seq_len)
  - max_new_tokens: How many tokens to generate
  - temperature: Distribution sharpness (default 1.0)
  - top_k: Filter to top k tokens (optional)
  - top_p: Nucleus sampling threshold (optional)
  - min_p: Relative probability floor, min_p * p_max (optional)
  - do_sample: If False, use greedy decoding
  - eos_token_id: Stop token (optional)
  - repetition_penalty: Penalize repeated tokens (default 1.0)

Summary

Key takeaways from this module:

  1. Autoregressive generation: Produce tokens one at a time, feeding each back as input
  2. Greedy decoding: Always pick the max - deterministic but often repetitive
  3. Temperature: Controls randomness - lower is more focused, higher is more diverse
  4. Top-k sampling: Limits choices to k most likely tokens
  5. Top-p (nucleus) sampling: Adapts to distribution shape - keeps more tokens when uncertain
  6. Min-p sampling: A relative probability floor (min_p × p_max) that tracks model confidence automatically - stays coherent at high temperature where top-p unravels
  7. Repetition penalty: Reduces probability of previously-generated tokens to prevent loops
  8. Stop conditions: Use EOS tokens and max length to control when generation ends
  9. Combine strategies: Temperature + top-p (or min-p) + repetition penalty is common in practice
  10. KV-cache: Essential optimization - trades memory for O(n) speedup per token

Common Pitfalls

Problem Cause Solution
Repetitive output Greedy decoding or low temperature Use sampling, repetition penalty
Incoherent nonsense Temperature too high Lower temperature, use top-p or min-p
Cuts off mid-sentence max_new_tokens too low Increase limit, ensure EOS handling
Slow generation No KV-cache Implement caching (production)
Out of memory Long context + large batch Reduce batch size or context

Conclusion

Congratulations! You’ve completed the Learn LLM series. You now understand all the building blocks of a language model:

  1. Tensors: The fundamental data structure
  2. Autograd: Automatic differentiation for training
  3. Tokenization: Converting text to numbers
  4. Embeddings: Learned vector representations
  5. Attention: The mechanism that lets tokens interact
  6. Transformer: The complete architecture
  7. Training: How models learn from data
  8. Generation: How to produce text from trained models

What’s Next?

  • Module 09: Efficient Attention — shrink the KV-cache you just built (MQA/GQA) and stream the softmax without the N×N matrix (FlashAttention)
  • Check out the minigpt directory to see everything assembled into a working model
  • Train your own small language model on real data
  • Explore the Going Deeper resources for advanced topics

Going Deeper

Core Papers:

Advanced Topics (not covered here):

  • Beam Search: Maintain k best partial sequences; better for translation, worse for open-ended generation
  • Speculative Decoding: Use a small draft model to propose tokens, verify with large model in parallel (built from scratch in Module 16)
  • Contrastive Decoding: Compare probabilities from expert and amateur models

Structured Generation — constraining outputs to valid JSON, code syntax, or grammar rules — is built from scratch above in Constrained Decoding:

Practical Resources: