Perplexity from scratch

Perplexity from scratch#

Track 06 - Evaluation · Notebook 01 · Runtime: ~30s on CPU

Prerequisites: basic probability, log and exp, softmax, light PyTorch.

Compute perplexity three ways on a tiny model trained inline, so no piece of the calculation is mysterious. The math, the implementation, and the two most common evaluation bugs (non-overlapping windows, mismatched tokenisers) are all covered.

A quick refresher - skip if you’re comfortable#

Language models are probability distributions over sequences#

A language model assigns a probability \(P(x_1, x_2, \dots, x_n)\) to every sequence of tokens. Using the chain rule of probability,

\[P(x_1, \dots, x_n) = \prod_{t=1}^{n} P(x_t \mid x_{<t})\]

so the model really only has to learn the conditional distribution \(P(x_t \mid x_{<t})\) - one step-ahead predictions. In a transformer that distribution is the softmax over the output logits at step \(t\).

Cross-entropy measures how well the model fits the data#

The per-token cross-entropy loss you minimise during training is

\[\mathcal{L}(x) = -\frac{1}{n} \sum_{t=1}^{n} \log P_\text{model}(x_t \mid x_{<t})\]

If the model predicts each token perfectly (\(P = 1\)), loss is \(0\). If it guesses uniformly among \(V\) tokens (\(P = 1/V\)), loss is \(\log V\).

Perplexity is just `exp(loss)`#

Perplexity is

\[\text{PPL}(x) = \exp\!\left(\mathcal{L}(x)\right) = \left(\prod_t P(x_t \mid x_{<t})\right)^{-1/n}\]

the geometric mean of \(1/P\). Two natural interpretations:

Uncertainty bound. A model with perplexity 20 behaves, on average, like a process that picks uniformly among 20 equally plausible alternatives at each step. Smaller is better.
Compression rate. \(\log_2 \text{PPL}\) is the number of bits the model needs per token (Shannon source-coding theorem). That’s why “bits-per-character” language-modelling benchmarks exist.

Three consequences you should internalise:

PPL is always \(\geq 1\).
A uniform model over \(V\) symbols has PPL \(= V\) (upper bound on any reasonable model evaluated on \(V\)-symbol text).
Comparing PPL between two models only makes sense if they share a tokeniser. Don’t compare byte-level PPL against BPE-level PPL and expect the answer to mean anything.

Step 1 - Perplexity from first principles#

Before we train any model, let’s check the two limit cases by hand: a uniform distribution and a delta (fully certain) distribution.

from llm_systems_cookbook.nb import bootstrap

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

s = bootstrap("06_eval_01_perplexity_from_scratch")

# Limit case A: uniform distribution over V symbols.
# Expected: loss = log V, PPL = V exactly.
V = 29  # we'll pick 29 because that's the alphabet we use later.

uniform_probs = torch.full((V,), 1.0 / V)
uniform_nll = -torch.log(uniform_probs).mean().item()
uniform_ppl = math.exp(uniform_nll)
print(f"uniform:  NLL={uniform_nll:.4f}   PPL={uniform_ppl:.4f}   (log V = {math.log(V):.4f})")

# Limit case B: delta distribution - one symbol with probability 1.
# Expected: NLL = 0, PPL = 1.
delta_probs = torch.zeros(V)
delta_probs[0] = 1.0
# log(0) is -inf, so we only evaluate the assigned symbol.
delta_nll = -math.log(delta_probs[0].item())
delta_ppl = math.exp(delta_nll)
print(f"delta:    NLL={delta_nll:.4f}   PPL={delta_ppl:.4f}")

s.assert_close("uniform_ppl_equals_vocab_size", actual=uniform_ppl, expected=float(V), rtol=1e-6)
s.assert_close("delta_ppl_equals_one",          actual=delta_ppl,   expected=1.0,         rtol=1e-6)

Checkpoint. You just proved that perplexity is bounded below by 1 and above (for reasonable models) by the vocabulary size. Any perplexity you ever read is a number between those two, usually closer to 1 than to \(V\).

Now let’s train a small language model so we can compute a real perplexity on real text.

Step 2 - A minimal character-level language model#

We’ll use characters instead of BPE tokens so everything stays tiny and legible. The model has three parts:

Embedding: char_id → vector. One nn.Embedding lookup.
Causal transformer block: a single layer (self-attention + MLP) that mixes context tokens into each position without peeking at the future.
Head: vector → logits over vocab. One linear projection.

That’s literally it. A real LLM just stacks 30-100 of these blocks with a larger vocabulary and context - the loss function and perplexity definition don’t change at all.

