Perplexity

Perplexity is the exponentiated cross-entropy per token. A perplexity of 12 means the model is, on average, as uncertain as if choosing uniformly among twelve options.

It is not comparable across tokenisers. A model with a larger vocabulary spends fewer tokens on the same text, so its per-token perplexity falls without any improvement in modelling. Bits per byte normalises this away and is the only fair cross-model comparison.

import torch
from torch import Tensor
import torch.nn.functional as F


@torch.no_grad()
def perplexity(model, ids: Tensor, window: int = 2048, stride: int = 512):
    nll, count = 0.0, 0
    for start in range(0, ids.size(1) - 1, stride):
        end = min(start + window, ids.size(1))
        chunk = ids[:, start:end]
        logits = model(chunk[:, :-1]).logits

        # only score the tokens this stride is responsible for
        scored = min(stride, chunk.size(1) - 1)
        loss = F.cross_entropy(
            logits[:, -scored:].flatten(0, 1),
            chunk[:, -scored:].flatten(),
            reduction="sum",
        )
        nll += loss.item()
        count += scored
    return torch.tensor(nll / count).exp().item()

Scoring only the last stride positions of each window is what makes the estimate a proper per-token average rather than a weighted one that double-counts the overlap. Sum the loss and divide once at the end; averaging per-chunk means and then averaging those weights short final chunks equally with full ones.