1.2 Probability, Statistics & Information Theory¶

Every LLM training run, every forward pass, every evaluation metric you encounter reduces to a single fundamental operation: computing a probability and then measuring how far it is from what you wanted. The model output is a probability distribution over the vocabulary. The training loss is the cross-entropy between that distribution and the ground truth. The evaluation metric is perplexity, itself an exponential of cross-entropy. To work fluently at any level of the stack, you need to understand the probability and information theory that underlie all of it — not just as formulas to memorize, but as interconnected ideas with a coherent geometry.

This chapter builds those ideas from the ground up. We start with random variables and the laws that govern them, move through the statistical estimation procedures that produce model parameters, and arrive at Shannon’s information theory — entropy, KL divergence, cross-entropy, mutual information — concluding with perplexity and a precise explanation of why cross-entropy is the right loss for language modeling. Every concept comes with a worked numerical example or code you can run.

Related chapters: Linear Algebra for Deep Learning covers the vector and matrix operations used throughout; Calculus, Optimization & Convexity develops the gradient mechanics that minimize the cross-entropy loss; and The Pretraining Objective & Loss applies everything here to real training pipelines.

Random Variables and Distributions¶

A random variable \(X\) is a function from a probability space (a set of outcomes) to a measurable space (usually \(\mathbb{R}\)). We typically care less about the formalism and more about the distribution it induces: the assignment of probability to events.

Discrete distributions are specified by a probability mass function (PMF):

Continuous distributions use a probability density function (PDF) \(p(x)\) where \(P(a \leq X \leq b) = \int_a^b p(x)\,dx\).

Distributions You Will Encounter in LLM Work¶

The categorical distribution is the star of this textbook. The softmax output of a language model is nothing but the parameter vector of a categorical distribution over the vocabulary of size \(V\) (typically on the order of 32,000–256,000 tokens).

Expectation, Variance, and Covariance¶

The expectation (mean) of \(X\) under distribution \(p\):

Variance measures spread:

Covariance between two random variables \(X\) and \(Y\):

The Pearson correlation is the normalized version: \(\rho = \operatorname{Cov}(X,Y) / (\sigma_X \sigma_Y)\).

Bayes’ Theorem¶

Bayes’ theorem is the mechanism by which we update beliefs in light of evidence:

where \(\theta\) are model parameters, \(\mathcal{D}\) is observed data, \(P(\theta)\) is the prior, \(P(\mathcal{D} \mid \theta)\) is the likelihood, \(P(\mathcal{D})\) is the marginal likelihood (evidence, a normalizing constant), and \(P(\theta \mid \mathcal{D})\) is the posterior. In LLM training we almost never compute the posterior directly — we optimize instead — but Bayes frames the conceptual story.

Maximum Likelihood and MAP Estimation¶

Maximum Likelihood Estimation¶

Suppose we have a parametric model \(p(x \mid \theta)\) and i.i.d. observations \(\mathcal{D} = \{x_1, \ldots, x_N\}\). The log-likelihood is:

Maximum likelihood estimation (MLE) finds:

For a language model, each observation \(x_i\) is a token, \(p(x_i \mid \theta) = p_\theta(x_i \mid x_{<i})\) is the model’s predicted probability of that token given context, and maximizing the log-likelihood is exactly minimizing the cross-entropy loss:

This is not a coincidence. Cross-entropy loss is negative log-likelihood. The two names refer to the same operation viewed from different disciplines.

MAP Estimation and Regularization¶

Maximum a posteriori (MAP) estimation incorporates the prior:

A Gaussian prior \(P(\theta) \propto \exp(-\lambda\|\theta\|^2)\) corresponds to \(\ell_2\) (weight decay) regularization. A Laplace prior corresponds to \(\ell_1\) (sparsity-inducing) regularization. When practitioners add weight_decay=0.1 to AdamW, they are performing MAP estimation with a Gaussian prior — this is discussed further in Calculus, Optimization & Convexity and Optimizers: SGD, Adam, Adafactor, Lion, Muon & Shampoo.

Shannon Information Theory¶

Claude Shannon’s 1948 paper “A Mathematical Theory of Communication” introduced the concepts we rely on every day in ML. The core idea: information content should be a function of probability, should be additive for independent events, and should be larger for rare events.

Self-Information and Entropy¶

The self-information (surprisal) of an event with probability \(p\):

