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, typical, ε/η, Mirostat)
  • Filter by information instead of probability with typical sampling — the one strategy that can drop the most-likely token
  • See truncation as desmoothing, and set an entropy-adaptive probability floor with η-sampling
  • Close a feedback loop on perplexity with Mirostat — set the target surprise directly instead of guessing a cutoff
  • Run beam search as a bounded tree search over sequences, and length-normalize it
  • Break repetition deterministically with contrastive search’s degeneration penalty
  • Use repetition penalties to prevent degenerate outputs
  • Constrain decoding to a grammar so output is valid JSON/dates/enums by construction
  • Compile any regex into a constraint FSM from scratch (Thompson NFA → subset-construction DFA)
  • Build a pushdown automaton so constrained output can nest to any depth (real JSON)
  • 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.9.1

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: [3, 23, 41, 36, 5]

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 27: 2.82%
  Token 18: 2.72%
  Token 45: 2.64%
  Token 14: 2.59%
  Token 0: 2.57%

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: [3, 23, 41, 36, 5]
Generated: [27, 44, 27, 44, 27, 44, 44, 44, 45, 45, 49, 49, 45, 49, 49]
# 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: [27, 44, 27, 44, 27, 44, 44, 44, 45, 45]
  Run 2: [27, 44, 27, 44, 27, 44, 44, 44, 45, 45]
  Run 3: [27, 44, 27, 44, 27, 44, 44, 44, 45, 45]

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: [25, 25, 12, 27, 11, 47, 10, 30, 14, 30]
  Sample 2: [18, 10, 25, 21, 18, 23, 27, 23, 45, 0]
  Sample 3: [31, 14, 15, 38, 47, 10, 45, 47, 23, 5]