The text we train on is repetitive on purpose: we want the model to memorise it so perplexity is low enough to be interesting, but not so low that sliding-window vs non-overlapping evaluation differences vanish.

TEXT = (
    "the quick brown fox jumps over the lazy dog. "
    "perplexity is defined as the exponential of the cross-entropy loss. "
    "a language model assigns a probability distribution over the next token. "
    "sliding-window evaluation reduces bias from hard document boundaries. "
    "non-overlapping windows give a pessimistic estimate because early tokens in each window see little context. "
    "tokens sampled with a fair coin have perplexity equal to the alphabet size. "
) * 12

CHARS = sorted(set(TEXT))
STOI = {c: i for i, c in enumerate(CHARS)}
ITOS = {i: c for c, i in STOI.items()}
VOCAB = len(CHARS)

IDS = torch.tensor([STOI[c] for c in TEXT], dtype=torch.long)
print(f"corpus length = {len(TEXT)} chars, vocab = {VOCAB}")
print(f"first 60 chars:  {TEXT[:60]!r}")
print(f"first 20 ids:    {IDS[:20].tolist()}")

class CharLM(nn.Module):
    """A minimal causal transformer language model.

    One encoder block, causal mask, learned positional embeddings — the
    smallest thing that legitimately deserves the name.
    """

    def __init__(self, vocab: int, d_model: int = 64, context: int = 32) -> None:
        super().__init__()
        self.context = context
        self.tok_emb = nn.Embedding(vocab, d_model)
        self.pos_emb = nn.Embedding(context, d_model)
        self.block = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=4, dim_feedforward=128,
            batch_first=True, activation="gelu",
        )
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        _B, T = x.shape
        pos = torch.arange(T, device=x.device)
        h = self.tok_emb(x) + self.pos_emb(pos)[None]
        # Recent torch versions require an explicit attn_mask alongside
        # is_causal=True.
        mask = nn.Transformer.generate_square_subsequent_mask(T, device=x.device)
        h = self.block(h, src_mask=mask, is_causal=True)
        return self.head(self.ln_f(h))


torch.manual_seed(0)
model = CharLM(VOCAB, d_model=64, context=32)
n_params = sum(p.numel() for p in model.parameters())
print(f"model: {n_params:,} parameters (~{n_params * 4 / 1024:.1f} KB in fp32)")

with torch.inference_mode():
    logits = model(IDS[:8].unsqueeze(0))
print(f"forward shape: {tuple(logits.shape)}  (expected (1, 8, {VOCAB}))")
s.check("model_forward_shape_correct", lambda: logits.shape == (1, 8, VOCAB))

Training objective#

We slide a window of length 32 over the text and ask the model to predict each next character. The loss is the standard per-token cross-entropy; the optimiser is AdamW. 60 steps is enough for this corpus.

The important thing to notice: we don’t need to think about perplexity during training. Perplexity is purely a post-hoc reporting metric - you evaluate it from the loss, you don’t optimise it directly.

opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
model.train()

N_STEPS = 400
losses: list[float] = []
for step in range(N_STEPS):
    i = step % (len(IDS) - 33)
    batch = IDS[i : i + 33].unsqueeze(0)
    x, y = batch[:, :-1], batch[:, 1:]
    logits = model(x)
    loss = F.cross_entropy(logits.reshape(-1, VOCAB), y.reshape(-1))
    opt.zero_grad()
    loss.backward()
    opt.step()
    losses.append(loss.item())

model.eval()
print(f"loss[0]   = {losses[0]:.3f}   (initial - should be near log V = {math.log(VOCAB):.3f})")
print(f"loss[{N_STEPS - 1}] = {losses[-1]:.3f}   (after {N_STEPS} steps)")

s.check(
    "training_loss_drops_below_log_vocab",
    lambda: losses[-1] < math.log(VOCAB),
    msg=f"final loss {losses[-1]:.3f} vs log V = {math.log(VOCAB):.3f}",
)

Checkpoint. The initial loss should be very close to \(\log V\), because at step 0 the model is essentially uniform (random-init + tiny vocabulary spread). After a few dozen steps the loss drops substantially. You’re now looking at a (tiny, memorised) language model.

Step 3 - Evaluating perplexity three ways#

In principle we’d score the whole corpus with a single forward pass: one big call, one big average. In practice two things get in the way:

Context length is finite. A model trained with context 32 cannot attend across 10 000 characters in one go. Real LLMs have the same problem at 4 K, 32 K, or 128 K tokens.
Memory is finite. Even if your model could attend over the whole corpus, a length-\(n\) forward costs \(\mathcal{O}(n^2)\) attention memory.

So we evaluate in chunks. Three strategies, each standard:

Strategy	How it splits	What it over- or under-counts
Non-overlapping	`[0:ctx]`, `[ctx:2·ctx]`, …	Early tokens in each window see no left context → pessimistic (inflated PPL).
Sliding window	windows of length `ctx` with stride `s < ctx`; score only the last `s` tokens of each	Almost every token is scored with a full context of left history. More compute, lower PPL.
Single pass	one forward over the whole thing	Reference implementation for short sequences. Fails for long corpora.

We’ll implement all three and verify that the sliding-window estimate is never worse than the non-overlapping one.

@torch.inference_mode()
def perplexity_non_overlapping(ids: torch.Tensor, ctx: int) -> float:
    '''Score ``ids`` in chunks of ``ctx`` with no overlap (fast, pessimistic).'''
    total_nll = 0.0
    total_tok = 0
    for i in range(0, len(ids) - 1, ctx):
        chunk = ids[i : i + ctx + 1]
        if len(chunk) < 2:
            break
        x, y = chunk[:-1].unsqueeze(0), chunk[1:]
        logits = model(x)[0]
        # sum-reduction so we can average correctly across unevenly-sized chunks.
        total_nll += F.cross_entropy(logits, y, reduction="sum").item()
        total_tok += len(y)
    return math.exp(total_nll / total_tok)


@torch.inference_mode()
def perplexity_sliding(ids: torch.Tensor, ctx: int, stride: int) -> float:
    '''Sliding window: overlap contexts, but only score each token once.

    For each window starting at ``begin``, we already scored everything up
    to ``prev_end`` in previous windows. So we only accumulate loss for
    positions in ``[prev_end, end)`` of the corpus.
    '''
    total_nll = 0.0
    total_tok = 0
    prev_end = 0
    for begin in range(0, len(ids), stride):
        end = min(begin + ctx, len(ids))
        if end - begin < 2:
            break
        chunk = ids[begin:end]
        x = chunk[:-1].unsqueeze(0)
        y = chunk[1:]
        logits = model(x)[0]
        score_from = max(0, prev_end - begin - 1)
        if score_from >= len(y):
            prev_end = end
            continue
        total_nll += F.cross_entropy(logits[score_from:], y[score_from:], reduction="sum").item()
        total_tok += len(y) - score_from
        prev_end = end
        if end == len(ids):
            break
    return math.exp(total_nll / total_tok)


@torch.inference_mode()
def perplexity_single_pass(ids: torch.Tensor, ctx: int) -> float:
    '''Single-pass perplexity on the first ``ctx + 1`` tokens - reference only.'''
    chunk = ids[: ctx + 1]
    x, y = chunk[:-1].unsqueeze(0), chunk[1:]
    logits = model(x)[0]
    return math.exp(F.cross_entropy(logits, y, reduction="mean").item())