If we use natural log instead of \(\log_2\), the unit is nats. Deep learning conventions use nats (because differentiation of \(\log\) is cleaner), while information-theoretic literature often uses bits. The conversion is \(1 \text{ nat} = \log_2 e \approx 1.4427 \text{ bits}\).

The Shannon entropy of a distribution \(p\) is the expected self-information — the average surprise:

Entropy is maximized by the uniform distribution (\(H = \log V\) for \(V\) outcomes) and minimized at zero by any degenerate (one-hot) distribution.

KL Divergence¶

The Kullback-Leibler (KL) divergence from distribution \(q\) to distribution \(p\) measures how much information is lost when \(q\) is used to approximate \(p\):

Key properties: - Non-negative: \(D_{\text{KL}}(p \,\|\, q) \geq 0\), with equality iff \(p = q\) (by Gibbs’ inequality) - Asymmetric: \(D_{\text{KL}}(p \,\|\, q) \neq D_{\text{KL}}(q \,\|\, p)\) in general - Not a metric: violates the triangle inequality - Infinite if \(q(x)=0\) where \(p(x)>0\): this is why label smoothing matters (more below)

The forward KL \(D_{\text{KL}}(p\|q)\) is called inclusive — minimizing it forces \(q\) to cover all modes of \(p\). The reverse KL \(D_{\text{KL}}(q\|p)\) is exclusive — minimizing it lets \(q\) concentrate on one mode of \(p\). This asymmetry is critical in variational inference and in RLHF/DPO, where we penalize the KL between the fine-tuned policy and the reference model.

One practical wrinkle is worth flagging here because it follows directly from the definition. Between two next-token distributions the KL is an exact finite sum over the \(V\) vocabulary entries — you have both logit vectors, so nothing is estimated. Between two policies over whole sequences the sum ranges over \(V^T\) possible continuations and is hopelessly intractable, so RL loops replace it with a Monte-Carlo estimate from the single token that was actually sampled. The naive estimator \(\hat{k}_1 = \log\frac{\pi_\theta(a)}{\pi_{\text{ref}}(a)}\) is unbiased but can be negative on any individual sample, which makes for a strange penalty; the standard fix is the unbiased and non-negative estimator \(\hat{k}_3 = e^{r} - 1 - r\) with \(r = \log\frac{\pi_{\text{ref}}(a)}{\pi_\theta(a)}\), derived and compared in Advantage Estimation, KL Control & Stability Tricks.

Cross-Entropy¶

Cross-entropy between distribution \(p\) (true) and \(q\) (model) is:

It relates to entropy and KL as:

This decomposition is crucial. Since \(H(p)\) is a constant with respect to \(q\) (it depends only on the true distribution), minimizing \(H(p,q)\) over \(q\) is equivalent to minimizing \(D_{\text{KL}}(p \,\|\, q)\). Training a language model by cross-entropy loss is fitting the model distribution to the data distribution in the KL-divergence sense.

In language modeling the true distribution \(p\) is the one-hot empirical distribution: all probability mass on the observed token \(x^*\). Therefore:

The cross-entropy loss is simply the negative log-probability the model assigns to the correct token. No sum needed — one-hot \(p\) kills all other terms.

Why Cross-Entropy is THE LLM Loss¶

Several properties conspire to make cross-entropy the right loss for language models.

1. It is the negative log-likelihood. For categorical outputs, minimizing cross-entropy is identical to MLE. MLE is consistent and asymptotically efficient (the Cramér-Rao bound) under standard regularity conditions.

2. It is a proper scoring rule. A loss \(\ell(p, y)\) is proper if the minimum is achieved exactly when \(p\) matches the true distribution. Cross-entropy (log loss) is strictly proper: the model is maximally rewarded for reporting its true beliefs.

3. It has well-behaved gradients. The gradient of cross-entropy with softmax output is \(\hat{p} - p_{\text{true}}\), the probability residual. No vanishing gradients even when the correct token probability is small (unlike squared error, whose gradient \((\hat{p} - p_{\text{true}}) \cdot \hat{p}(1-\hat{p})\) vanishes when \(\hat{p} \to 0\)).

4. It can handle uncertainty. If the true label distribution is not one-hot (label smoothing, soft targets from a teacher), cross-entropy still works: \(-\sum_k p_k \log q_k\).

5. Minimizing it minimizes KL. As shown above, it directly minimizes \(D_{\text{KL}}(p_{\text{data}} \| p_\theta)\), which is the information-theoretic statement that the model matches the data-generating process.

The pretraining loss used by GPT-2, GPT-3, LLaMA, Gemini, and essentially every modern LLM is precisely:

See The Pretraining Objective & Loss for exactly how this is computed in practice at scale, and The Pretraining Run: A Complete Single-GPU Training Loop for the loss curve this objective produces when you actually run it end to end on a ~100M-parameter model.

Mutual Information and Conditional Entropy¶

Conditional entropy measures remaining uncertainty about \(X\) given knowledge of \(Y\):

Mutual information (MI) quantifies how much knowing \(Y\) reduces uncertainty about \(X\):

Equivalently:

MI is symmetric: \(I(X;Y) = I(Y;X)\). It is zero iff \(X\) and \(Y\) are independent. In information-theoretic terms, two variables are dependent iff their joint distribution differs from the product of marginals.

Where does MI appear in LLMs? Several places:

• InfoNCE loss (used in contrastive pretraining like SimCLR, and in some embedding models) is a lower bound on MI.
• Probing classifiers measure MI between internal representations and linguistic features to understand what a model has learned.
• Feature selection in attention analysis: the attention pattern can be viewed as routing information; MI-based measures quantify how much information flows through a head.
• Tokenization: BPE merge criteria approximate minimizing description length, which is related to maximizing MI between pairs of adjacent subwords — see Tokenization: BPE, WordPiece, Unigram & Byte-Level.

Perplexity¶

Perplexity (PPL) is the standard metric for evaluating language models. It is the exponentiation of the average cross-entropy loss:

Intuition: a perplexity of \(K\) means the model is “as confused as if it had to choose uniformly among \(K\) options at each step.” A random model over a vocabulary of size \(V\) achieves \(\text{PPL} = V\). A perfect model achieves \(\text{PPL} = 1\).

Perplexity Pitfalls¶

Perplexity depends strongly on tokenization. A model using a byte-level tokenizer will report higher perplexity than one using a word-level tokenizer because more decisions are made per word. Always compare perplexity numbers computed with the same tokenizer on the same test set.

The fix, when tokenizers differ, is to normalize by a unit the tokenizer cannot change: the raw bytes of the text. Bits per byte (BPB) divides the total negative log-likelihood of the passage, converted to bits, by the number of UTF-8 bytes in the original string:

where \(\mathcal{L}_{\text{CE}}\) is the mean per-token loss in nats and \(T\) the token count. The numerator is precisely the number of bits an optimal arithmetic coder would need to store the text given the model, so BPB is a compression rate — and unlike perplexity it is comparable across any two models regardless of vocabulary size. Modern LLMs land below 1 bit per byte on ordinary English (i.e. better than 8:1 compression). This is the metric lm-evaluation-harness reports as bits_per_byte; see The Pretraining Objective & Loss for the training-loop implementation and Evaluation & Serving: Honest Benchmarks, int4 Quantization, and Running on a Laptop for how Stack-100M is scored with it.

Perplexity is also not a direct proxy for downstream task quality — a model can have lower perplexity than another while being worse at instruction following or reasoning. See The Evaluation Problem & Benchmark Landscape for a fuller treatment.

Label Smoothing and Soft Targets¶

Label smoothing replaces the one-hot target with a smoothed distribution:

where \(\varepsilon\) is a small constant (commonly 0.1) and \(V\) is the vocabulary size. The loss becomes:

Why does this help? Three reasons:

1. Calibration: pure cross-entropy with hard targets can push logits for the correct class to \(+\infty\), producing overconfident models. Smoothing prevents this.
2. KL guard: when \(p(x) = 0\) and \(q(x) = 0\), the KL term \(0 \cdot \log 0 / 0\) is undefined; smoothing ensures \(p(x) > 0\) everywhere, making KL well-defined.
3. Regularization: the smoothed objective implicitly penalizes the entropy of the output distribution, preventing the model from assigning zero probability to any class and encouraging more distributed predictions.

In the original Transformer paper (Vaswani et al., “Attention Is All You Need”, 2017), label smoothing of \(\varepsilon = 0.1\) was used and attributed a significant improvement in BLEU score. LLM pretraining today typically does not use label smoothing (the model scale provides sufficient implicit regularization), but it remains common in fine-tuning and machine translation.

Code: Computing Perplexity From Scratch¶

The following is a complete, runnable Python module that computes perplexity at multiple levels of abstraction — from raw logits to a full dataset evaluation loop. It mirrors what frameworks like lm-evaluation-harness do internally.

"""
perplexity.py — compute LLM perplexity from scratch
Requires: torch, transformers (for a demo tokenizer/model)
Run: python perplexity.py
"""

import math
import torch
import torch.nn.functional as F
from typing import List, Optional


# ─────────────────────────────────────────────────────
# 1. Low-level: perplexity from a list of log-probs
# ─────────────────────────────────────────────────────

def perplexity_from_log_probs(log_probs: List[float]) -> float:
    """
    Given a list of natural-log probabilities log p(x_t | x_{<t}),
    compute perplexity = exp( -mean(log_probs) ).

    Args:
        log_probs: list of floats, each <= 0.

    Returns:
        Perplexity as a float.
    """
    if not log_probs:
        raise ValueError("Empty log-prob list")
    avg_nll = -sum(log_probs) / len(log_probs)   # average negative log-likelihood
    return math.exp(avg_nll)


# ─────────────────────────────────────────────────────
# 2. From raw logits and target token ids
# ─────────────────────────────────────────────────────

def perplexity_from_logits(
    logits: torch.Tensor,   # shape (T, V) — one logit vector per time step
    targets: torch.Tensor,  # shape (T,)   — ground-truth token ids
    ignore_index: int = -100,
) -> float:
    """
    Compute perplexity from raw (unnormalized) logits.

    This is exactly what a training loop would call after the forward pass,
    except we expose each step for clarity.

    Steps:
      1. Apply log-softmax to get log-probabilities.
      2. Gather the log-prob of the correct token at each step.
      3. Mask out padding tokens (ignore_index).
      4. Average and exponentiate.
    """
    # (T, V) → (T, V) in log-probability space
    log_probs = F.log_softmax(logits, dim=-1)   # numerically stable via LogSumExp

    # Gather log-prob of the target token at each position.
    # targets shape: (T,) → unsqueeze to (T, 1) for gather
    valid_mask = targets != ignore_index
    safe_targets = targets.clone()
    safe_targets[~valid_mask] = 0  # avoid index error on masked positions

    # Shape: (T, 1) → squeeze to (T,)
    token_log_probs = log_probs.gather(1, safe_targets.unsqueeze(1)).squeeze(1)

    # Zero out masked positions before averaging
    token_log_probs = token_log_probs * valid_mask.float()

    n_valid = valid_mask.sum().item()
    avg_nll = -token_log_probs.sum().item() / n_valid
    return math.exp(avg_nll)


# ─────────────────────────────────────────────────────
# 3. Dataset-level perplexity with sliding window
#    (handles sequences longer than the model context)
# ─────────────────────────────────────────────────────

def sliding_window_perplexity(
    token_ids: torch.Tensor,   # shape (N_total,)
    logit_fn,                  # callable: (T,) -> (T, V) logits
    max_length: int = 512,
    stride: int = 256,
) -> float:
    """
    Compute perplexity over a long token sequence using a sliding window.

    The 'stride' trick ensures each token is scored in context,
    not at the very beginning of a truncated window where the model
    has no context. This is the method used by Radford et al. (GPT-2)
    for WikiText-103 evaluation.

    Args:
        token_ids:  1-D tensor of all token ids in the test set.
        logit_fn:   function mapping a 1-D token tensor of length <= max_length
                    to a logit tensor of shape (len, vocab_size).
        max_length: model's maximum context length.
        stride:     step between windows; lower stride = more context overlap
                    = slightly slower but more accurate for longer texts.

    Returns:
        Perplexity (float).
    """
    seq_len = token_ids.size(0)
    total_nll = 0.0
    total_tokens = 0
    prev_end = 0

    for begin in range(0, seq_len, stride):
        end = min(begin + max_length, seq_len)
        window = token_ids[begin:end]          # shape: (window_size,)

        # The model scores positions [0, window_size-1];
        # for LM evaluation we predict token t+1 from tokens 0..t.
        with torch.no_grad():
            logits = logit_fn(window)          # shape: (window_size, V)

        # Targets are the next token at each position:
        # position i predicts token i+1, so we shift by 1.
        # In a causal LM, logits[i] predicts token i+1 (the standard convention).
        # We count only tokens that were *not* already counted in the previous window.
        target_ids = token_ids[begin + 1 : end + 1]
        # Use only new positions (stride steps from the right of the window)
        count_from = max(prev_end - begin, 1)  # at least predict 1 token

        logits_new = logits[count_from - 1 : len(target_ids)]
        targets_new = target_ids[count_from - 1 :]

        if logits_new.shape[0] == 0:
            prev_end = end
            continue

        log_probs = F.log_softmax(logits_new, dim=-1)
        token_nlls = -log_probs.gather(
            1, targets_new[:logits_new.shape[0]].unsqueeze(1)
        ).squeeze(1)

        total_nll += token_nlls.sum().item()
        total_tokens += token_nlls.shape[0]
        prev_end = end

        if end == seq_len:
            break

    return math.exp(total_nll / total_tokens)


# ─────────────────────────────────────────────────────
# 4. Demonstration: cross-entropy decomposition
# ─────────────────────────────────────────────────────

def demonstrate_ce_kl_relationship():
    """
    Show numerically that H(p, q) = H(p) + KL(p || q).
    Uses a small 4-class example.
    """
    # True distribution p (label smoothed example)
    p = torch.tensor([0.7, 0.1, 0.1, 0.1])
    # Model distribution q
    q = torch.tensor([0.5, 0.2, 0.2, 0.1])

    assert abs(p.sum().item() - 1.0) < 1e-6, "p must be a valid distribution"
    assert abs(q.sum().item() - 1.0) < 1e-6, "q must be a valid distribution"

    # Shannon entropy H(p)
    # Convention: 0 * log(0) = 0
    H_p = -(p * torch.log(p.clamp(min=1e-12))).sum().item()

    # KL divergence KL(p || q)
    # Sum over positions where p > 0
    kl_pq = (p * torch.log((p / q.clamp(min=1e-12)).clamp(min=1e-12))).sum().item()

    # Cross-entropy H(p, q) = -sum_x p(x) log q(x)
    ce = -(p * torch.log(q.clamp(min=1e-12))).sum().item()

    print(f"H(p)              = {H_p:.4f} nats")
    print(f"KL(p || q)        = {kl_pq:.4f} nats")
    print(f"H(p) + KL(p||q)   = {H_p + kl_pq:.4f} nats")
    print(f"H(p, q) directly  = {ce:.4f} nats")
    print(f"Match: {abs(ce - (H_p + kl_pq)) < 1e-5}")

    # In the one-hot case (standard training), H(p) = 0 so CE = KL
    p_onehot = torch.tensor([1.0, 0.0, 0.0, 0.0])
    H_onehot = 0.0  # entropy of a degenerate distribution
    ce_onehot = -(p_onehot * torch.log(q.clamp(min=1e-12))).sum().item()
    kl_onehot = ce_onehot - H_onehot  # = ce_onehot
    print(f"\nOne-hot target:")
    print(f"H(p_onehot, q) = -log q(k*) = {ce_onehot:.4f} = {-math.log(0.5):.4f}")
    print(f"This equals -log(q[0]) = -log(0.5) = {-math.log(0.5):.4f}")


# ─────────────────────────────────────────────────────
# 5. Label smoothing loss
# ─────────────────────────────────────────────────────

def label_smoothed_cross_entropy(
    logits: torch.Tensor,    # (B, T, V) or (T, V)
    targets: torch.Tensor,   # (B, T) or (T,) long tensor
    smoothing: float = 0.1,
    ignore_index: int = -100,
) -> torch.Tensor:
    """
    Cross-entropy with label smoothing.

    Equivalent to PyTorch's CrossEntropyLoss(label_smoothing=smoothing),
    but written out explicitly for teaching purposes.

    The loss is:
        L = (1 - eps) * CE_hard(logits, targets)
          + eps * mean_over_classes( -log_softmax(logits) )
    """
    vocab_size = logits.shape[-1]
    flat_logits = logits.reshape(-1, vocab_size)   # (B*T, V)
    flat_targets = targets.reshape(-1)              # (B*T,)

    log_probs = F.log_softmax(flat_logits, dim=-1)  # (B*T, V)

    # Hard-target negative log-likelihood
    nll_loss = F.nll_loss(
        log_probs,
        flat_targets,
        ignore_index=ignore_index,
        reduction='mean',
    )

    # Soft-target component: uniform over all classes
    # -mean_k log_prob_k, averaged over valid positions
    smooth_loss = -log_probs.mean(dim=-1)  # (B*T,)
    mask = flat_targets != ignore_index
    smooth_loss = smooth_loss[mask].mean()

    return (1.0 - smoothing) * nll_loss + smoothing * smooth_loss


# ─────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────

if __name__ == "__main__":
    print("=== CE / KL decomposition demo ===")
    demonstrate_ce_kl_relationship()

    print("\n=== Perplexity from raw log-probs ===")
    # Toy sequence: 5 tokens with model-assigned probabilities
    # as in the worked example above
    probs = [0.20, 0.05, 0.30, 0.60, 0.15]
    log_probs_list = [math.log(p) for p in probs]
    ppl = perplexity_from_log_probs(log_probs_list)
    print(f"Perplexity = {ppl:.3f}  (expected ~5.17)")

    print("\n=== Perplexity from logits ===")
    torch.manual_seed(42)
    V, T = 32, 10
    # Simulate a model that has reasonably high confidence on the correct tokens
    logits = torch.randn(T, V)
    targets = torch.randint(0, V, (T,))
    # Boost the correct token logits by 2 so the model looks "smart"
    for t in range(T):
        logits[t, targets[t]] += 2.0
    ppl_from_logits = perplexity_from_logits(logits, targets)
    print(f"Perplexity from logits = {ppl_from_logits:.3f}")

    print("\n=== Label-smoothed loss ===")
    logits_batch = torch.randn(2, 5, 100)  # batch=2, seq_len=5, vocab=100
    targets_batch = torch.randint(0, 100, (2, 5))
    ls_loss = label_smoothed_cross_entropy(logits_batch, targets_batch, smoothing=0.1)
    hard_loss = F.cross_entropy(logits_batch.reshape(-1, 100), targets_batch.reshape(-1))
    print(f"Label-smoothed loss = {ls_loss.item():.4f}")
    print(f"Hard cross-entropy  = {hard_loss.item():.4f}")
    print(f"Label smoothing adds regularization: LS loss > hard CE = {ls_loss.item() > hard_loss.item()}")

Running the script produces output like:

=== CE / KL decomposition demo ===
H(p)              = 0.9404 nats
KL(p || q)        = 0.0969 nats
H(p) + KL(p||q)   = 1.0373 nats
H(p, q) directly  = 1.0373 nats
Match: True

One-hot target:
H(p_onehot, q) = -log q(k*) = 0.6931 = 0.6931
This equals -log(q[0]) = -log(0.5) = 0.6931

=== Perplexity from raw log-probs ===
Perplexity = 5.173  (expected ~5.17)

=== Perplexity from logits ===
Perplexity from logits = 6.538

=== Label-smoothed loss ===
Label-smoothed loss = 5.6843
Hard cross-entropy  = 5.7464
Label smoothing adds regularization: LS loss > hard CE = False

The first block (CE/KL decomposition) is deterministic — p and q are fixed tensors, so those four numbers are exact and will match on any machine. The last two blocks depend on torch.manual_seed(42)-seeded random logits, so the specific numbers (and even the True/False outcome of the last comparison) can differ across PyTorch versions and hardware. The point the code demonstrates is architectural — H(p,q) = H(p) + D_KL(p‖q) holds exactly — not that label smoothing is guaranteed to increase the loss on any single random draw of untrained logits.

The Library Versions of Everything Above¶

You write these by hand once, to know what is inside them, and then you call the library for the rest of your career. Every quantity in this chapter has a one-line implementation in PyTorch — each one a fused, numerically stable kernel built on the logsumexp shift trick derived in Numerical Computing, Floating Point & Precision — and knowing which to reach for, and which has a hostile signature, is part of the working knowledge.

"""
library_versions.py — the same quantities via the standard PyTorch APIs.
Requires: torch. Run: python library_versions.py
"""

import torch
import torch.nn.functional as F
from torch.distributions import Categorical, kl_divergence

torch.manual_seed(0)
logits = torch.randn(4, 8)                # (T=4 positions, V=8 vocabulary)
targets = torch.tensor([2, 5, 0, 7])      # ground-truth next-token ids

# 1. Cross-entropy = NLL = the training loss. Takes RAW logits, never probabilities:
#    it fuses log_softmax + gather + mean, using the logsumexp shift trick derived in
#    the Numerical Computing chapter to stay stable when logits are large.
loss = F.cross_entropy(logits, targets, reduction="mean")   # nats per token
ppl = loss.exp()                                            # perplexity
# Label smoothing and padding are flags here, not separate code paths:
loss_ls = F.cross_entropy(logits, targets, label_smoothing=0.1, ignore_index=-100)

# 2. Entropy of the model's OWN predictive distribution, one value per position.
#    Build the distribution from logits= (not probs=) so the log-softmax stays stable.
dist = Categorical(logits=logits)
H = dist.entropy()               # (4,) nats; the ceiling here is log(8) = 2.079
#    This is exactly the quantity an RL "entropy bonus" maximizes to delay mode collapse.

# 3. KL between two categorical distributions — the RLHF/DPO reference-model penalty.
ref_logits = logits + 0.3 * torch.randn_like(logits)
kl = kl_divergence(Categorical(logits=ref_logits),          # p = reference
                   Categorical(logits=logits))              # q = policy  -> D_KL(p||q)

# 4. The same number via F.kl_div, whose signature is a well-known trap (see below).
kl_f = F.kl_div(
    F.log_softmax(logits, dim=-1),         # `input`  MUST already be LOG-probs (q)
    F.log_softmax(ref_logits, dim=-1),     # `target` (p) — log-probs only because...
    log_target=True,                       # ...we set log_target=True
    reduction="none",
).sum(dim=-1)                              # sum over vocabulary -> per-position KL

if __name__ == "__main__":
    print(f"loss      = {loss.item():.4f} nats   ppl = {ppl.item():.3f}")
    print(f"loss (LS) = {loss_ls.item():.4f} nats")
    print(f"mean H(q) = {H.mean().item():.4f} nats")
    print(f"KL agrees: {torch.allclose(kl, kl_f, atol=1e-5)}")   # True

At the evaluation layer the standard open-source tool is EleutherAI’s lm-evaluation-harness, which wraps the same sliding-window loop written out above and applies it across hundreds of tasks with a fixed, citable methodology:

pip install lm-eval
lm_eval --model hf \
        --model_args pretrained=EleutherAI/pythia-160m \
        --tasks wikitext \
        --batch_size 8

Its wikitext task deliberately reports word-level perplexity, byte-level perplexity, and bits_per_byte side by side — three normalizations of the same total negative log-likelihood — precisely so that the tokenizer-dependence discussed above is visible rather than hidden. Building your own version of this harness is covered in Building Eval Harnesses.

A Unified Picture¶

The diagram below shows how the concepts of this chapter interrelate:

This picture has three take-aways:

1. Cross-entropy loss and KL minimization are the same thing for a fixed data distribution.
2. Perplexity is just a human-readable transformation of the loss.
3. KL reappears in post-training alignment — but in the reverse direction, as a constraint preventing the policy from drifting too far from the reference model.

Additional Interview Questions¶

Further Reading¶

• Shannon, C. E. “A Mathematical Theory of Communication.” Bell System Technical Journal, 1948. The original paper; remarkably readable.
• Cover, T. M., and Thomas, J. A. Elements of Information Theory, 2^nd ed. Wiley, 2006. The definitive graduate textbook.
• Goodfellow, I., Bengio, Y., and Courville, A. Deep Learning. MIT Press, 2016. Chapter 3 (Probability and Information Theory) and Chapter 5 (Machine Learning Basics).
• Bishop, C. M. Pattern Recognition and Machine Learning. Springer, 2006. Chapters 1–2 for distributions and Bayesian estimation.
• Vaswani, A. et al. “Attention Is All You Need.” NeurIPS, 2017. The Transformer paper; introduces label smoothing in the context of machine translation.
• Müller, R., Kornblith, S., and Hinton, G. “When Does Label Smoothing Help?” NeurIPS, 2019. Empirical analysis of label smoothing’s effects on calibration and distillation.
• Radford, A. et al. “Language Models are Unsupervised Multitask Learners.” OpenAI, 2019. (GPT-2 paper) — describes the sliding-window perplexity evaluation methodology.
• lm-evaluation-harness (EleutherAI): open-source framework for evaluating language models, including perplexity across many benchmarks. Available at github.com/EleutherAI/lm-evaluation-harness.

Exercises¶

1. (Conceptual) The chapter states that cross-entropy loss for a language model is exactly the negative log-likelihood, and that in language modeling the true distribution \(p\) is one-hot. Explain in your own words why, under a one-hot target on the observed token \(x^*\), the cross-entropy \(H(p,q) = -\sum_x p(x)\log q(x)\) collapses to \(-\log q(x^*)\), and what the decomposition \(H(p,q)=H(p)+D_{\text{KL}}(p\|q)\) becomes in this one-hot case.

2. (Quantitative — entropy and perplexity) A toy language model, at a single prediction step, outputs the categorical distribution \(q = [0.5,\ 0.25,\ 0.125,\ 0.125]\) over a 4-token vocabulary.