Temperature = 0.7:
  Sample 1: [25, 25, 29, 27, 11, 47, 10, 30, 14, 30]
  Sample 2: [18, 10, 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: [25, 25, 29, 27, 11, 47, 10, 30, 14, 30]
  Sample 2: [18, 10, 25, 21, 18, 23, 27, 23, 45, 0]
  Sample 3: [2, 8, 15, 38, 47, 10, 24, 31, 23, 38]

Temperature = 1.5:
  Sample 1: [37, 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.

6. Typical Sampling

Top-k, top-p, and min-p all rank tokens by probability and keep a prefix from the top. Locally typical sampling (Meister et al., TACL 2023) is the one member of the family that ranks by information instead — and it is the only filter that will ever throw away the most likely token.

The idea is information-theoretic. Fluent human text tends to carry a fairly steady amount of information per word: writers are neither maximally surprising (gibberish) nor maximally predictable (a word everyone already expects). Formally, the per-token surprisal -\log p(x) tends to sit close to the distribution’s own conditional entropy

H(p) = -\sum_{x} p(x)\,\log p(x),

which is just the expected surprisal. So instead of keeping the highest-probability tokens, typical sampling keeps the tokens whose surprisal is closest to the entropy — the smallest set, ordered by ascending deviation

\text{dev}(x) = \bigl|\,\underbrace{-\log p(x)}_{\text{surprisal}} - \underbrace{H(p)}_{\text{expected surprisal}}\,\bigr|,

whose cumulative probability first reaches a mass \tau (typical_p). Then it renormalizes and samples from that set.

The consequence is the surprising part. When the distribution is peaked, the top token’s surprisal is far below H — it is “too predictable” — so typical sampling can drop it, exactly as readily as it drops the deep tail whose surprisal is far above H. Top-p and top-k, which always take the highest-probability tokens first, can never do this.

# typical_filtering lives in generation.py alongside top_k / top_p / min_p
from generation import typical_filtering, demonstrate_typical

# A high-entropy step: token 0 alone holds ~42% of the mass, over a band of 24
# equally-plausible tokens. Its surprisal sits far below the entropy.
out = demonstrate_typical()
============================================================
TYPICAL SAMPLING vs TOP-P (mass = 0.5)
============================================================

Entropy H(p) = 2.527 nats
Argmax = token 0 (p=0.419), surprisal 0.871 (far below H — 'too predictable')

Top-p   kept  5 tokens, argmax kept: True
Typical kept 21 tokens, argmax kept: False  <- drops the peak

The argmax carries 42% of the probability, yet typical sampling excludes it — its surprisal (0.87 nats) is nowhere near the distribution’s entropy (2.53 nats), so it is too predictable to be typical. Top-p, at the same mass, keeps it without question. Watch the two filters disagree as you reshape the distribution:

NoteKey Insight

Top-k, top-p, and min-p all measure tokens by probability and keep a prefix from the top. Typical sampling measures them by information — distance of surprisal from the entropy — so it trims tokens from both ends: the deep tail (too surprising) and the over-confident head (too predictable). It is the only strategy that can refuse the argmax.

TipTry It!
  1. Raise peakedness. As one token takes over, its surprisal drops far below H, its bar turns dim (dropped) while the outline (the argmax) still wears a top-p ●. That gap — kept by top-p, dropped by typical — is the whole idea.
  2. Lower peakedness. On a flat distribution every surprisal is near H, the two sets nearly agree, and the argmax stays typical. Typical sampling only diverges from top-p when the model is confident.
  3. Sweep \tau. Smaller \tau keeps a tighter band around the entropy; the paper recommends about 0.2 for focused generation (e.g. stories) up to 0.95 for abstractive summarization.

7. Entropy-Aware Truncation: ε- and η-sampling

Every filter so far reshapes the model’s distribution and hopes the reshaping helps. Truncation as desmoothing (Hewitt, Manning & Liang, 2022) says why it helps, and turns that into two more filters.

The idea starts one module back. A trained model is a smoothed estimate of the true next-token distribution. Cross-entropy training (m07) never lets a token reach probability exactly zero — the loss -\log p explodes if a token that does occur was assigned zero — so the model spreads a thin film of probability across the whole vocabulary, including tokens that should have had none. Sampling from the raw distribution therefore samples that smoothing noise: the long tail of not-actually-plausible tokens is where degenerate, incoherent text comes from. In this light every truncation method — top-p, min-p, typical — is doing the same job: desmoothing, cutting the tokens the model only kept alive to satisfy the loss. ε- and η-sampling do it with the most direct rule of all — a probability floor.

ε-sampling keeps every token whose probability clears a fixed absolute threshold:

\mathcal{A}_\varepsilon = \{\, x : p(x) > \varepsilon \,\}.

This is the whole method. It differs from min-p in one crucial way: min-p’s floor is relative (min_p × p_max, so it scales with confidence), while ε’s floor is absolute — the same number at every step. That is its strength (nothing simpler) and its flaw. The right floor depends on how uncertain the step is: a fixed \varepsilon that cleanly desmooths a peaked step will over-truncate a legitimately flat one, where every genuine continuation is individually improbable.

η-sampling fixes exactly that flaw by letting the floor breathe with the entropy. It keeps p(x) > \eta where

\eta = \min\!\bigl(\varepsilon,\; \sqrt{\varepsilon}\,\cdot e^{-H(p)}\bigr), \qquad H(p) = -\sum_x p(x)\log p(x)

is the same conditional entropy typical sampling used one section up (the paper fixes \alpha = \sqrt{\varepsilon} to leave a single knob). Two clean facts fall straight out of that \min, and they are all you need to reason about η:

  • η is never above ε, so \mathcal{A}_\eta \supseteq \mathcal{A}_\varepsilon — η always keeps a superset of what ε keeps.
  • The two are identical below a crossover entropy H^\star = \tfrac{1}{2}\ln\frac{1}{\varepsilon}. When H \le H^\star the term \sqrt{\varepsilon}\,e^{-H} \ge \varepsilon, the \min binds at \varepsilon, and η is ε. Only once the step’s entropy climbs past H^\star does η lower its floor below \varepsilon and start admitting more tokens.

So η desmooths exactly as hard as ε when the model is confident enough to warrant it, and eases off precisely when a fixed floor would do damage. That makes η the safer default of the two.

# epsilon_filtering and eta_filtering live in generation.py beside the others
from generation import epsilon_filtering, eta_filtering, demonstrate_eta

# Two steps at the same epsilon: a confident one (below the crossover) and an
# uncertain one (above it). eta matches epsilon on the first, relaxes on the second.
out = demonstrate_eta(epsilon=0.02)
============================================================
EPSILON vs ETA SAMPLING (epsilon = 0.02)
============================================================

Crossover entropy H* = 0.5 ln(1/epsilon) = 1.956 nats
Below H* the two are identical; above it, eta relaxes the floor.

 peaked: H = 1.535 nats  |  epsilon kept  3, eta kept  3  (H < H* -> identical)
   flat: H = 4.159 nats  |  epsilon kept  1, eta kept 64  (H > H* -> eta relaxes)

On the peaked step the entropy sits below H^\star, so ε and η keep the same three tokens. On the flat step (64 near-uniform tokens, each below the floor) ε collapses to a single survivor while η, having dropped its floor, keeps all 64 — the plausible continuations a fixed floor would have thrown away as noise. Drag the two knobs and watch the floors move:

NoteKey Insight

A model smears a little probability onto tokens that should have none, just to keep the training loss finite. Truncation is desmoothing — cutting that film before you sample. ε-sampling cuts with a fixed floor; η-sampling cuts with a floor that lowers itself as the step’s entropy rises, so it stops desmoothing precisely when the model’s uncertainty is real rather than noise. Because \eta \le \varepsilon always, η never truncates more than ε — only less, and only where it should.

TipTry It!
  1. Turn concentration up. The distribution peaks, H drops below H^\star, the two floor lines snap together, and every dim (η-only) bar disappears — ε and η agree exactly when the model is confident.
  2. Turn it down. As the step flattens and H climbs past H^\star, the η line slides below the ε line and a band of dim bars appears: the plausible tail η rescues that a fixed floor would have cut.
  3. Sweep ε. A larger ε raises both floors (more aggressive desmoothing) and lowers the crossover H^\star = \tfrac12\ln\frac1\varepsilon, so η starts relaxing sooner. The paper uses tiny values (~3\times10^{-4} to 2\times10^{-3}) on a full 50k vocabulary; here ε is larger so the effect is visible on 24 tokens.

8. Mirostat: Closing the Loop on Perplexity

Every strategy so far is open-loop: you pick a cutoff — a k, a p, a floor — apply it to each step in isolation, and hope the text that comes out has the perplexity you wanted. You never actually measure the result and correct. Mirostat (Basu et al., ICLR 2021) is the one decoder in this chapter that does: it is a feedback controller that names the target perplexity directly and adjusts the cutoff on the fly to hit it.

The problem it solves is that a fixed truncation gives you the wrong perplexity in two opposite ways, and the paper names both. Set the cutoff too tight and you fall into the boredom trap — perplexity drifts down as the text lengthens, the model recycles a shrinking pool of safe tokens, and you get repetition. Set it too loose and you fall into the confusion trap — perplexity climbs, the tail leaks in, and coherence breaks down. The healthy zone is a target surprise in between, and it is different for every prompt and every position. No single k or p holds it.

NoteKey Insight

Think of a thermostat. You don’t set the furnace power (the cutoff k); you set the temperature you want (the target surprise τ) and let a loop adjust the furnace to hold it. Mirostat’s “furnace power” is a running surprise budget μ, and its thermostat is a one-line feedback rule.

The Math: a budget and a feedback loop

Work in bits. A token’s surprise is its information content

S(x) = -\log_2 p(x),

and the perplexity of a stretch of text is 2^{\overline{S}}, the exponential of the mean surprise — so controlling the average surprise to a target \tau controls perplexity to 2^{\tau} exactly. Mirostat carries a surprise budget \mu (initialized to \mu = 2\tau) and, at each step, runs four moves:

  1. Truncate to the tokens the budget can afford — those whose surprise is below it: \mathcal{A}_\mu = \{\, x : S(x) = -\log_2 p(x) < \mu \,\}. This is an adaptive top-k whose cutoff is a variable: a larger \mu admits more of the surprising tail, a smaller \mu clamps to the confident head.
  2. Sample a token X from that (renormalized) set.
  3. Measure how surprising the draw actually was, under the original distribution: S = -\log_2 p(X).
  4. Correct the budget by the error e = S - \tau: \mu \;\leftarrow\; \mu - \eta\,(S - \tau).

That last line is the whole method, and it is exactly one step of gradient descent on the surprise error. If the draw came out too surprising (S > \tau, so e > 0), the budget shrinks — next step truncates harder and pulls surprise back down. If it was too boring (e < 0), the budget grows and lets more of the tail through. The rate \eta is a learning rate: large \eta settles fast but jitters, small \eta is smooth but slow.

NoteKey Insight

The original Mirostat 1.0 got its truncation set by estimating the distribution’s Zipf exponent each step and solving for a top-k. Mirostat 2.0 — what we build, and what llama.cpp and every UI ship — throws that assumption away: it truncates directly on surprise, so it works on any distribution. The feedback update in step 4 is identical in both.

Code: the controller from scratch

The truncation is a stateless filter like the others — mirostat_truncate in generation.py keeps the tokens whose surprise is below the budget:

import torch
from generation import mirostat_truncate

# A peaked step over 8 tokens. A tight budget affords only the confident head;
# a generous budget affords more of the surprising tail.
logits = torch.tensor([[4.0, 2.0, 1.0, 0.5, 0.0, 0.0, 0.0, 0.0]])
for mu in (1.0, 2.0, 4.0):
    kept = int((mirostat_truncate(logits, mu) > float("-inf")).sum())
    print(f"budget mu = {mu:>3} bits  ->  {kept} token(s) affordable")
budget mu = 1.0 bits  ->  1 token(s) affordable
budget mu = 2.0 bits  ->  1 token(s) affordable
budget mu = 4.0 bits  ->  2 token(s) affordable

The state — the budget and its feedback update — lives in the MirostatSampler class. Each .sample() call runs the four moves and nudges μ:

from generation import MirostatSampler

sampler = MirostatSampler(tau=3.0, eta=0.1)   # target 3 bits (perplexity 8)
# A fixed, high-entropy source (Zipf over 200 tokens) stands in for a model step.
probs = 1.0 / torch.arange(1, 201, dtype=torch.float32)
source = (probs / probs.sum()).log().unsqueeze(0)

gen = torch.Generator().manual_seed(0)
print(f"start:            mu = {sampler.mu.item():.2f} bits  (= 2*tau)")
for i in range(100):
    _, rec = sampler.sample(source, generator=gen)
print(f"after 100 steps:  mu = {sampler.mu.item():.2f} bits  "
      f"(the loop settled the budget to hold the 3-bit target)")
start:            mu = 6.00 bits  (= 2*tau)
after 100 steps:  mu = 4.16 bits  (the loop settled the budget to hold the 3-bit target)

Run the full controller over a stationary source and watch the achieved surprise land on the target, even though the budget started at 2\tau:

from generation import demonstrate_mirostat

out = demonstrate_mirostat(tau=3.0, eta=0.1)
============================================================
MIROSTAT 2.0 FEEDBACK CONTROL (tau = 3.0 bits, eta = 0.1)
============================================================

Source entropy         : 6.85 bits (raw perplexity 115)
Initial budget  mu_0   : 6.00 bits  (= 2*tau)
Settled budget  mu     : 3.89 bits  (2nd-half mean)
Target surprise tau    : 3.00 bits  (perplexity 8.0)
Achieved surprise      : 3.01 bits  (2nd-half mean)  <- tracks tau

The source pours out ~6.8 bits of raw surprise per token (perplexity ~115), but Mirostat holds the generated stream at the 3-bit target — a perplexity of 8 — by settling its budget wherever it needs to. Change tau and the achieved surprise moves with it: it is a genuine dial on perplexity, not a proxy.

Interactive Exploration

Drive the loop. Each dot is one step’s observed surprise; the bright line is the running mean, the flat accent line is your target \tau, and the second curve is the budget \mu hunting for the value that holds the mean on target. The dashed line is what a fixed top-k (k = 40) would average — the perplexity you’d be stuck with, with no way to set it.

TipTry It!
  1. Drag \tau. The whole trajectory re-settles: the running mean (bright line) climbs or drops to meet the new target within a few dozen steps, and the budget \mu follows it there. This is the payoff — perplexity is a setting, not an outcome you discover afterward.
  2. Compare to the dotted top-k line. That fixed baseline sits wherever k = 40 happens to land it. To move it you’d have to guess a new k; Mirostat just moves \tau.
  3. Crank \eta up. The mean snaps to target faster but the budget jitters — a hot controller overshoots. Turn \eta down and the settling is smooth but slow. That trade-off (responsiveness vs. stability) is the universal signature of a feedback loop.

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:
  [27, 44, 27, 44, 27, 44, 44, 44, 45, 45]
  [27, 44, 27, 44, 27, 44, 44, 44, 45, 45]
  [27, 44, 27, 44, 27, 44, 44, 44, 45, 45]

Temperature=0.5:
  [44, 34, 14, 27, 27, 8, 24, 28, 9, 18]
  [10, 29, 33, 13, 17, 21, 46, 16, 23, 45]
  [0, 47, 13, 32, 49, 10, 28, 46, 44, 27]

Temperature=1.0:
  [3, 34, 14, 27, 27, 8, 24, 28, 9, 18]
  [10, 29, 33, 13, 30, 11, 46, 16, 23, 45]
  [0, 47, 13, 32, 49, 10, 28, 46, 44, 27]

Top-k=5:
  [45, 45, 44, 27, 27, 49, 18, 27, 45, 18]
  [14, 0, 27, 19, 14, 0, 44, 27, 0, 49]
  [0, 18, 45, 45, 18, 44, 44, 27, 44, 27]

Top-p=0.9:
  [3, 34, 14, 27, 27, 8, 24, 28, 9, 18]
  [10, 29, 33, 42, 17, 21, 46, 16, 23, 45]
  [0, 47, 13, 17, 49, 10, 28, 46, 44, 27]

Min-p=0.1:
  [3, 34, 14, 27, 27, 8, 24, 28, 9, 18]
  [10, 29, 33, 13, 30, 11, 46, 16, 23, 45]
  [0, 47, 13, 32, 49, 10, 28, 46, 44, 27]

Combined (T=0.7, k=20, p=0.9):
  [44, 9, 14, 27, 27, 21, 17, 27, 9, 18]
  [10, 0, 47, 17, 14, 21, 44, 16, 23, 45]
  [0, 47, 25, 17, 49, 10, 44, 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:
  [25, 25, 12, 27, 27, 47, 10, 46, 14, 23, 17, 14]
  [18, 10, 25, 21, 18, 23, 27, 23, 45, 0, 18, 27]

Balanced chat:
  [25, 25, 29, 27, 11, 47, 10, 46, 14, 23, 17, 14]
  [18, 10, 25, 21, 18, 23, 27, 23, 45, 0, 46, 27]

Creative writing:
  [25, 25, 29, 27, 11, 47, 10, 2, 14, 41, 17, 14]
  [18, 10, 25, 21, 18, 23, 27, 23, 45, 0, 39, 27]

Brainstorming:
  [37, 25, 29, 27, 11, 47, 10, 2, 14, 41, 17, 14]
  [18, 10, 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: [27, 44, 27, 44, 27, 44, 44, 44, 45, 45, 49, 49, 45, 49, 49, 45, 45, 45, 45, 49, 45, 49, 49, 49, 49, 49, 49, 49, 49, 49]
  Most common: [(49, 14), (45, 8), (44, 5)]

Generation with repetition penalty (1.3):
  Tokens: [27, 44, 45, 18, 49, 44, 44, 44, 44, 44, 27, 49, 49, 49, 49, 49, 49, 49, 49, 49, 45, 49, 49, 49, 49, 49, 49, 49, 49, 49]
  Most common: [(49, 19), (44, 6), (27, 2)]

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: [27, 44, 27, 44, 27, 44, 44, 44, 45, 45, 49, 49, 45, 49, 49, 45, 45, 45, 45, 49]

With EOS check: 20 new tokens generated
  Tokens: [27, 44, 27, 44, 27, 44, 44, 44, 45, 45, 49, 49, 45, 49, 49, 45, 45, 45, 45, 49]

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)

Contrastive Search: Determinism Without Degeneration

The beam-search warning leaves an uncomfortable gap. Maximizing probability (greedy, beam) degenerates into loops; the escape route the whole lesson took was to sample. But sampling buys diversity by giving up determinism — run it twice and you get two different answers. Is repetition really the price of a deterministic decoder, or a bug we can fix directly?

Contrastive search (Su et al., 2022) fixes it directly. It stays fully deterministic — argmax, no randomness — but adds a term that actively pushes back on repetition. It is the last decoder in this lesson, and the only one that is both deterministic and built for open-ended text.

Intuition: repetition is a geometry problem

Su et al. traced degeneration to the shape of a model’s representation space. As a Transformer trains, the hidden vectors it assigns to tokens drift into a narrow cone — they become anisotropic, all pointing roughly the same way. When representations are that crowded, the cosine similarity between the current state and a token it just produced is high, and the decoder is nudged to emit it again. The loop is a feedback effect of a collapsed geometry.

That diagnosis suggests the cure. At each step, look at the model’s hidden state h_v for each candidate token v, and compare it against the states of the tokens already generated. If a candidate’s representation is nearly identical to something in the context, it is probably a repeat — so penalize it, no matter how probable the model thinks it is. Keep the model’s confidence, but subtract a degeneration penalty. Two forces, one score.

The Math: confidence minus a degeneration penalty

From the model’s top-k most probable tokens V^{(k)}, pick the one maximizing

x_t \;=\; \operatorname*{arg\,max}_{v \in V^{(k)}} \Big\{\, (1 - \alpha)\, \underbrace{p_\theta(v \mid x_{<t})}_{\text{model confidence}} \;-\; \alpha \underbrace{\max_{1 \le j \le t-1} s\big(h_v,\, h_{x_j}\big)}_{\text{degeneration penalty}} \Big\}.

The pieces:

  • Model confidence p_\theta(v \mid x_{<t}) — the usual next-token probability. Left alone, its argmax over V^{(k)} is just greedy.
  • Degeneration penalty \max_j s(h_v, h_{x_j}) — the largest cosine similarity s between the candidate’s representation h_v and any token already in the sequence. Near 1 means “you have seen this before.”
  • \alpha \in [0, 1] balances them. At \alpha = 0 the penalty vanishes and contrastive search is greedy over the top-k (i.e. plain greedy); as \alpha \to 1 it chases novelty. The paper uses k = 8, \alpha = 0.6.

Two gates matter. The top-k restriction keeps the search honest: a wildly novel but implausible token is never even considered, so coherence is preserved. And h_v is the representation the model produces after appending v — the same forward pass, one token longer — so the penalty reads the model’s own geometry, not a bolt-on heuristic. We build it from scratch in contrastive.py.

Code: contrastive search from scratch

We reuse the scripted-model trick from beam search, but the toy now returns a representation for every token as well as a distribution. In ScriptedReprModel.repetition_trap() the representations are the one-hot basis vectors, so a candidate’s degeneration penalty is exactly 1.0 when that token already appears in the context and 0.0 otherwise — an exact repeat detector that makes every number below checkable by hand. The distribution is the same everywhere: “the” (0.5) beats “cat” (0.3) beats “sat” (0.2), so greedy loops forever:

from contrastive import (
    contrastive_search,
    degeneration_penalty,
    contrastive_score,
    mean_context_similarity,
    ScriptedReprModel,
    demonstrate_contrastive_search,
)

trap = ScriptedReprModel.repetition_trap()   # "the" is always the argmax
prompt = torch.tensor([[0]])                  # BOS
names = {0: "BOS", 1: "the", 2: "cat", 3: "sat"}

greedy = contrastive_search(trap, prompt, top_k=3, alpha=0.0, max_new_tokens=5)
contra = contrastive_search(trap, prompt, top_k=3, alpha=0.6, max_new_tokens=5)
print("greedy      (α=0.0):", [names[t] for t in greedy[0, 1:].tolist()])
print("contrastive (α=0.6):", [names[t] for t in contra[0, 1:].tolist()])
greedy      (α=0.0): ['the', 'the', 'the', 'the', 'the']
contrastive (α=0.6): ['the', 'cat', 'sat', 'the', 'the']

Same model, same start: greedy collapses to a single word, contrastive search walks through the vocabulary. The mechanism is the penalty flipping on once a token is in the context — and, exactly as for beam search’s length penalty, there is a clean threshold where the winner changes. At step 2 the context already holds “the”, so its penalty is 1.0 while fresh “cat” pays 0. Setting their scores equal, (1-\alpha)\,0.5 - \alpha = (1-\alpha)\,0.3, gives \alpha = 1/6:

for a in (0.16, 0.17):
    seq = contrastive_search(trap, prompt, top_k=3, alpha=a, max_new_tokens=2)
    print(f"α={a}: second token is '{names[seq[0, 2].item()]}'")
α=0.16: second token is 'the'
α=0.17: second token is 'cat'

Below 1/6 the model’s confidence in “the” still wins; above it the penalty does. The contrastive.py implementation is model-agnostic — it takes any model(tokens) -> (logits, hidden). On a real GPTModel (hidden states from return_hidden_states=True), the \alpha = 0 case must reproduce greedy exactly, because the penalty is multiplied away:

from m06_transformer.transformer import GPTModel

torch.manual_seed(0)
gpt = GPTModel(vocab_size=100, embed_dim=32, num_heads=2, num_layers=2,
               max_seq_len=64, dropout=0.0)

def gpt_with_hidden(tokens):                       # adapt to (logits, hidden)
    logits, hs = gpt(tokens, return_hidden_states=True)
    return logits, hs[-1]                          # last layer's hidden state

p = torch.tensor([[2, 7, 11]])
cs0 = contrastive_search(gpt_with_hidden, p, top_k=5, alpha=0.0, max_new_tokens=6)
gd = generate_greedy(gpt, p, max_new_tokens=6)
print("α=0 contrastive:", cs0[0].tolist())
print("greedy:         ", gd[0].tolist())
print("identical:      ", torch.equal(cs0, gd))
α=0 contrastive: [2, 7, 11, 8, 62, 12, 62, 12, 8]
greedy:          [2, 7, 11, 8, 62, 12, 62, 12, 8]
identical:       True

Contrastive search is not a different family of decoder — at \alpha = 0 it is greedy. The penalty is a dial you turn up.

Interactive: the α dial — confidence vs. novelty

Below is the exact decision contrastive search faces at step 2 of the trap, with “the” already generated. Each candidate contributes a confidence bar (1-\alpha)\,p (it earned) and a penalty bar \alpha \cdot \text{penalty} (it owes); the net score is the difference, and the highlighted candidate is the one selected. Drag \alpha and watch the winner flip from the repeated “the” to the fresh “cat” as it crosses 1/6 \approx 0.17.

TipTry This
  1. Start at \alpha = 0. Only the green confidence bars show; “the” (highest p) wins. This is greedy.
  2. Raise \alpha past \approx 0.17. “the” grows a full red penalty bar (it repeats a context token, penalty 1.0) while “cat” and “sat” stay clean. The winner flips to “cat” — the loop is broken deterministically.
  3. Push \alpha \to 1. Confidence vanishes entirely; the decoder now picks purely by novelty, ignoring what the model thinks is likely. This is why \alpha is kept moderate (\approx 0.6): too high and coherence goes with it.

Interactive: watch the full decode stay isotropic

The dial showed one step; here is the whole run at \alpha = 0.6. Step through it and watch each chosen token — the sequence spreads across the vocabulary instead of collapsing, and the mean self-similarity of the generated representations stays low (greedy’s would be a flat 1.0).

NoteKey Insight

Contrastive search and repetition penalty attack the same symptom from opposite ends. Repetition penalty (earlier in this lesson) edits the logits by token identity — “you emitted this id, so lower it.” Contrastive search edits the ranking by representation geometry — “this candidate looks like something you emitted, so lower it.” The geometric view catches near-repeats and paraphrases that a literal id-match misses, and it needs no tunable per-token bookkeeping — just \alpha, k, and the model’s own hidden states.

The paper pairs contrastive search with a light contrastive training objective, SimCTG, that spreads the token representations apart (raising their pairwise distance, margin \rho = 0.5) so the geometry the decoder relies on is well-behaved from the start. The decoder works on off-the-shelf models too — the training just makes the isotropy the penalty exploits more pronounced.

Warningα is a coherence knob, not a “more is better” dial

Turning \alpha toward 1 does not make text “more diverse and just as good” — past a point the decoder ignores the model’s probabilities and picks near-random in-vocabulary tokens, and fluency collapses. Keep it moderate (0.5–0.7) so confidence still leads and the penalty only breaks ties toward novelty.

WarningContrastive search needs a well-shaped representation space

The penalty is only as meaningful as the geometry it reads. On a model whose representations are pathologically anisotropic (everything cosine-similar to everything), even novel tokens carry a high penalty and the signal washes out — which is exactly why Su et al. pair the decoder with SimCTG training. If contrastive search underwhelms, the representations, not the decoder, are usually the problem.

Constrained Decoding: Guaranteeing Valid Output

Nearly every strategy so far — greedy, temperature, top-k, top-p, min-p, repetition penalty — is the same move: reshape the logits, then sample (beam search and contrastive search were the exceptions: beam searches over whole sequences, and contrastive search re-ranks the top-k by representation geometry). Constrained decoding returns to the per-step pattern — 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.

Compiling a Regex: NFA → DFA

The template toolkit built those machines by hand, and its own docstring admits the catch: it works “as long as a variable-length field and the text following it start with different characters — no NFA subset construction required.” That is a shortcut, not the general road. Write a genuinely ambiguous pattern — a starred alternation like a(b|c)*d, where after an a the next character could open another b/c loop or be the closing d — and the left-to-right trick breaks: the machine would need to be in two places at once.

The real compiler that every structured-generation library runs (Outlines, above) turns any regular expression into a deterministic FSM in three moves — and we build all three from scratch in regex_fsm.py:

regex  ──parse──▶  AST  ──Thompson──▶  ε-NFA  ──subset──▶  DFA (an FSM)

Two classic theorems do the work. Thompson’s construction (1968) maps every regex operator to a tiny automaton fragment glued with ε-transitions (edges taken without consuming a character), so the NFA has O(len(regex)) states. Subset construction (Rabin & Scott, 1959) then removes the nondeterminism: the set of NFA states the machine could be in is itself one DFA state — proving an NFA and a DFA recognize exactly the same language. The payoff is immediate: the output is a plain FSM, so it drops straight into build_token_index and constrained_decode with zero new machinery.

from regex_fsm import compile_regex

# A starred alternation the left-to-right template DSL cannot express:
fsm = compile_regex("a(b|c)*d")
print(fsm.accepts("abcbcd"))   # True: any run of b/c, then d
print(fsm.accepts("abcbc"))    # False: must end in d

# The flagship — a date, compiled straight from a regex, no template needed:
date = compile_regex(r"\d{4}-\d{2}-\d{2}")
print(date.accepts("2024-01-31"))  # True
print(date.accepts("2024-1-31"))   # False: month needs two digits
True
False
True
False

That second machine accepts the exact same language as the hand-built date_fsm() from earlier — the compiler simply derives it for you. Because the compiled object is an ordinary FSM, constraining a model to a regex is now a one-liner: constrained_decode(logit_fn, compile_regex(r"\d{4}-\d{2}-\d{2}"), vocab, ...).

Step 1 — Thompson’s ε-NFA. Below is the NFA for a(b|c)*d. Solid edges consume a character; dashed ε-edges are the free glue that stitches the fragments together — the | forks into two branches, and the * loops back for another repetition or skips ahead. The nondeterminism lives entirely in those ε-forks.

Step 2 — subset construction. Now collapse the NFA into a DFA. Each DFA state is a set of NFA states — the machine’s “superposition.” We start from the ε-closure of the NFA start, then for each character compute the closure of where it leads; a set we have not seen becomes a new DFA state. Step the worklist below and watch the powerset get discovered one state at a time.

Step 3 — minimize, then the state-count story. Subset construction can leave redundant DFA states (two that no future string can tell apart). Moore’s partition refinement merges them into the unique smallest DFA (Myhill–Nerode). The bars below trace the whole pipeline’s state count for two patterns: nondeterminism is free (the NFA is smallest), determinizing can cost states, and minimization wins some back.

NoteKey Insight

A regular grammar is exactly what a finite automaton can recognize — and a regex is a regular grammar. Thompson’s construction and subset construction turn any regex into a DFA mechanically, and because a DFA has finitely many states you can precompute its allowed-token index once (the Outlines road) and pay only a dict lookup per step. The template DSL was a hand-tuned special case; this is the general engine underneath it.

TipTry This
  1. Force the nondeterminism. Step the subset-construction widget for a(b|c)*d: after the a, a single DFA state stands for both “inside the loop” and “ready for d”. That merged set is why one deterministic pass suffices.
  2. Watch a machine shrink. In a {python} cell, compile compile_regex("aa|ab", minimize_dfa=False) and compile_regex("aa|ab") and compare len(fsm.transitions) — minimization merges the two post-a tails into one.
  3. Constrain to your own regex. Swap the date grammar in the decode cell for compile_regex(r"(true|false)") or compile_regex(r"[a-z]+@[a-z]+") and confirm every sampled run matches — the model can only spell strings the regex accepts.

Beyond Regular: Nested JSON Needs a Stack

The FSM above handles a fixed schema. But real JSON nests to any depth — objects inside arrays inside objects — and here a finite-state machine hits a hard wall. To close a } correctly the machine must remember how many braces are still open, and “how many” can be any number. A machine with a fixed number of states cannot count without bound.

This is not a limitation of our FSM; it is a theorem. The language of n opening brackets followed by n closing brackets,

L = \{\, [^{\,n}\,]^{\,n} : n \ge 0 \,\} = \{\, \texttt{[]},\ \texttt{[[]]},\ \texttt{[[[]]]},\ \dots \,\},

is the textbook non-regular language. Any DFA with k states, run on [^{\,k+1}, must revisit a state (pigeonhole), so it can no longer tell depth i from depth j — and it will accept some unbalanced string. dfa_cannot_count builds exactly such a counter and catches it in the act:

from pushdown import (
    json_pda_start, pda_step, pda_accepts, pda_can_end,
    pda_allowed_token_ids, pda_constrained_decode,
    stack_trace, dfa_cannot_count, demonstrate_pushdown,
)

# Our fixed JSON *schema* FSM has no state to remember an open brace:
flat = json_answer_fsm()
print("FSM accepts a nested value?",
      flat.accepts('{"answer": {"n": 1}, "correct": true}'))   # -> False

# A k-state bracket counter is provably fooled by depth k+1:
demo = dfa_cannot_count(3)
print(f"\nA {demo['states']}-state bracket DFA:")
print(f"  balanced   {demo['balanced']!r}: DFA accepts={demo['dfa_balanced']}, really valid=True")
print(f"  UNbalanced {demo['false_accept']!r}: DFA accepts={demo['dfa_false']}  <- WRONG")
print(f"  the same unbalanced string, PDA verdict: {demo['pda_false']}  <- correct")
FSM accepts a nested value? False

A 4-state bracket DFA:
  balanced   '[[[]]]': DFA accepts=True, really valid=True
  UNbalanced '[[[[]]]': DFA accepts=True  <- WRONG
  the same unbalanced string, PDA verdict: False  <- correct

Add one thing — a stack — and the machine can count. A finite control plus an unbounded stack is a pushdown automaton (PDA), and PDAs recognize exactly the context-free languages, JSON among them. The rule is simple: push when you open a container, pop-and-match when you close one.

config = (control mode, stack of open containers)
    '{'  ->  push OBJ        '}'  ->  pop, only if the top is OBJ
    '['  ->  push ARR        ']'  ->  pop, only if the top is ARR
accept  <=>  input ended, machine alive, and the stack is empty

This is built from scratch in pushdown.py. pda_step(config, char) is the whole JSON grammar as a transition function; the stack’s height is the current nesting depth:

cfg = json_pda_start()
for ch in '{"a": [1]}':
    cfg = pda_step(cfg, ch)          # push on '{' and '[', pop on ']' and '}'
print("stack after a full document:", cfg.stack, "| complete?", pda_can_end(cfg))

# The stack has no ceiling, so nesting has no ceiling:
print("accepts 200-deep nesting:", pda_accepts("[" * 200 + "]" * 200))
print("rejects a bracket mismatch:", pda_accepts("[1, 2}"))
print("rejects a trailing comma: ", pda_accepts('{"a": 1,}'))
stack after a full document: () | complete? True
accepts 200-deep nesting: True
rejects a bracket mismatch: False
rejects a trailing comma:  False
NoteKey Insight

The FSM and the PDA mask logits with the same -inf trick; the only difference is what “the state” is. For a regular grammar the state is finite, so Outlines can precompute an allowed-token index once (build_token_index). For a context-free grammar the configuration includes an unbounded stack — there are infinitely many — so there is no table to precompute; the decoder walks the machine per step (an incremental parse). Regular ⇒ index; context-free ⇒ stack. That single line is the whole difference between JSON-schema mode and full GBNF-grammar mode.

The token layer is otherwise identical: a token is legal iff feeding its characters keeps the PDA alive, and pda_constrained_decode masks the rest. So a model — random logits and all — can only ever emit valid, arbitrarily-nested JSON:

import json, torch

# A tiny JSON token vocabulary (characters + a couple of words).
pv = {i: s for i, s in enumerate(
    ['{', '}', '[', ']', '"', ':', ',', 'k', 'true', 'null'] + list("0123456789"))}
peos = len(pv)

# Bias the "model" toward opening a container for its first few steps, so we can
# watch the stack work at depth. The constraint does the rest.
open_id = next(t for t, s in pv.items() if s == "[")
calls = {"n": 0}
def biased_logits(_ids):
    logits = torch.zeros(peos + 1)
    calls["n"] += 1
    if calls["n"] <= 6:
        logits[open_id] = 100.0          # prefer '[' early
    return logits

result = demonstrate_pushdown(biased_logits, pv, eos_token_id=peos, max_tokens=40)
print("generated:", result["text"])                 # -> [[[[[[{}]]]]]]
print("nesting depth reached:", result["max_depth"]) # deeper than any fixed FSM
print("json.loads parses it:", result["parses"])     # always True
generated: [[[[[[{}]]]]]]
nesting depth reached: 7
json.loads parses it: True

No finite template_fsm could produce that six-deep string — it would need a distinct state for every depth. The PDA needs one stack.

# Bridge two traces to the visualizations below.
_pda_trace = stack_trace('{"user": {"tags": ["a", "b"]}}')
_count_demo = [dfa_cannot_count(d) for d in [1, 2, 3, 4, 5]]
ojs_define(pdTrace = _pda_trace, pdCount = _count_demo)

Interactive: Why a Finite Machine Can’t Count

Drag the depth. A k-state bracket counter accepts the balanced string it was sized for — but saturates one level too shallow, so it also accepts an unbalanced string one opener too long. The PDA, with an unbounded stack, is never fooled.

Interactive: Watch the Stack

Step character by character through {"user": {"tags": ["a", "b"]}}. Each { or [ pushes a block onto the stack; each } or ] pops one. The stack height is the nesting depth — the piece of memory no finite machine has.

TipTry This
  1. Watch the stack breathe. Step through the trace and note the depth hit 3 at ["a" (object → array → … ), then unwind to 0. A completed document always ends with an empty stack.
  2. Go arbitrarily deep. In the decode cell, raise the <= 6 bias cutoff to <= 20; the output nests 20-plus deep and still parses. No FSM could enumerate those states.
  3. Break the balance. Call pda_accepts("[" * 5 + "]" * 4) — one closer short — and confirm it is rejected. The stack never empties, so the document is never “complete.”

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) — which we build from scratch above in Beyond Regular: Nested JSON Needs a Stack. The token-masking idea is identical; only the “what’s legal next” bookkeeping grows a stack, and the precomputed FSM index gives way to an incremental parse.

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.19e-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|: 5.96e-08  (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: [8, 8, 8, 30]

Temperature = 0.5:
  5/5 unique sequences
  Sample: [11, 47, 10, 30]

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

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

Temperature = 5.0:
  5/5 unique sequences
  Sample: [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 2b: Make Typical Sampling Drop the Argmax

Top-k and top-p can never exclude the most likely token; typical sampling can. Find a typical_p that does, then confirm top-p keeps the same token:

from generation import typical_filtering

# A high-entropy step: token 0 holds ~42% of the mass but a wide band of
# equally-plausible tokens keeps the entropy high.
logits = torch.tensor([[2.85] + [0.0] * 24])
probs = F.softmax(logits, dim=-1)[0]
argmax = int(probs.argmax())

logp = F.log_softmax(logits, dim=-1)[0]
H = -(probs * logp).sum().item()
print(f"Entropy H(p) = {H:.3f} nats,  argmax p = {probs[argmax]:.3f}, "
      f"surprisal = {-logp[argmax].item():.3f}")

for tau in (0.3, 0.5, 0.9):
    typ = typical_filtering(logits.clone(), tau) > float("-inf")
    top = top_p_filtering(logits.clone(), tau) > float("-inf")
    print(f"  typical_p={tau}: typical keeps argmax = {bool(typ[0, argmax])}"
          f"  |  top-p keeps argmax = {bool(top[0, argmax])}")

# TODO: which typical_p values drop the argmax? Why does top-p never?
# TRY: raise the peak logit (2.85 -> 5.0). At what point is the peak SO dominant
# that it becomes 'typical' again (its surprisal returns near H)?
Entropy H(p) = 2.527 nats,  argmax p = 0.419, surprisal = 0.871
  typical_p=0.3: typical keeps argmax = False  |  top-p keeps argmax = True
  typical_p=0.5: typical keeps argmax = False  |  top-p keeps argmax = True
  typical_p=0.9: typical keeps argmax = True  |  top-p keeps argmax = True

Exercise 2c: Find η’s Crossover Entropy

η-sampling equals ε-sampling below the crossover entropy H^\star = \tfrac12\ln\frac1\varepsilon and only relaxes above it. Locate that boundary empirically by flattening a distribution until η pulls ahead of ε:

import math
from generation import epsilon_filtering, eta_filtering

eps = 0.02
crossover = 0.5 * math.log(1 / eps)
print(f"Predicted crossover H* = {crossover:.3f} nats")

# A uniform distribution over V tokens has entropy ln(V) and per-token prob 1/V.
# Grow V and watch eta start keeping more than epsilon exactly past H*.
for V in (4, 8, 16, 32, 64, 128):
    logits = torch.zeros(1, V)
    H = math.log(V)
    eps_kept = int((epsilon_filtering(logits.clone(), eps) > float("-inf")).sum())
    eta_kept = int((eta_filtering(logits.clone(), eps) > float("-inf")).sum())
    flag = "eta relaxes" if eta_kept > eps_kept else "identical"
    print(f"  V={V:>3}  H={H:.2f}  ε kept {eps_kept:>3}, η kept {eta_kept:>3}  ({flag})")

# TODO: at which V does eta first pull ahead? Compare ln(V) to H*.
# TRY: raise eps to 0.05. H* drops — does eta start relaxing at a smaller V?
Predicted crossover H* = 1.956 nats
  V=  4  H=1.39  ε kept   4, η kept   4  (identical)
  V=  8  H=2.08  ε kept   8, η kept   8  (identical)
  V= 16  H=2.77  ε kept  16, η kept  16  (identical)
  V= 32  H=3.47  ε kept  32, η kept  32  (identical)
  V= 64  H=4.16  ε kept   1, η kept  64  (eta relaxes)
  V=128  H=4.85  ε kept   1, η kept 128  (eta relaxes)

Exercise 2d: Watch Mirostat Settle on a Target

Mirostat’s promise is that the observed surprise settles at whatever τ you ask for. Verify it, and find the source’s floor — the target it can’t go below.

from generation import demonstrate_mirostat

# The loop should drive the mean observed surprise to tau from mu_0 = 2*tau.
for tau in (3.0, 4.0, 5.0):
    out = demonstrate_mirostat(tau=tau, eta=0.1, verbose=False)
    err = out["mean_surprise_second_half"] - tau
    print(f"tau = {tau}:  achieved {out['mean_surprise_second_half']:.2f} bits  "
          f"(error {err:+.2f})")

# TODO: lower tau toward 2.0. The Zipf source's least-surprising token has
# surprise -log2(p_0) ~= 2.8 bits, so the mean CANNOT drop below it — the
# controller saturates. At which tau does 'achieved' stop tracking the target?
# TRY: raise eta to 0.4. Does it settle faster, or does the surprise trace get noisier?
# (Inspect out["surprise_trace"] to see the per-step draws.)
tau = 3.0:  achieved 3.01 bits  (error +0.01)
tau = 4.0:  achieved 4.02 bits  (error +0.02)
tau = 5.0:  achieved 5.02 bits  (error +0.02)

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: 1/36 (3%)
  First 15: [15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15]

Penalty = 1.1:
  Unique tokens: 2/36 (6%)
  First 15: [15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15]

Penalty = 1.3:
  Unique tokens: 4/36 (11%)
  First 15: [15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15]

Penalty = 1.5:
  Unique tokens: 4/36 (11%)
  First 15: [8, 49, 45, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8]

Penalty = 2.0:
  Unique tokens: 7/36 (19%)
  First 15: [45, 49, 25, 8, 8, 42, 43, 8, 8, 8, 48, 8, 8, 8, 8]

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 (51 tokens):
[0, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 5, 5]

Most common tokens:
  Token 15: 48 times (94.1%)
  Token 5: 2 times (3.9%)
  Token 0: 1 times (2.0%)

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: int | None = None, top_p: float | None = None, min_p: float | None = None, typical_p: float | None = None, epsilon: float | None = None, eta: float | None = None, mirostat_tau: float | None = None, mirostat_eta: float = 0.1, do_sample: bool = True, eos_token_id: int | None = 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)

Exercise 5: Compile and Constrain to Your Own Regex

Use the from-scratch compiler in regex_fsm.py to turn a regex into a constraint FSM, confirm it agrees with Python’s re, and watch subset construction shrink the machine.

from regex_fsm import compile_regex, subset_construction, regex_to_nfa, minimize, compare_with_re

# 1. Compile a pattern the template DSL can't express (a starred alternation).
fsm = compile_regex("a(b|c)*d")
print("abcbcd:", fsm.accepts("abcbcd"), "| abcbc:", fsm.accepts("abcbc"))

# 2. Cross-check the compiled DFA against the standard library.
report = compare_with_re("a(b|c)*d", ["ad", "abcd", "abcbc", "", "acd"])
print("agrees with re.fullmatch:", report["agree"])

# 3. See minimization merge redundant states.
raw, _ = subset_construction(regex_to_nfa("aa|ab"))
small = minimize(raw)
print(f"aa|ab states: {len(raw.transitions)} → {len(small.transitions)} after minimization")

# Your turn: write a regex for a US phone number `\d{3}-\d{3}-\d{4}` and
# confirm compile_regex(...).accepts("555-867-5309") is True.
abcbcd: True | abcbc: False
agrees with re.fullmatch: True
aa|ab states: 4 → 3 after minimization

Exercise 6: Beam Width and Length

Explore how the beam width and the length penalty change what beam search returns, using the scripted models from beam.py.

import math
from beam import beam_search, sequence_logprob, ScriptedModel

# 1. Sweep the beam width on the garden-path model. Where does the myopic
#    greedy answer (b=1) give way to the better B -> JACKPOT sequence?
garden = ScriptedModel.garden_path()
p = torch.tensor([[0]])
for b in (1, 2, 3):
    seq, _ = beam_search(garden, p, beam_width=b, max_new_tokens=2)
    lp = sequence_logprob(garden, seq, prompt_len=1)
    print(f"b={b}: {seq[0, 1:].tolist()}  p={math.exp(lp):.3f}")

# Your turn:
# 2. On ScriptedModel.length_corpus(), find the smallest length_alpha at which
#    beam_search stops returning the short sequence and returns the long one.
# 3. (Harder) beam_search here handles a single sequence (batch=1). Sketch what
#    would change to run B prompts at once: each prompt needs its own beam, and
#    the top-b prune must stay *within* a prompt, never mixing beams across
#    prompts. This bookkeeping is why production beam search looks so different
#    from the loop above.
b=1: [1, 3]  p=0.300
b=2: [2, 4]  p=0.360
b=3: [2, 4]  p=0.360

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. Typical sampling: Ranks tokens by information, not probability — keeps the smallest set whose surprisal −log p is closest to the distribution’s entropy H(p). The only strategy that can drop the most-likely token (too predictable) as readily as the deep tail (too surprising)
  8. ε- and η-sampling: Truncation seen as desmoothing — cutting the film of probability a model smears onto tokens that should have none. ε-sampling keeps p > ε (a fixed absolute floor); η-sampling keeps p > min(ε, √ε·e^(−H)) (an entropy-adaptive floor that lowers itself past the crossover H* = ½ln(1/ε)), so A_η ⊇ A_ε always and the two agree whenever the model is confident
  9. Mirostat: The one closed-loop decoder — set a target surprise τ (bits; perplexity 2^τ) and a feedback loop holds the generated stream there. Each step truncates to a running budget {x : −log₂ p(x) < μ}, samples, measures the draw’s surprise S, and updates μ ← μ − η(S − τ). Directly targets perplexity where every other filter only sets a cutoff and hopes; avoids the boredom trap (perplexity collapses) and the confusion trap (perplexity explodes)
  10. Repetition penalty: Reduces probability of previously-generated tokens to prevent loops
  11. Stop conditions: Use EOS tokens and max length to control when generation ends
  12. Combine strategies: Temperature + top-p (or min-p) + repetition penalty is common in practice
  13. Beam search: A bounded breadth-first search that keeps the b best partial sequences by cumulative log-probability; b=1 is greedy, and length normalization (score / |y|^α) undoes its bias toward short sequences. Right for search problems (translation), wrong for open-ended text (it degenerates)
  14. Contrastive search: A deterministic decoder for open-ended text — from the top-k, pick argmax (1−α)·p − α·max cos(h_v, h_context), trading model confidence against a degeneration penalty (max cosine similarity to the context). α=0 (or k=1) is greedy; the penalty breaks loops by representation geometry, so repetition is not the price of determinism
  15. Constrained decoding: Mask logits to a grammar’s legal tokens - output is valid by construction, not by re-checking; a regular grammar precomputes an allowed-token index (Outlines)
  16. Regex → FSM compilation: Any regex compiles mechanically to a DFA — Thompson’s construction builds an ε-NFA, subset construction removes the nondeterminism — and the result is a drop-in FSM, so compile_regex(...) constrains a model with zero new decoding machinery
  17. Pushdown constrained decoding: Nested JSON is context-free, not regular - a finite machine cannot count brackets, so add a stack; the same mask, an unbounded configuration, arbitrary nesting (the idea behind GBNF grammars)
  18. 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
Bland, looping beam-search text Beam search on open-ended generation Use sampling for open-ended tasks; reserve beam search for search problems
Beam search always stops short Raw scores favor shorter sequences Add a length penalty (length_alpha / GNMT (5+|y|)^α)
Want deterministic non-repetitive text Greedy loops; sampling isn’t reproducible Contrastive search: keep argmax, add a degeneration penalty (alpha≈0.6, top_k≈8)
Contrastive search too incoherent alpha pushed too high Lower alpha to 0.5–0.7 so confidence still leads and the penalty only breaks ties
Incoherent nonsense Temperature too high Lower temperature, use top-p or min-p
Output feels bland / too “safe” Probability-ordered filters always keep the over-confident head Try typical sampling (typical_p≈0.2–0.95): it can drop the too-predictable token
A fixed probability floor kills good text on uncertain steps ε-sampling’s absolute floor over-truncates high-entropy distributions Use η-sampling (eta≈6e-4–2e-3): the floor lowers itself past H* = ½ln(1/ε), so it only desmooths where the tail is genuinely noise
Perplexity drifts over a long generation (bland then rambling, or vice-versa) A fixed cutoff can’t hold surprise steady as context grows Use Mirostat (mirostat_tau≈3, mirostat_eta≈0.1): the feedback loop re-targets the cutoff every step to hold perplexity at 2^τ
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):

  • Diverse / batched beam search: run many prompts at once, or force the beams apart so they don’t collapse to near-duplicates (Vijayakumar et al., 2016)
  • 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: