4. Probability and information

Why this layer exists

The last thing a language model computes is not a token. It is a vector of $V$ real numbers, one per vocabulary entry, called logits, which a softmax turns into a probability distribution over the next token. Everything downstream consumes that distribution. Generation draws a token from it, after reshaping it with temperature, top-k or top-p (chapter 12). Training compares it with the token that actually came next and adjusts the weights to raise that token’s probability; the quantity minimised is a cross-entropy (chapter 8). The number reported when models are compared on held-out text is perplexity, the same cross-entropy passed through an exponential, and it is how chapter 14 measures what quantization costs. When you send "logprobs": true to llama-server’s OpenAI-compatible /v1/chat/completions, the values that come back are natural logs of entries of that distribution.

KL divergence measures how far one distribution is from another, and it appears wherever one model is trained against another: the penalty that keeps a fine-tuned model near its starting point in RLHF-style training, the objective DPO is derived from (chapter 13), and distillation, where a small model is trained to match a large one’s outputs. This chapter builds all of these from the definition of a distribution, and the lab fits a first language model to data and measures its perplexity.

Mechanism

Distributions

A sample space is the set of things that can happen; for next-token prediction it is the vocabulary, $V$ token ids, and every sample space in this chapter is finite. A probability distribution over it assigns each outcome $i$ a number $p_i$ with

$p_i \ge 0 \quad \text{for every } i , \qquad \sum_{i=1}^{V} p_i = 1 .$

As a vector, $\mathbf{p}$ has shape (V,): a vector in the sense of chapter 2, restricted to non-negative entries that sum to 1. A distribution over a finite set with one free probability per outcome is a categorical distribution, and it is what a language model outputs at each position.

Two variables $X$ and $Y$ have a joint distribution $p(x, y)$, the probability that $X = x$ and $Y = y$ together. Summing out one variable gives the marginal of the other, $p(x) = \sum_y p(x, y)$. The conditional distribution of $Y$ given $X = x$ restricts to the cases where $X = x$ and renormalises:

$p(y \mid x) = \frac{p(x, y)}{p(x)} , \qquad p(x) > 0 .$

A language model is a conditional distribution, $p(\text{next token} \mid \text{context})$, one categorical per context. $X$ and $Y$ are independent if $p(x, y) = p(x)\,p(y)$ for every pair, equivalently $p(y \mid x) = p(y)$: knowing $X$ says nothing about $Y$. Exercise (e) measures how far co-occurring words are from independence.

Expectation and variance

A random variable attaches a number $f(i)$ to each outcome. Its expectation under $\mathbf{p}$ is the probability-weighted average

$\mathbb{E}_{\mathbf{p}}[f] = \sum_{i=1}^{V} p_i\, f(i) ,$

the dot product of $\mathbf{p}$ with the vector of values $(f(1), \ldots, f(V))$, both of shape (V,). Expectation is therefore linear, $\mathbb{E}[a f + b g] = a\,\mathbb{E}[f] + b\,\mathbb{E}[g]$, because the dot product is.

From here on, name a random variable by a capital letter, $X$, and work directly with the values $x$ it takes and their probabilities $p(x)$, so that $\mathbb{E}[X] = \sum_x p(x)\, x$: the same weighted average, summed over values instead of outcomes. The variance measures spread around the mean $\mu = \mathbb{E}[X]$. Expanding the square and using linearity,

$\mathrm{Var}[X] = \mathbb{E}\big[(X - \mu)^2\big] = \mathbb{E}[X^2] - 2\mu^2 + \mu^2 = \mathbb{E}[X^2] - \mu^2 .$

For independent $X$ and $Y$ the expectation of the product factors, because the joint probability does:

$\mathbb{E}[XY] = \sum_{x}\sum_{y} p(x)\,p(y)\, x y = \Big(\sum_x p(x)\, x\Big)\Big(\sum_y p(y)\, y\Big) = \mathbb{E}[X]\,\mathbb{E}[Y] .$

Expanding $\mathrm{Var}[X + Y] = \mathbb{E}[(X + Y)^2] - (\mathbb{E}[X] + \mathbb{E}[Y])^2$ gives $\mathrm{Var}[X] + \mathrm{Var}[Y] + 2(\mathbb{E}[XY] - \mathbb{E}[X]\,\mathbb{E}[Y])$, and the last term vanishes for independent variables: variances of independent terms add. Two consequences come back later. A sum of $n$ independent terms of variance $\sigma^2$ has variance $n\sigma^2$, so its standard deviation grows as $\sqrt{n}$, not $n$; chapter 9 uses this to set a layer’s initial weights. And since scaling a variable by $1/n$ scales its variance by $1/n^2$, the average of $n$ independent draws has variance $\sigma^2 / n$: averages concentrate around the expectation, with spread shrinking as $1/\sqrt{n}$. That is why an average over a held-out set estimates the expectations the definitions below are written in.

Logits and softmax

A network’s last layer is an affine map (chapter 3), and its output $\mathbf{z}$, shape (V,), can hold any real numbers. Exponentiating makes every entry positive and dividing by the total makes them sum to 1. That is the softmax:

$\mathrm{softmax}(\mathbf{z})_i = \frac{e^{z_i}}{\sum_{j=1}^{V} e^{z_j}} ,$

shape (V,) in and out. Every output is strictly positive, and since $e^x$ is increasing, softmax preserves the order of its inputs. Taking logs,

$\log \mathrm{softmax}(\mathbf{z})_i = z_i - \log \sum_{j=1}^{V} e^{z_j} ,$

so the logits are log-probabilities up to one constant shared by every entry. That constant does not matter. Subtract any $c$ from every logit; since $e^{z_i - c} = e^{-c} e^{z_i}$,

$\frac{e^{z_i - c}}{\sum_j e^{z_j - c}} = \frac{e^{-c}\, e^{z_i}}{e^{-c} \sum_j e^{z_j}} = \frac{e^{z_i}}{\sum_j e^{z_j}} ,$

because the common factor comes out of the sum and cancels. Softmax sees only differences between logits.

That invariance is what makes softmax computable. The largest float64 is about $1.8 \times 10^{308}$, whose natural log is 709.78, so $e^{z}$ overflows to inf for $z > 709.78$; the limit is 88.72 in float32 and 11.09 in float16. Logits of 1000, 1001 and 1002 have a perfectly good softmax, (0.090, 0.245, 0.665), but computed as written every exponential is inf, inf / inf is nan, and numpy returns [nan nan nan]. Choose $c = \max_j z_j$. After the shift the largest logit is 0, so the largest exponential is 1 and none can overflow; the sum contains that 1, so the division is safe. Shifted logits below about $-745$ still underflow: their exponentials, below about $e^{-744}$, round to 0. As a probability that is 0 for any practical purpose; its log is not, and the walkthrough shows how to keep it. For logits of shape (B, T, V) the max and the sum are taken along the last axis, giving shape (B, T, 1), which broadcasts back against (B, T, V) by the rules of chapter 3.

Surprise

How surprising is an outcome of probability $p$? Call the answer $h(p)$ and ask three things of it: it depends only on $p$; it decreases as $p$ grows, with a certain event carrying none; and for independent events, whose joint probability is a product, surprises add:

$h(pq) = h(p) + h(q) .$

$h(p) = -\log p$ satisfies all three, and it is the only choice up to the base. Write $g(x) = h(e^{-x})$ for $x \ge 0$; the product rule becomes $g(x + y) = g(x) + g(y)$. Applied repeatedly it gives $g(nx) = n\,g(x)$ for positive integers $n$, and with $x = 1/m$, $g(n/m) = (n/m)\,g(1)$ for every positive rational. A monotone function that matches the line $g(1)\,x$ at every rational matches it everywhere, since every real is squeezed between rationals on both sides. So $h(p) = -g(1)\log p$, and the constant $g(1)$ only sets the unit.

$-\log p$ is the surprisal (or information content) of the outcome. With $\log_2$ the unit is the bit: a fair coin flip carries $-\log_2 \frac{1}{2} = 1$ bit. With the natural log it is the nat, and the coin carries $\log 2 \approx 0.693$ nats. Converting is a constant factor, $\log_2 p = \log p / \log 2$, so 1 nat is $1/\log 2 \approx 1.4427$ bits. In this book $\log$ without a base is the natural log, the one np.log computes and training losses are reported in. An outcome of probability 0.01 carries 4.605 nats, or 6.644 bits.

Entropy

The entropy of a distribution is its expected surprise:

$H(\mathbf{p}) = \mathbb{E}_{\mathbf{p}}[-\log p_i] = -\sum_{i=1}^{V} p_i \log p_i ,$

a scalar for $\mathbf{p}$ of shape (V,). A term with $p_i = 0$ is taken as 0, which is the limit: with $p = e^{-x}$, $p \log p = -x\,e^{-x}$, and the exponential beats the linear factor as $x \to \infty$.

Entropy is 0 exactly when one outcome has probability 1: every term is then $1 \log 1$ or $0 \log 0$, both 0, and otherwise some $p_i$ lies strictly between 0 and 1 and contributes a positive term. The uniform distribution over $V$ outcomes has entropy $-\sum_i \frac{1}{V}\log\frac{1}{V} = \log V$, and the KL section proves no distribution on $V$ outcomes has more. For two outcomes with probabilities $p$ and $1 - p$ the entropy is a curve:

It peaks at $p = 1/2$ with $\log 2 \approx 0.693$ nats (1 bit) and falls steeply near the ends: a coin that lands heads 90% of the time still carries 0.325 nats per flip, nearly half the maximum, while one at 99% carries 0.056.

Entropy’s meaning is compression. For symbols drawn independently from $\mathbf{p}$, no uniquely decodable code averages fewer than $H(\mathbf{p})$ bits per symbol (entropy in bits); for such codes this lower bound follows from McMillan’s inequality (1956). Shannon’s source coding theorem (“A Mathematical Theory of Communication”, 1948) supplies the other direction: coding one symbol at a time can get within 1 bit of $H(\mathbf{p})$, and coding long blocks of symbols together gets arbitrarily close per symbol. The ideal codeword for outcome $i$ has $-\log_2 p_i$ bits. When those are integers the ideal is exact: four symbols with probabilities $(\frac12, \frac14, \frac18, \frac18)$ get the prefix-free codewords 0, 10, 110, 111, with average length

$\frac12 \cdot 1 + \frac14 \cdot 2 + \frac18 \cdot 3 + \frac18 \cdot 3 = 1.75 \text{ bits} = H(\mathbf{p}) ,$

against 2 bits for a fixed-width code.

Cross-entropy

Now code data from $\mathbf{p}$ with codeword lengths $-\log q_i$ designed for a different distribution $\mathbf{q}$. The expected length is the cross-entropy

$H(\mathbf{p}, \mathbf{q}) = \mathbb{E}_{\mathbf{p}}[-\log q_i] = -\sum_{i=1}^{V} p_i \log q_i ,$

with both vectors of shape (V,): the expectation is over the truth, the surprise measured by the model. The fixed-width code above is $\mathbf{q}$ uniform, $-\log_2 q_i = 2$, and $H(\mathbf{p}, \mathbf{q}) = 2$ bits against $H(\mathbf{p}) = 1.75$. Cross-entropy is infinite if $q_i = 0$ for an outcome with $p_i > 0$: an outcome with no codeword cannot be sent at any length.

This is a language model’s training loss. At one position the “true distribution” is the token that actually came next, $t$: a one-hot $\mathbf{p}$ with $p_t = 1$. The sum has one nonzero term,

$H(\mathbf{p}, \mathbf{q}) = -\log q_t ,$

the surprisal the model assigned to what happened. Averaged over the positions of a batch, it is the loss minimised in chapter 8.

KL divergence and Gibbs’ inequality

The extra cost of coding $\mathbf{p}$-data with a $\mathbf{q}$-code is

$\mathrm{KL}(\mathbf{p} \Vert \mathbf{q}) = H(\mathbf{p}, \mathbf{q}) - H(\mathbf{p}) = -\sum_i p_i \log q_i + \sum_i p_i \log p_i = \sum_{i=1}^{V} p_i \log \frac{p_i}{q_i} ,$

the Kullback–Leibler divergence of $\mathbf{q}$ from $\mathbf{p}$, with terms where $p_i = 0$ counted as 0. In the four-symbol example it is $2 - 1.75 = 0.25$ bits.

Gibbs’ inequality: $\mathrm{KL}(\mathbf{p} \Vert \mathbf{q}) \ge 0$, with equality only when $\mathbf{p} = \mathbf{q}$. The proof needs one fact:

$\log x \le x - 1 \quad \text{for all } x > 0, \text{ with equality only at } x = 1 .$

Let $f(x) = x - 1 - \log x$. Its derivative $f'(x) = 1 - 1/x$ is negative for $x < 1$ and positive for $x > 1$, so $f$ has its minimum at $x = 1$, where $f(1) = 0$, and is positive everywhere else.

If $q_i = 0$ for some $i$ with $p_i > 0$, the divergence is $+\infty$ and there is nothing to prove. Otherwise let $S$ be the set of outcomes with $p_i > 0$, and apply the fact with $x = q_i / p_i$:

$-\mathrm{KL}(\mathbf{p} \Vert \mathbf{q}) = \sum_{i \in S} p_i \log \frac{q_i}{p_i} \le \sum_{i \in S} p_i \left( \frac{q_i}{p_i} - 1 \right) = \sum_{i \in S} q_i - \sum_{i \in S} p_i \le 1 - 1 = 0 .$

The last step uses $\sum_{i \in S} p_i = 1$, since $S$ holds all of $\mathbf{p}$‘s mass, and $\sum_{i \in S} q_i \le 1$. Equality needs both inequalities tight: the first forces $q_i = p_i$ on $S$, the second forces $\mathbf{q}$ to have no mass outside $S$, so $\mathbf{q} = \mathbf{p}$. Three consequences:

1. No code beats the truth: $H(\mathbf{p}, \mathbf{q}) \ge H(\mathbf{p})$. A model’s cross-entropy on data is bounded below by the entropy of the process that generated it, and reaches the bound only by matching the process.
2. Uniform has maximum entropy. For $\mathbf{u}$ uniform on $V$ outcomes, $\mathrm{KL}(\mathbf{p} \Vert \mathbf{u}) = \sum_i p_i \log (p_i V) = \log V - H(\mathbf{p}) \ge 0$.
3. KL is not a distance. It is not symmetric: for $\mathbf{p} = (0.1, 0.9)$ and $\mathbf{q} = (0.5, 0.5)$, $\mathrm{KL}(\mathbf{p} \Vert \mathbf{q}) = 0.368$ nats and $\mathrm{KL}(\mathbf{q} \Vert \mathbf{p}) = 0.511$. Exercise (c) shows what each order penalises.

Maximum likelihood

Usually the distribution is unknown and there is data instead: $N$ samples $x_1, \ldots, x_N$, assumed drawn independently from an unknown categorical. Which distribution $\theta$, shape (V,), best explains them? Its likelihood is the probability it gives the observed data. Independence makes that a product, and grouping equal outcomes turns it into powers of the counts $n_i$ (the number of samples equal to $i$, so $\sum_i n_i = N$):

$L(\theta) = \prod_{m=1}^{N} \theta_{x_m} = \prod_{i=1}^{V} \theta_i^{\,n_i} .$

Maximum likelihood picks the $\theta$ that makes the data most probable. Since the log is increasing, maximise the log-likelihood instead, which has the same maximiser and does not underflow:

$\ell(\theta) = \sum_{i=1}^{V} n_i \log \theta_i .$

The maximisation is constrained, $\sum_i \theta_i = 1$; without that, $\ell$ grows without bound. The tool for an equality constraint is a Lagrange multiplier, and chapter 2’s geometry is enough to build it.

Assume first that every $n_i > 0$. Then $\ell = -\infty$ if any $\theta_i = 0$, so a maximum has every $\theta_i > 0$. Move it by a small step $\delta$, shape (V,). To first order,

$\ell(\theta + \delta) - \ell(\theta) \approx \sum_{i=1}^{V} \frac{\partial \ell}{\partial \theta_i}\, \delta_i = \nabla \ell \cdot \delta ,$

where $\nabla \ell$, shape (V,), collects the partial derivatives $\partial \ell / \partial \theta_i = n_i / \theta_i$, each the ordinary derivative of $n_i \log \theta_i$ with the other coordinates held fixed (chapter 5 treats these properly). The step keeps the constraint only if $\sum_i \delta_i = 0$, that is, $\delta \cdot \mathbf{1} = 0$ for the all-ones vector $\mathbf{1}$: allowed steps are exactly those orthogonal to $\mathbf{1}$. At a maximum $\nabla \ell \cdot \delta = 0$ for every allowed step; if it were positive, a small step along $\delta$ would increase $\ell$, and if negative, a step along $-\delta$ would.

That forces $\nabla \ell$ to be a multiple of $\mathbf{1}$. Write $\nabla \ell = \lambda \mathbf{1} + \mathbf{r}$ with $\lambda$ the mean of its entries, so that $\mathbf{r} \cdot \mathbf{1} = 0$. Then $\mathbf{r}$ is an allowed step, and $0 = \nabla \ell \cdot \mathbf{r} = \lambda\,(\mathbf{1} \cdot \mathbf{r}) + \mathbf{r} \cdot \mathbf{r} = \lVert \mathbf{r} \rVert^2$, so $\mathbf{r} = \mathbf{0}$ and

$\nabla \ell = \lambda\, \mathbf{1} .$

The scalar $\lambda$ is the Lagrange multiplier: at a constrained maximum the gradient points straight out of the constraint surface, so no allowed move improves the objective. The usual packaging is $\mathcal{L}(\theta, \lambda) = \ell(\theta) - \lambda\,(\sum_i \theta_i - 1)$ with every partial derivative set to zero; $\partial \mathcal{L} / \partial \theta_i = 0$ is the condition above and $\partial \mathcal{L} / \partial \lambda = 0$ is the constraint.

Here the condition reads $n_i / \theta_i = \lambda$, so $\theta_i = n_i / \lambda$. Summing and using the constraint, $1 = N / \lambda$, so $\lambda = N$ and

$\hat{\theta}_i = \frac{n_i}{N} .$

The maximum-likelihood estimate of a categorical is the normalised counts. An outcome with $n_i = 0$ does not appear in $\ell$ at all, so probability placed on it is taken from outcomes that do appear, and the maximum puts 0 there, which the formula also gives. The multiplier argument locates the stationary point; Gibbs’ inequality shows it is the global maximum, and what maximum likelihood is doing.

Maximising likelihood is minimising cross-entropy

Let $\hat{\mathbf{p}}$, shape (V,), be the empirical distribution of the data, $\hat{p}_i = n_i / N$. Dividing the log-likelihood by $N$,

$\frac{1}{N}\,\ell(\theta) = \sum_{i=1}^{V} \frac{n_i}{N} \log \theta_i = \sum_{i=1}^{V} \hat{p}_i \log \theta_i = -H(\hat{\mathbf{p}}, \theta) .$

Maximising likelihood is minimising the cross-entropy from the empirical distribution to the model. Since $H(\hat{\mathbf{p}}, \theta) = H(\hat{\mathbf{p}}) + \mathrm{KL}(\hat{\mathbf{p}} \Vert \theta)$ and $H(\hat{\mathbf{p}})$ is fixed by the data, it is also minimising $\mathrm{KL}(\hat{\mathbf{p}} \Vert \theta)$, which Gibbs says is zero, its minimum, exactly at $\theta = \hat{\mathbf{p}}$.

The same identity holds for a model with context. The average of $-\log q(x_t \mid \text{context}_t)$ over the positions of a text is the negative log-likelihood (NLL) per token, and it is the cross-entropy loss of the previous section averaged over positions. A network cannot set each conditional freely, since all of them come out of the same weights, so there is no closed form and the minimisation is done by gradient descent (chapter 5). The objective is the same.

Zero probabilities and smoothing

Maximum likelihood trusts the sample completely. An outcome absent from $N$ draws gets $\hat{\theta}_i = 0$, and if it turns up in new data the likelihood of that data is 0 and its NLL infinite. With a large vocabulary and a small sample this is the normal case. Add-$\alpha$ smoothing adds a pseudo-count $\alpha > 0$ to every outcome before normalising:

$\theta_i = \frac{n_i + \alpha}{N + \alpha V} ,$

where the denominator is $\sum_i (n_i + \alpha)$. With $\alpha = 1$ this is Laplace’s rule. Every outcome gets at least $\alpha / (N + \alpha V) > 0$. The price is a pull toward uniform, which matters when $\alpha V$ is comparable to $N$ and vanishes when $N \gg \alpha V$; the lab shows both regimes. A softmax has the protection built in, since $e^{z_i} > 0$ for every real $z_i$: a neural language model never assigns exactly zero probability, except by floating-point underflow.

Perplexity

NLL in nats per token is hard to read. Its exponential is not:

$\mathrm{PPL} = e^{\mathrm{NLL}} = \exp\Big(-\frac{1}{N}\sum_{m=1}^{N} \log q_{x_m}\Big) = \Big(\prod_{m=1}^{N} q_{x_m}\Big)^{-1/N} ,$

using $e^{a \log b} = b^{a}$ and turning the exponential of a sum into a product. Perplexity is the reciprocal of the geometric mean of the probabilities the model gave the tokens that occurred. A uniform model over $K$ outcomes has NLL $\log K$ and perplexity exactly $K$, which fixes the reading: perplexity $K$ means the model is, on average in log space, as uncertain as a uniform choice among $K$ options, its effective branching factor. It is at least 1, reached only by a model that gives every occurring token probability 1.

Perplexity does not depend on the log base, since $2^{\text{NLL in bits}} = e^{\text{NLL in nats}}$; the NLL does. The trap is the unit of text: per-token NLL and perplexity are comparable only between models that share a tokenizer and a test set, because a tokenizer that splits the same text into more tokens spreads the same total surprise over more positions. Chapter 14 compares a model with its own quantized copy, which shares the tokenizer, so the comparison is sound.

Walkthrough

The reference module is py/tinygpt/prob.py. All logs are natural, so every function returns nats.

def softmax(logits, axis=-1):
    """exp(z_i) / sum_j exp(z_j) along axis, computed after subtracting the max for stability."""
    z = np.asarray(logits, dtype=np.float64)
    z = z - z.max(axis=axis, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=axis, keepdims=True)

keepdims=True keeps a size-1 axis where the max and sum were taken: for logits of shape (B, T, V) and axis=-1, both have shape (B, T, 1) and broadcast back across the last axis, one per row. Without it the max would have shape (B, T), which usually fails to broadcast against (B, T, V). The dangerous case is a square (V, V) input, where a (V,) max broadcasts as a row without complaint and subtracts row $j$‘s maximum from column $j$: a different constant for each entry of a row, which softmax is not invariant to.

A loss needs log-probabilities, and np.log(softmax(z)) loses them when a probability underflows: the logits (0, -800) have log-probabilities (0, -800) to within rounding, but their softmax is (1.0, 0.0) and its log is (0, -inf). Computing $z_i - \max_j z_j - \log \sum_j e^{z_j - \max_j z_j}$ directly, the mechanism’s log-softmax with the same shift, returns -800 exactly. Code that computes a loss from logits should use that form.

def entropy(p):
    """H(p) = -sum p_i log p_i in nats, with 0 log 0 taken as 0."""
    p = _distribution(p)
    nz = p > 0
    return float(-(p[nz] * np.log(p[nz])).sum())


def cross_entropy(p, q):
    """H(p, q) = -sum p_i log q_i: the average code length when data from p is coded for q."""
    p, q = _distribution(p), _distribution(q)
    support = p > 0
    if np.any(q[support] == 0):
        return math.inf
    return float(-(p[support] * np.log(q[support])).sum())


def kl(p, q):
    """KL(p || q) = H(p, q) - H(p): the extra nats paid for using q when the truth is p."""
    return cross_entropy(p, q) - entropy(p)

entropy and cross_entropy pass their arguments through _distribution, a private helper just above the listing, which converts to float64 and raises ValueError unless the input is one-dimensional, has no negative entry, and sums to 1 within $10^{-9}$. The tolerance allows for rounding: a computed softmax need not sum to exactly 1. The check matters because none of these formulas fails loudly on a non-distribution: $-\sum p_i \log p_i$ of [0.5, 0.6] is a finite number, just not an entropy.

The mask nz = p > 0 implements $0 \log 0 = 0$; without it, 0.0 * np.log(0.0) is 0.0 * -inf, which is nan, and one impossible outcome would poison the sum. cross_entropy restricts both vectors to the support of $\mathbf{p}$ for the same reason, and returns math.inf explicitly when $\mathbf{q}$ is zero anywhere on it. kl is cross-entropy minus entropy, the first form in the derivation, so the infinite case passes through: inf - H is inf. When the divergence is tiny compared with the entropies, the direct form $\sum_i p_i \log(p_i / q_i)$ is more accurate, since subtracting two nearly equal floats keeps only the digits in which they differ; the tests compare kl(p, p) with zero to a tolerance for that reason.

def fit_categorical(samples, vocab_size, alpha=0.0):
    """Maximum-likelihood (alpha=0) or add-alpha smoothed estimate of a categorical distribution."""
    counts = np.bincount(np.asarray(samples, dtype=np.int64), minlength=vocab_size).astype(np.float64)
    counts += alpha
    return counts / counts.sum()


def nll(probs, samples):
    """Mean negative log-likelihood of samples under probs, in nats per sample."""
    probs = np.asarray(probs, dtype=np.float64)
    p = probs[np.asarray(samples, dtype=np.int64)]
    if np.any(p == 0):
        return math.inf
    return float(-np.log(p).mean())


def perplexity(probs, samples):
    """exp(NLL): the effective number of equally likely choices the model is confused between."""
    return math.exp(nll(probs, samples))

fit_categorical is the maximum-likelihood result plus the pseudo-count. np.bincount counts occurrences of each integer, and minlength=vocab_size extends the result to the whole vocabulary, so unseen ids get count 0. Samples must be ids in [0, vocab_size); bincount silently lengthens its output for a larger id rather than raising. Adding alpha everywhere and dividing by the new total is the smoothing formula.

nll gathers each sample’s probability with one indexing operation, probs[samples], shape (N,), and checks for a zero before taking any log, so an impossible sample returns math.inf without a divide-by-zero warning. perplexity is math.exp of that, and math.exp(math.inf) is inf, so the zero case needs no special code.

Exercises

(a) What is the entropy of a fair six-sided die, in bits and in nats?

(b) Show that $\mathrm{softmax}(\mathbf{z} + c\mathbf{1}) = \mathrm{softmax}(\mathbf{z})$ for any scalar $c$, but $\mathrm{softmax}(s\,\mathbf{z}) \ne \mathrm{softmax}(\mathbf{z})$ in general for a scalar $s > 0$. What does scaling do as $s \to \infty$ and as $s \to 0$?

s = 0.5   (0.1863, 0.3072, 0.5065)
s = 1     (0.0900, 0.2447, 0.6652)
s = 2     (0.0159, 0.1173, 0.8668)
s = 10    (0.0000, 0.0000, 1.0000)   to four places

(c) Let $\mathbf{p} = (0.5, 0.5)$ and $\mathbf{q} = (0.99, 0.01)$. Compute $\mathrm{KL}(\mathbf{p} \Vert \mathbf{q})$ and $\mathrm{KL}(\mathbf{q} \Vert \mathbf{p})$, and explain why they differ.

(d) A coin lands heads with unknown probability $\theta$. In $N$ flips you see $k$ heads. Derive the maximum-likelihood $\theta$ by differentiating the log-likelihood.

(e) Chapter 2 found that raw co-occurrence counts make cat and king look similar (cosine 0.878) because both rows are dominated by the columns for the, a, . and the adverbs. The pointwise mutual information of a word $w$ and a context word $c$,

$\mathrm{PMI}(w, c) = \log \frac{p(w, c)}{p(w)\,p(c)} ,$

is the log of how much more often they co-occur than they would if independent, and positive PMI clips it at zero: $\mathrm{PPMI}(w, c) = \max(\mathrm{PMI}(w, c), 0)$. Compute PPMI from chapter 2’s co-occurrence matrix of py/data/tiny-corpus.txt (window 2) and compare the nearest results with those from the raw counts.

import pathlib

import numpy as np

from tinygpt.vectors import cooccurrence, nearest

tokens = pathlib.Path("data/tiny-corpus.txt").read_text().split()
vocab, counts = cooccurrence(tokens, window=2)
joint = counts / counts.sum()                  # (V, V): p(w, c)
pw = joint.sum(axis=1, keepdims=True)          # (V, 1): p(w)
pc = joint.sum(axis=0, keepdims=True)          # (1, V): p(c)
with np.errstate(divide="ignore"):
    pmi = np.log(joint / (pw * pc))            # (V, V); log 0 = -inf
ppmi = np.maximum(pmi, 0.0)
for name, m in [("counts", counts), ("ppmi", ppmi)]:
    print(name)
    for word in ["king", "cat", "quickly", "the"]:
        hits = nearest(m, vocab, word, k=3)
        print(f"  {word:8}", "  ".join(f"{w} {s:.3f}" for w, s in hits))

counts
  king     queen 1.000  prince 0.999  dog 0.879
  cat      dog 0.999  bird 0.999  queen 0.879
  quickly  quietly 0.998  . 0.935  queen 0.726
  the      a 0.999  bird 0.557  king 0.556
ppmi
  king     queen 0.999  prince 0.998  quickly 0.721
  cat      dog 0.998  bird 0.998  quietly 0.734
  quickly  quietly 0.968  prince 0.725  queen 0.725
  the      a 0.968  bites 0.658  chases 0.640

(f) A model has perplexity 20 on held-out text, with a vocabulary of 50,000 tokens. What does that mean operationally?

Lab

Implement every function in py/labs/ch04/starter.py, then run

make lab CH=04 IMPL=mine

until it passes. Read py/labs/ch04/test_lab.py first. It pins what the prose leaves open: softmax must give the same result for logits (1, 2, 3) and (1001, 1002, 1003) with no inf or nan, and must work along the last axis of a matrix; entropy must raise ValueError for a vector that does not sum to 1 and for one with a negative entry even when it does; cross_entropy must return math.inf when $\mathbf{q}$ misses part of $\mathbf{p}$‘s support; fit_categorical must give $(3/8, 2/8, 2/8, 1/8)$ for samples [0, 0, 1, 2] with $V = 4$ and $\alpha = 1$; and perplexity must be $V$, to within rounding, for a uniform model and inf for a sample the model gives probability 0.

Then fit the simplest language model that learns anything from data: one categorical distribution over words, the unigram model, which predicts every token the same way regardless of context. Save this script outside the repository, for example as /tmp/unigram.py:

import math
import pathlib

from tinygpt.prob import entropy, fit_categorical, perplexity