CTX = model.context
ppl_non = perplexity_non_overlapping(IDS, ctx=CTX)
ppl_slid = perplexity_sliding(IDS, ctx=CTX, stride=CTX // 2)
ppl_ref = perplexity_single_pass(IDS, ctx=CTX)

print(f"non-overlapping      PPL = {ppl_non:.3f}")
print(f"sliding (stride=½ctx) PPL = {ppl_slid:.3f}")
print(f"single-pass (ref)     PPL = {ppl_ref:.3f}")

Reading the three numbers#

Sliding ≤ non-overlapping on real text. Sliding gives every token more left context, so the conditional probabilities are better and the average loss is lower. On WikiText-103 with a 7 B model the sliding-window estimate is routinely 5-15 % lower.
On our memorised corpus the two nearly coincide. The text here loops every ≈480 characters and the model has seen it end-to-end many times. Once the model has memorised the repetition, extra left context doesn’t help much - the “benefit of more history” story is drowned out by positional-embedding quirks and stochastic training effects. You will occasionally see the sliding estimate come in a fraction of a percent above non-overlapping on this toy corpus. That tells you the test corpus is too easy, not that the sliding algorithm is broken.
Single-pass is a reference for short sequences. It’s the most accurate in principle because it uses the longest possible context, but it can only score up to the model’s own context tokens and requires \(n\) to fit in memory. Useful for debugging, unusable for production eval.

The scoring below allows sliding and non-overlapping to differ by up to 5 % in either direction - matching the real-world signal you can expect from a toy corpus. On a proper benchmark (WikiText-2, a few hundred thousand tokens of diverse text) the 5 % budget becomes a strict inequality.

s.check(
    "sliding_and_non_overlapping_agree_within_5pct",
    lambda: abs(ppl_slid - ppl_non) / ppl_non <= 0.05,
    msg=f"sliding={ppl_slid:.3f}  non-overlapping={ppl_non:.3f}  rel_diff={abs(ppl_slid - ppl_non) / ppl_non:.2%}",
)
s.check(
    "corpus_ppl_well_below_random",
    lambda: max(ppl_non, ppl_slid) < VOCAB,
    msg=f"non={ppl_non:.3f}  slid={ppl_slid:.3f}  uniform-baseline={VOCAB}",
)
# Note: single-pass PPL over only ``ctx`` characters can exceed V on short
# samples - the V upper bound is a statement about long-corpus averages,
# not about every 32-token slice. That's one of several reasons nobody
# uses single-pass PPL as a serious evaluation metric.
s.check(
    "every_ppl_above_one",
    lambda: all(p > 1.0 and math.isfinite(p) for p in [ppl_non, ppl_slid, ppl_ref]),
    msg=f"values=({ppl_non:.3f}, {ppl_slid:.3f}, {ppl_ref:.3f})",
)

# Determinism: rerun the same evaluation and demand identical output.
ppl_non_again = perplexity_non_overlapping(IDS, ctx=CTX)
s.assert_close("perplexity_is_deterministic", actual=ppl_non_again, expected=ppl_non, rtol=1e-5)

# Bits-per-character: a nice CS-flavoured sanity check. log₂(PPL) is the
# number of bits the model spends on each character (Shannon source coding).
# For context: an uncompressed ASCII stream is 8 bits/char; gzip on English
# text gets around 2.5 bits/char; a well-trained LLM can get under 1.
bpc = math.log2(ppl_non)
print(f"bits-per-character = log₂(PPL) = {bpc:.3f} bits/char")
s.check(
    "bpc_better_than_ascii_trivial_bound",
    lambda: bpc < 8.0,
    msg=f"bpc={bpc:.3f}  (uncompressed ASCII = 8.0 bits/char)",
)

One picture: the three estimates side by side#

A bar chart puts the three evaluation strategies on the same axis, with the uniform-model upper bound (\(V\)) and the trivial lower bound (1) drawn in. This is the figure to include in any PPL report: it tells the reader at a glance how the estimator was chosen and how close the number is to the two endpoints of the meaningful range.

import matplotlib.pyplot as plt

labels = ["non-overlapping", "sliding (½ctx)", "single-pass"]
values = [ppl_non, ppl_slid, ppl_ref]
colors = ["tab:blue", "tab:green", "tab:gray"]

fig, ax = plt.subplots(figsize=(6.5, 3.6))
bars = ax.bar(labels, values, color=colors)
ax.axhline(VOCAB, color="tab:red",  linestyle="--", label=f"uniform baseline (V={VOCAB})")
ax.axhline(1.0,   color="black",    linestyle=":",  label="perfect model (PPL=1)")
ax.set_ylabel("perplexity")
ax.set_title(f"perplexity across evaluation strategies  (bpc={math.log2(ppl_non):.2f})")
for b, v in zip(bars, values):
    ax.text(b.get_x() + b.get_width() / 2, v, f"{v:.2f}", ha="center", va="bottom", fontsize=9)
ax.set_ylim(0, max(VOCAB, max(values)) * 1.15)
ax.legend(loc="upper right", fontsize=9)
fig.tight_layout()
plt.show()

Exercises#

Smaller stride. Rerun perplexity_sliding with stride=1. Best PPL you’ll see, at the cost of one forward per token.
Different context lengths. Train with context=64 and context=16. Does doubling context always lower PPL?
Temperature scaling. Divide the logits by τ before softmax; at what τ does PPL stop improving?
Break the grader. Construct a model whose non_overlapping PPL beats sliding. (You shouldn’t be able to — if you can, the sliding implementation has a bug.)

Perplexity from scratch

Contents

Perplexity from scratch#

A quick refresher - skip if you’re comfortable#

Language models are probability distributions over sequences#

Cross-entropy measures how well the model fits the data#

Perplexity is just `exp(loss)`#

Step 1 - Perplexity from first principles#

Step 2 - A minimal character-level language model#

Training objective#

Step 3 - Evaluating perplexity three ways#

Reading the three numbers#

One picture: the three estimates side by side#

Exercises#

Further reading#

Perplexity from scratch

Contents

Perplexity from scratch#

A quick refresher - skip if you’re comfortable#

Language models are probability distributions over sequences#

Cross-entropy measures how well the model fits the data#

Perplexity is just exp(loss)#

Step 1 - Perplexity from first principles#

Step 2 - A minimal character-level language model#

Training objective#

Step 3 - Evaluating perplexity three ways#

Reading the three numbers#

One picture: the three estimates side by side#

Exercises#

Further reading#

Perplexity is just `exp(loss)`#