Token Cross-Entropy and Perplexity - LLM Internals

Token Cross-Entropy and Perplexity

What you'll build

The language-modeling loss and its standard evaluation metric: per-token negative log-likelihood from logits and targets, with a padding mask, averaged into a mean NLL and exponentiated into perplexity.

Background

Every LLM is pretrained by minimizing next-token cross-entropy, and perplexity is how papers report it: "the effective number of tokens the model is choosing between". A model that is uniform over V tokens has perplexity exactly V; a perfect model has perplexity 1. Batches are padded to a common length, so padded positions must be excluded, or the reported loss depends on how much padding the batch happened to have.

Formula

log_probs  = log_softmax(logits, axis=-1)          # z - max - log(sum(exp(z - max)))
token_nll  = -log_probs[..., target]               # one value per position
mean_nll   = sum(token_nll * mask) / sum(mask)     # nats per real token
perplexity = exp(mean_nll)

Inputs

• logits: float array (..., V), e.g. (T, V) or (B, T, V)
• targets: integer array (...) of token ids in [0, V); targets[b, t] is the token that follows position t (already shifted)
• mask: array (...) with 1 for real tokens and 0 for padding, or None (every position counts)

Expected Return Format

A dictionary with these keys:

• 'token_nll': per-position NLL, same shape as targets, 0 where masked
• 'mean_nll': mean NLL over unmasked positions (Python float, in nats)
• 'perplexity': exp(mean_nll) (Python float)

Edge cases and conventions

• Use a stable log-softmax: logits around 1000 must not overflow.
• Divide by the number of real tokens (mask.sum()), not by the number of positions.
• Natural log (nats). Bits per token would be mean_nll / ln 2.

Constraints

NumPy only, vectorised (np.take_along_axis for the target gather).

Follow-up

• Why is perplexity only comparable between models that share a tokenizer? How does bits-per-byte fix that?
• Training code often averages per sequence first, then over the batch. When does that differ from the per-token mean here?

Input: perplexity(np.zeros((3, 5)), np.array([0, 1, 2]))
Output: {'token_nll': [1.6094, 1.6094, 1.6094], 'mean_nll': 1.6094, 'perplexity': 5.0}

Input: perplexity(np.array([[2.0, 1.0, 0.0], [0.0, 0.0, 3.0]]), np.array([0, 2]))
Output: token_nll = [0.4076, 0.0949]

import numpy as np

def perplexity(logits, targets, mask=None):
    """
    Next-token cross-entropy and perplexity.

    Args:
        logits: (..., V)
        targets: (...) integer token ids
        mask: (...) 1 for real tokens, 0 for padding (None = all real)

    Returns:
        dict with 'token_nll' (same shape as targets, 0 where masked),
        'mean_nll' (float, nats) and 'perplexity' (float, exp(mean_nll))
    """
    # Your code here
    pass