Easy Primitives

Perplexity from Logits

Perplexity summarizes how much probability a model assigns to the supplied target tokens. Compute the average negative log-probability of those targets, then exponentiate. Lower perplexity means the model assigned greater geometric mean probability to those same targets.

Your task

Implement perplexity(logits, targets).

logits is a finite floating-point tensor with shape (N, T, V).
targets has shape (N, T) and contains integer-valued IDs in [0, V). The runtime delivers them as floats; cast them to integers for indexing.
N, T and V are positive. Test results are representable in the input dtype.
Return one Python float, leaving both inputs unchanged.

Targets are already aligned: targets[n, t] is the correct token for logits[n, t, :]. Do not shift either tensor, discard the last position or insert a start token. Every one of the N*T positions counts equally; there is no padding mask or ignored target ID in this exercise.

Calculation

Compute log-softmax over the final vocabulary axis for each position.
Select the log-probability at that position's target ID.
Negate the mean of the selected values across all N*T positions.
Return the natural exponential of that mean loss.

mean_nll = -sum(target_log_probabilities) / (N * T)
perplexity = exp(mean_nll)

Use a stable log-softmax. One formula is shifted - log(sum(exp(shifted))), where shifted is a vocabulary row minus its maximum. Computing softmax and then taking its logarithm can lose very small target probabilities to underflow. Do not average per-token perplexities or exponentiate the total loss.

Example and interpretation

If two target probabilities are 1/2 and 1/8, the mean loss is (log(2) + log(8)) / 2, giving perplexity 4. It is the geometric mean of inverse target probabilities, not their arithmetic mean of 5.

Uniform predictions over V tokens give perplexity V. Perplexity approaches 1 as every target probability approaches 1, but it can exceed V when the model assigns targets very low probability. It is not a literal count of options considered at every position.

Compare perplexities using the same evaluation data, tokenization and scoring rules. A lower value under a different tokenization is not directly comparable.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Match each prediction to the token it should predict

A causal language model predicts the next token from the tokens before it. For sequence [A, B, C], logits produced after reading A are trained against B, and logits after reading B are trained against C. This is why input and target sequences are usually shifted by one position. A causal attention mask prevents access to later tokens that would reveal the answer.

A masked language model instead predicts selected tokens from a corrupted sequence and can use context on both sides. The supervised positions are not necessarily just those showing a mask token. A corruption rule may leave some selected tokens unchanged or replace them randomly while still training on them. Compute the loss over the specified supervised positions only.

Teacher forcing supplies the correct previous tokens

In encoder-decoder training, teacher forcing feeds the decoder the correct previous target tokens. During generation, it usually receives its own previous choices instead. A mistake can therefore change later inputs, which helps explain why training loss alone does not describe generation quality.

Convert mean token loss to perplexity

For M valid target tokens, mean negative log likelihood is

$$L=-\frac1M\sum_t\log p(y_t).$$

With natural logarithms, perplexity is exp(L). A model that gives the correct token probability 1/4 at every step has perplexity 4. A uniform distribution over a vocabulary of size V has perplexity V.

When batches contain different numbers of valid tokens, add their token losses and valid-token counts before dividing. Averaging batch means would give small batches too much weight. Exclude padding and other unsupervised positions as specified, and define what happens when no tokens are supervised.

Perplexity also depends on tokenization. Scores based on different token units are not directly comparable just because both are called perplexity.

Detecting runtime…