tokens = pathlib.Path("data/tiny-corpus.txt").read_text().split()
vocab = sorted(set(tokens))
ids = [vocab.index(t) for t in tokens]
cut = int(0.9 * len(ids))
train, test = ids[:cut], ids[cut:]
print(len(train), "training tokens,", len(test), "held-out tokens,", len(vocab), "words")
for n in [len(train), 40]:
    for alpha in [0.0, 1.0]:
        probs = fit_categorical(train[:n], len(vocab), alpha=alpha)
        h = entropy(probs)
        print(
            f"train {n:5d}  alpha={alpha:.0f}  entropy {h:.4f} nats ({h / math.log(2):.4f} bits)"
            f"  held-out perplexity {perplexity(probs, test):.4f}"
        )

and run it from py/ with PYTHONPATH=. uv run python /tmp/unigram.py. The first 90% of the token stream is the training set and the last 10% is held out. The vocabulary comes from the whole corpus so every held-out word has an id; in this corpus every held-out word also occurs in training. The script fits on all the training tokens and then on only the first 40, each with $\alpha = 0$ and $\alpha = 1$. Change the import to from labs.ch04.starter import entropy, fit_categorical, perplexity to run yours; the output must be identical:

11998 training tokens, 1334 held-out tokens, 17 words
train 11998  alpha=0  entropy 2.6045 nats (3.7576 bits)  held-out perplexity 13.5865
train 11998  alpha=1  entropy 2.6052 nats (3.7585 bits)  held-out perplexity 13.5864
train    40  alpha=0  entropy 2.4329 nats (3.5099 bits)  held-out perplexity inf
train    40  alpha=1  entropy 2.6578 nats (3.8344 bits)  held-out perplexity 14.9375

Explain each pair of lines before moving on.

• With all the training data, smoothing makes no visible difference. Every word occurs at least 275 times in the 11,998 training tokens (summons is the rarest). One pseudo-count per word adds 17 to a total of 11,998 and moves the held-out perplexity in the fourth decimal place: the $N \gg \alpha V$ regime.
• The numbers agree with the theory. The entropy, 2.6045 nats, is below the 17-word maximum $\log 17 = 2.833$ because a, . and the account for 45% of the training tokens. Gibbs’ inequality bounds the held-out score from below by the held-out text’s own entropy: the held-out tokens’ empirical distribution has entropy 2.5968 nats, perplexity 13.42, and no unigram model can score better on them. The fitted model scores 13.59, a cross-entropy of 2.6091 nats; the 0.0122-nat gap is $\mathrm{KL}(\hat{\mathbf{p}}_{\text{test}} \Vert \theta)$, the cost of the held-out frequencies differing slightly from the training ones. A uniform model would score 17. Word frequency alone cuts the effective branching factor from 17 to 13.6.
• With 40 training tokens, maximum likelihood fails outright. The first 40 tokens contain 14 of the 17 words; chases, quietly and sniffs are missing. With $\alpha = 0$ they get probability 0, they occur in the held-out set, and the perplexity is infinite. With $\alpha = 1$, 17 pseudo-counts against 40 real tokens pull hard toward uniform, each missing word gets $1/57$, and the perplexity is 14.94: worse than the full-data model, better than uniform, finite.
• The small unsmoothed model is also overconfident. Its entropy, 2.4329 nats, is below the full-data value: this small sample makes the distribution look more concentrated than it is, and maximum likelihood takes it at its word. Smoothing pushes the entropy the other way, to 2.6578, because it pulls toward uniform.

The language models of chapters 8, 9, 11 and 14 are scored the same way: fit on one part of the data, report perplexity on another. Chapter 8’s bigram model conditions on the previous token and should beat 13.59 on this corpus by a wide margin (an add-one bigram fitted on the same split scores about 3.9), because the grammar that generated the corpus makes most next words predictable from the one before.

Further reading

• Claude E. Shannon, “A Mathematical Theory of Communication” (1948). Entropy, the source coding theorem, and statistical approximations to English text built from letter and word frequencies.
• Thomas M. Cover and Joy A. Thomas, Elements of Information Theory (2nd ed., 2006), chapter 2. Entropy, relative entropy, and the inequalities between them, with proofs.
• David J. C. MacKay, Information Theory, Inference, and Learning Algorithms (Cambridge University Press, 2003), chapters 2 and 4. Probability, entropy and the source coding theorem, with many worked exercises.
• Kenneth Ward Church and Patrick Hanks, “Word Association Norms, Mutual Information, and Lexicography” (1990). Pointwise mutual information applied to word co-occurrence in a corpus, the basis of the reweighting in exercise (e).