(a) Compute the entropy \(H(q)\) in bits. (b) If the true next token is token 1 (the one with probability \(0.25\)), what is the single-step cross-entropy loss \(-\log q(x^*)\), in bits and in nats? © What single-step perplexity does that loss correspond to?

3. (Quantitative — KL divergence and its asymmetry) Let \(p = [0.5,\ 0.5]\) and \(q = [0.9,\ 0.1]\) be two distributions over a binary outcome. Using natural logs (nats):

(a) Compute \(D_{\text{KL}}(p\|q)\). (b) Compute \(D_{\text{KL}}(q\|p)\). © Are they equal? Use the numbers to state concretely which direction penalizes putting mass where the other distribution has little.

4. (Conceptual — perplexity and tokenization) Two language models are evaluated on the same held-out English text. Model A uses a word-level tokenizer; Model B uses a byte-level tokenizer. Model B reports a lower per-token perplexity than Model A. A colleague concludes Model B is the better language model. Explain, using the chapter’s discussion of perplexity pitfalls, why this comparison is invalid, and describe what would make it a fair comparison.

5. (Implementation) The chapter’s perplexity.py exposes perplexity_from_log_probs and perplexity_from_logits, but no function that computes perplexity directly from a probability distribution and target ids without going through logits. Write a function perplexity_from_probs(probs, targets) in the chapter’s style: probs is a (T, V) tensor of already-normalized categorical distributions (each row sums to 1), targets is a (T,) long tensor of correct token ids. It should gather the probability of the correct token at each step, take the mean negative log, and exponentiate. Guard against \(\log 0\) with a small clamp, consistent with the chapter’s clamp(min=1e-12) idiom. Then show the call that reproduces the chapter’s toy 5-token result (~5.17).

import math
import torch

def perplexity_from_probs(
    probs: torch.Tensor,    # shape (T, V), each row sums to 1
    targets: torch.Tensor,  # shape (T,), ground-truth token ids
) -> float:
    """
    Compute perplexity directly from a categorical probability
    distribution (already softmaxed) and target token ids.

    Steps:
      1. Gather p(x_t | x_{<t}) for the correct token at each step.
      2. Take the natural log (clamped to avoid log(0)).
      3. Average the negative log-probs and exponentiate.
    """
    if probs.dim() != 2:
        raise ValueError("probs must be 2-D (T, V)")
    # Gather the probability assigned to each correct token: (T,)
    token_probs = probs.gather(1, targets.unsqueeze(1)).squeeze(1)
    # Clamp before log for numerical safety, matching the chapter's idiom.
    token_log_probs = torch.log(token_probs.clamp(min=1e-12))
    avg_nll = -token_log_probs.mean().item()
    return math.exp(avg_nll)

probs_correct = [0.20, 0.05, 0.30, 0.60, 0.15]
T, V = 5, 4
probs = torch.zeros(T, V)
targets = torch.tensor([0, 1, 2, 3, 0])  # any valid ids; one per row
for t, pc in enumerate(probs_correct):
    probs[t, targets[t]] = pc
    # distribute the rest arbitrarily over other columns (rows sum to 1)
    rest = (1.0 - pc) / (V - 1)
    for v in range(V):
        if v != targets[t].item():
            probs[t, v] = rest

print(perplexity_from_probs(probs, targets))  # -> ~5.173

6. (Implementation + conceptual — label smoothing) The chapter gives the label-smoothed loss

$$ \mathcal{L}_{\text{LS}} = -(1-\varepsilon)\log q(k^*) - \frac{\varepsilon}{V}\sum_k \log q(k). $$

(a) Write a small function smoothed_target(k_star, V, eps) returning the smoothed target distribution \(p_{\text{smooth}}\) as a length-\(V\) list, per the chapter’s formula \(p_{\text{smooth}}(k) = (1-\varepsilon)\mathbf{1}[k=k^*] + \varepsilon/V\). (b) By hand, for \(V = 4\), \(k^* = 0\), \(\varepsilon = 0.1\): write out \(p_{\text{smooth}}\) and verify it sums to 1. © Explain, using the chapter’s “KL guard” reasoning, why \(p_{\text{smooth}}(k) > 0\) for every \(k\) matters for keeping \(D_{\text{KL}}(p\|q)\) well-defined.

def smoothed_target(k_star: int, V: int, eps: float = 0.1):
    """Smoothed target distribution per the chapter's label-smoothing formula."""
    base = eps / V
    p = [base] * V
    p[k_star] += (1.0 - eps)   # add the (1-eps) mass onto the correct class
    return p