A language model is a probability distribution. Given the text so far, it assigns a probability to every possible next token, and generating text means drawing from that distribution, one token at a time. Training it means choosing parameters that make the real text in the training set as probable as possible. Evaluating it means asking how probable it finds text it has never seen.

So probability is not background material for modern AI. It is the language the models are written in. This section covers the parts you will use constantly (expectation and variance, the three distributions that do most of the work, Bayes rule, maximum likelihood, priors as regularizers, and sampling), and it ends with a numerical trick, the log-sum-exp, that every loss function in this book depends on.

Random variables, expectation, and variance

A random variable is a quantity whose value is uncertain, described by a distribution. A discrete one, like the next token, has a probability mass function $p(x) = \mathbb{P}(X = x)$ that sums to 1. A continuous one, like a weight drawn at initialization, has a density $p(x)$ that integrates to 1. (A density can exceed 1. Only its integral over an interval is a probability.)

The expectation is the probability-weighted average,

In plain words, the expectation is the long-run average value, and to average a function of $X$ you weight the function's values by the probabilities, with no need to work out the distribution of $f(X)$ first.

The most useful property is linearity. $\mathbb{E}[aX + bY] = a\,\mathbb{E}[X] + b\,\mathbb{E}[Y]$ always, whether or not $X$ and $Y$ are independent. Most expectation calculations in machine learning are linearity plus bookkeeping.

The variance measures spread around the mean, $\operatorname{Var}(X) = \mathbb{E}[(X - \mathbb{E}X)^2] = \mathbb{E}[X^2] - (\mathbb{E}X)^2$, and its square root is the standard deviation. Three rules carry most of the weight.

• $\operatorname{Var}(aX + b) = a^2 \operatorname{Var}(X)$. Shifting does nothing, and scaling by $a$ scales the variance by $a^2$.
• If $X$ and $Y$ are independent, $\operatorname{Var}(X + Y) = \operatorname{Var}(X) + \operatorname{Var}(Y)$. In general there is an extra $2\operatorname{Cov}(X, Y)$, where $\operatorname{Cov}(X, Y) = \mathbb{E}[(X - \mathbb{E}X)(Y - \mathbb{E}Y)]$.
• If $X$ and $Y$ are independent with mean zero, $\operatorname{Var}(XY) = \mathbb{E}[X^2 Y^2] = \mathbb{E}[X^2]\,\mathbb{E}[Y^2] = \operatorname{Var}(X)\operatorname{Var}(Y)$.

Worked example, the variance argument behind attention

Here is a calculation you will use twice in this book. Draw two vectors $\mathbf{q}, \mathbf{k} \in \mathbb{R}^d$ with independent entries of mean 0 and variance 1. What is the variance of their dot product $\mathbf{q}^\top\mathbf{k} = \sum_{j=1}^d q_j k_j$?

Each term $q_j k_j$ has mean $\mathbb{E}[q_j]\mathbb{E}[k_j] = 0$ and, by the third rule, variance $1 \cdot 1 = 1$. The $d$ terms are independent, so by the second rule the variances add, and $\operatorname{Var}(\mathbf{q}^\top\mathbf{k}) = d$. The typical size of a dot product is therefore $\sqrt{d}$, which for $d = 512$ is about 23. Dividing by $\sqrt{d}$ brings the variance back to 1, by the first rule. That division is exactly the scaling in scaled dot-product attention (§4.3), where unscaled scores would push the softmax into saturation. The same three rules, applied to a layer's pre-activations, give the Xavier and He initializations of §0.7.

Averages get more precise, at a rate of one over root n

Average $n$ independent copies of $X$, each with variance $\sigma^2$. By the rules above, the average has variance $n\sigma^2 / n^2 = \sigma^2/n$, so its standard deviation, the standard error, is $\sigma/\sqrt{n}$. The law of large numbers says the average converges to $\mathbb{E}[X]$, and the central limit theorem says its fluctuations around the mean look Gaussian with that standard deviation, whatever the distribution of $X$ (as long as its variance is finite).

That $1/\sqrt n$ rate shows up all over machine learning. A minibatch gradient over $B$ examples has noise that shrinks like $1/\sqrt{B}$ (§0.6). An accuracy measured on 1,000 test questions has a standard error of about $\sqrt{0.5 \cdot 0.5/1000} \approx 1.6$ percentage points near 50%, so a 1-point gap between two models on that benchmark is noise. Confidence intervals, paired comparisons between two models, and how many test questions you need are developed in §0.11, and applied to model evaluations in §14.7.

The distributions that do most of the work

Bernoulli. A single yes-or-no outcome, $X \in \{0, 1\}$, with $\mathbb{P}(X = 1) = p$. Its mean is $p$ and its variance is $p(1 - p)$. Binary classifiers, reward models comparing two answers, and dropout masks are all Bernoulli. A classifier usually outputs a logit $z$ and sets $p = \sigma(z) = 1/(1 + e^{-z})$, where in this section $\sigma(\cdot)$ with an argument is the sigmoid and a bare $\sigma$ is a standard deviation.

Categorical. One of $K$ outcomes, with probabilities $\mathbf{p} = (p_1, \dots, p_K)$ summing to 1. Every softmax output is a categorical distribution, and every language model's next-token prediction is one with $K = V$, the vocabulary size. If you encode the outcome as a one-hot vector $\mathbf{e}_X$, its mean is $\mathbf{p}$ and its covariance is $\operatorname{diag}(\mathbf{p}) - \mathbf{p}\mathbf{p}^\top$. That is exactly the softmax Jacobian from §0.2, which is no coincidence. For every distribution in this family (the exponential families), the derivative of the mean with respect to the natural parameters, which for the categorical are the logits, is the covariance.

Gaussian (normal). A continuous value with mean $\mu$ and variance $\sigma^2$, written $\mathcal{N}(\mu, \sigma^2)$, with density

In plain words, the log-density is a downward parabola centered at the mean, and its curvature is set by the variance. Gaussians appear for three reasons. The central limit theorem makes sums of many small effects Gaussian. They are the maximum-entropy distribution for a given mean and variance (§0.5), the "least assuming" choice. And they are algebraically convenient, since sums and linear maps of Gaussians stay Gaussian.

Multivariate Gaussian. A vector $\mathbf{x} \in \mathbb{R}^d$ with mean $\boldsymbol{\mu} \in \mathbb{R}^d$ and covariance matrix $\boldsymbol{\Sigma} \in \mathbb{R}^{d \times d}$ (symmetric and positive definite), with density

The quadratic form in the exponent is the squared Mahalanobis distance, a distance that accounts for correlations and different scales. Its level sets are ellipsoids (ellipses in 2-D) whose axes point along the eigenvectors of $\boldsymbol{\Sigma}$, with half-lengths proportional to the square roots of its eigenvalues. (Eigenvectors of symmetric matrices, and what positive definite means geometrically, are the subject of §0.10.) To sample, factor $\boldsymbol{\Sigma} = \mathbf{L}\mathbf{L}^\top$ with the Cholesky decomposition ($\mathbf{L}$ lower triangular), draw $\mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$, and set $\mathbf{x} = \boldsymbol{\mu} + \mathbf{L}\mathbf{z}$. Its covariance is $\mathbb{E}[\mathbf{L}\mathbf{z}\mathbf{z}^\top\mathbf{L}^\top] = \mathbf{L}\mathbf{I}\mathbf{L}^\top = \boldsymbol{\Sigma}$, as required.

Joint, marginal, and conditional probability

Most questions worth asking involve two or more random variables at once, such as a word and the word that follows it. Their joint distribution $p(x, y) = \mathbb{P}(X = x, Y = y)$ gives the probability of every combination. Summing out one variable gives the marginal distribution of the other, $p(x) = \sum_y p(x, y)$. Dividing the joint by a marginal gives the conditional distribution, the probability of $x$ once you know $y$,

In plain words, to condition on $Y = y$, keep only the slice of the joint table where $Y = y$ and rescale it so it sums to 1. The second form, the product rule, builds a joint out of a marginal and a conditional.

A small table makes this concrete. Let $X$ be whether it rained and $Y$ whether the grass is wet, with joint probabilities $p(\text{rain}, \text{wet}) = 0.30$, $p(\text{rain}, \text{dry}) = 0.02$, $p(\text{none}, \text{wet}) = 0.08$, and $p(\text{none}, \text{dry}) = 0.60$ (sprinklers explain the 0.08).

Rain makes wet grass very likely, $p(\text{wet} \mid \text{rain}) = 0.30/0.32 = 0.94$. Wet grass makes rain likely but less so, $p(\text{rain} \mid \text{wet}) = 0.30/0.38 = 0.79$, because the sprinklers also wet the grass. The two conditionals are different numbers, and confusing them is one of the most common errors in reasoning about evidence. Bayes rule, below, is the bridge between them. We will meet this same table again in §0.5, where we measure how much the grass tells us about the rain.

Two variables are independent when the joint factors into its marginals, $p(x, y) = p(x)\,p(y)$, which is the same as $p(x \mid y) = p(x)$ for every $y$. Knowing one tells you nothing about the other. Rain and wet grass are not independent ($0.32 \times 0.38 = 0.12$, not 0.30). Independence is the assumption that makes the variance rules above work, and it is what turns the probability of a whole data set into a product, as we will use for maximum likelihood.

Applying the product rule over and over factors the joint distribution of any sequence,

where $x_{<t}$ means all the elements before position $t$. This is the chain rule of probability, and it holds for any distribution, with no assumptions at all. An autoregressive language model is this identity with a neural network supplying each factor, predicting every token from the ones before it (§3.1).

Bayes rule

Bayes rule turns a model of how data arise from causes into a belief about the causes given the data. It follows in one line from the product rule, because the joint can be factored either way, $\mathbb{P}(H, D) = \mathbb{P}(D \mid H)\,\mathbb{P}(H) = \mathbb{P}(H \mid D)\,\mathbb{P}(D)$. Divide by $\mathbb{P}(D)$. For a hypothesis $H$ and observed data $D$,

In plain words, the posterior is the prior times the likelihood, renormalized. The denominator is the marginal probability of the data, the joint summed over every hypothesis (sometimes called the law of total probability). The names matter. $\mathbb{P}(H)$ is the prior (belief before the data), $\mathbb{P}(D \mid H)$ is the likelihood (how well the hypothesis predicts the data), and $\mathbb{P}(H \mid D)$ is the posterior.

The classic worked problem shows why the prior cannot be ignored. A disease affects 1% of people. A test catches 99% of cases and has a 1% false-positive rate. You test positive. How likely are you to have the disease? The numerator is $0.99 \times 0.01 = 0.0099$. The denominator adds the false positives, $0.0099 + 0.01 \times 0.99 = 0.0198$. The posterior is $0.0099/0.0198 = 0.5$. A "99% accurate" test gives a coin flip, because the disease is rare, so the healthy people's false positives are as numerous as the true positives. Make the disease ten times rarer (0.1%) and the posterior drops to 9%.

This is the reasoning behind any detector of rare events, including the safety classifiers of §19.7. A classifier with a low false-positive rate can still produce mostly false alarms if what it detects is rare enough. In the language of §0.11, the posterior here is the classifier's precision, and it depends on the base rate in a way that recall and the false-positive rate do not.

import numpy as np

def posterior(prior, lik):
    raise NotImplementedError

Maximum likelihood, where losses come from

Suppose a model with parameters $\theta$ assigns probability $p_\theta(x)$ to each possible data point, and we observe independent data $x_1, \dots, x_N$. The likelihood is the probability of the whole data set as a function of $\theta$, and maximum likelihood estimation (MLE) picks the $\theta$ that makes it largest,

In plain words, choose the parameters under which the data you actually saw are as probable as possible, which is the same as minimizing the average negative log-probability. We take logs for two reasons. A product of thousands of probabilities underflows (you just saw this), while a sum of logs does not. And sums are much easier to differentiate.

Almost every loss in this book is a negative log-likelihood in disguise.

Bernoulli gives binary cross-entropy. With $p_\theta(y \mid \mathbf{x}) = \hat p^{\,y}(1 - \hat p)^{1-y}$ where $\hat p = \sigma(z_\theta(\mathbf{x}))$, the negative log-likelihood is $-[y\log\hat p + (1-y)\log(1-\hat p)]$, the binary cross-entropy. Reward models (§16.3) and DPO (§17.2) are trained with exactly this loss.

Categorical gives cross-entropy. With $p_\theta(y \mid \mathbf{x}) = \operatorname{softmax}(\mathbf{z})_y$, the negative log-likelihood is $-\log \operatorname{softmax}(\mathbf{z})_y$, the cross-entropy of §0.2. Language-model pretraining is maximum likelihood on next tokens, nothing more.

Gaussian gives squared error. With $y = f_\theta(\mathbf{x}) + \text{noise}$ and Gaussian noise of fixed variance $\sigma^2$, the negative log-likelihood is $\frac{1}{2\sigma^2}(y - f_\theta(\mathbf{x}))^2$ plus a constant. Minimizing it is least squares. So "use MSE" is secretly the assumption "my errors are Gaussian", which tells you when it is a bad idea (heavy-tailed noise, where one outlier dominates).

Worked example, the MLE of a Gaussian

Given samples $x_1, \dots, x_N$, maximize $\ell(\mu, \sigma^2) = \sum_n \log \mathcal{N}(x_n; \mu, \sigma^2) = -\frac{1}{2\sigma^2}\sum_n (x_n - \mu)^2 - \frac N2\log(2\pi\sigma^2)$. Setting the derivatives to zero,

In plain words, the MLE of the mean is the sample mean, and the MLE of the variance is the average squared deviation from it.

That variance estimate is biased. Its expectation is $\frac{N-1}{N}\sigma^2$, slightly too small, because $\hat\mu$ was fitted to the same data and sits closer to the samples than the true mean does. Dividing by $N - 1$ instead (Bessel's correction) removes the bias. With $N = 5$, the MLE underestimates the variance by 20% on average. For the large $N$ of machine learning it hardly matters, but it is a favorite interview question, and you will meet the $1/N$ versus $1/(N-1)$ choice again inside BatchNorm (§0.7).

import numpy as np
from scipy.linalg import solve_triangular

def fit_gaussian(X):
    raise NotImplementedError

def mvn_logpdf(X, mu, Sigma):
    raise NotImplementedError

MAP, and priors as regularizers

Maximum likelihood trusts the data completely. With little data and a flexible model, that is a problem. Flip a coin three times, see three heads, and the MLE says the coin always lands heads. Fit a degree-9 polynomial to 10 points and the MLE passes through every point while oscillating wildly between them.

The Bayesian fix is to put a prior on the parameters and maximize the posterior instead, which is called maximum a posteriori (MAP) estimation. By Bayes rule, $p(\theta \mid \mathcal{D}) \propto p(\mathcal{D} \mid \theta)\, p(\theta)$, so

In plain words, MAP is maximum likelihood plus a penalty, and the penalty is the negative log-prior.

Now look at what common priors become.

• A Gaussian prior $\theta_j \sim \mathcal{N}(0, \tau^2)$ has $-\log p(\theta) = \frac{1}{2\tau^2}\lVert\theta\rVert^2 + \text{const}$. That is the L2 penalty, called weight decay in deep learning (§0.6 and §0.7).
• A Laplace prior $p(\theta_j) \propto e^{-|\theta_j|/b}$ has $-\log p(\theta) = \frac1b\lVert\theta\rVert_1 + \text{const}$, the L1 penalty, which drives many weights to exactly zero.

Here is the derivation for linear regression, which is worth doing once in full. Let $\mathbf{y} = \mathbf{X}\mathbf{w} + \boldsymbol{\epsilon}$ with $\mathbf{X} \in \mathbb{R}^{N \times d}$, noise $\epsilon_n \sim \mathcal{N}(0, \sigma^2)$, and prior $w_j \sim \mathcal{N}(0, \tau^2)$. The negative log-posterior, dropping constants, is

Its gradient is $-\frac{1}{\sigma^2}\mathbf{X}^\top(\mathbf{y} - \mathbf{X}\mathbf{w}) + \frac{1}{\tau^2}\mathbf{w}$. Setting it to zero and multiplying by $\sigma^2$,

In plain words, MAP with a Gaussian prior is ridge regression, and the regularization strength is the ratio of noise variance to prior variance. A tight prior (small $\tau$) or noisy data (large $\sigma$) means strong regularization. As $\tau \to \infty$ the prior becomes flat and you recover least squares. A bonus is that $\mathbf{X}^\top\mathbf{X} + \lambda\mathbf{I}$ is always invertible for $\lambda > 0$, even when $d > N$.

The MLE fit needs enormous coefficients (the largest is about 370,000) to thread every point, and it generalizes badly. A modest prior, $\lambda = 0.001$, shrinks the largest coefficient to about 10 and cuts the test error by more than half. A stronger one, $\lambda = 0.1$, already flattens the fit too much and the error rises again. That is the bias side of the trade-off you'll meet in §0.7.

The coin makes the same point with a prior you can see. For a coin with heads probability $\theta$, the Beta distribution $\text{Beta}(a, b)$, with density proportional to $\theta^{a-1}(1-\theta)^{b-1}$, is a convenient prior. After $h$ heads and $t$ tails the posterior is $\text{Beta}(a + h, b + t)$, so the prior acts like $a$ imaginary heads and $b$ imaginary tails added to the data (for the MAP estimate, the most probable value, it is $a - 1$ and $b - 1$). The MLE is $h/(h + t)$, and the MAP estimate is $(h + a - 1)/(h + t + a + b - 2)$. Try three heads and no tails with a weak prior, then add data and watch the prior's influence fade.

import numpy as np

def map_linear_regression(X, y, sigma2, tau2):
    raise NotImplementedError

Sampling

Generating from a model means sampling, and several of the most important gradient estimators in machine learning are really facts about sampling.

Inverse CDF. If $U$ is uniform on $[0, 1]$ and $F$ is a cumulative distribution function, then $F^{-1}(U)$ has distribution $F$. For a categorical distribution, $F$ is the running sum of the probabilities, and inverting it is a binary search. The word2vec negative sampler of §1.5 is built exactly this way.

The Gumbel-max trick. To sample from $\operatorname{softmax}(\mathbf{z})$ without computing the softmax, add independent Gumbel noise $g_k = -\log(-\log U_k)$ to every logit and take the argmax,

In plain words, perturb every logit with the right kind of noise and pick the largest. It never exponentiates, so it cannot overflow. It parallelizes trivially, and it underlies the Gumbel-softmax (or Concrete) relaxation that makes discrete choices differentiable (Jang et al., 2017, and Maddison, Mnih & Teh, 2017). The trick itself is much older, a consequence of the extreme-value statistics named after Emil Gumbel.

Temperature. Divide the logits by a temperature $\tau > 0$ before the softmax and sample from $\operatorname{softmax}(\mathbf{z}/\tau)$. (In this paragraph $\tau$ is a temperature, not the prior scale of the MAP discussion.) A temperature below 1 sharpens the distribution toward the argmax, and one above 1 flattens it toward uniform. This is the main knob of text generation (§6.4).

import numpy as np

def gumbel_max_sample(logits, n, rng):
    raise NotImplementedError

Two ways to differentiate through randomness

Often we need the gradient of an expectation, $\nabla_\theta\, \mathbb{E}_{x \sim p_\theta}[f(x)]$, where the distribution itself depends on $\theta$. Reinforcement learning needs this (the expected reward of a policy), and so do variational autoencoders. There are two standard estimators, and the contrast between them is a classic interview topic.

The score-function estimator, also called REINFORCE or the log-derivative trick, uses $\nabla_\theta p_\theta = p_\theta\nabla_\theta\log p_\theta$,

In plain words, sample, then weight each sample's log-probability gradient by how good the sample was. It works for any $f$, even a non-differentiable one like a reward from a human or a unit test, which is why it is the foundation of policy gradients (§15.3). Its weakness is high variance.

The reparameterization estimator writes the sample as a differentiable function of $\theta$ and fixed noise. For a Gaussian, $x = \mu + \sigma\epsilon$ with $\epsilon \sim \mathcal{N}(0, 1)$, so $\nabla_\mu \mathbb{E}[f(x)] = \mathbb{E}_\epsilon[f'(\mu + \sigma\epsilon)]$. It needs $f$ to be differentiable and the distribution to be reparameterizable, but its variance is usually far lower. It is the trick that lets variational autoencoders train by backpropagation through a random sampling step.

Both are unbiased. But with $\mu = 1$ the per-sample variance of the score-function estimator is exactly 30 (expand $(1 + \epsilon)^2\epsilon$ and use the Gaussian moments $\mathbb{E}[\epsilon^4] = 3$ and $\mathbb{E}[\epsilon^6] = 15$), against 4 for the reparameterized one. So it needs about 7.5 times as many samples for the same precision. The cell's per-sample standard deviations, about 5.8 and 2.0, are close to the theoretical $\sqrt{30} \approx 5.5$ and $2$. Much of §15.4 is about cutting that variance with baselines.

The log-sum-exp trick and the stable softmax

Every softmax, every cross-entropy, and every log-likelihood of a categorical model computes $\log\sum_k e^{z_k}$, the log-sum-exp, which we write $\operatorname{LSE}(\mathbf{z})$ for logits $\mathbf{z} \in \mathbb{R}^K$. You met it, and the trick that makes it safe, in §0.2. This section completes the story, with the limits of each floating-point format, the two facts that make LSE a friendly function, and the edge cases (masked rows, infinities) that production code must handle. Computed naively, it breaks. $e^{z}$ overflows to infinity once $z$ exceeds about 88.7 in float32 or 709.8 in float64, and in float16 it overflows once $z$ passes about 11, so in low precision perfectly ordinary logits already break it. In the other direction, if every $z_k$ is very negative, every $e^{z_k}$ underflows to 0 and the log returns $-\infty$.

As in §0.2, the fix factors out the largest term. Let $m = \max_k z_k$. Then

In plain words, subtract the maximum before exponentiating and add it back after the log. Now every exponent is at most 0, so nothing overflows, and the largest term is exactly $e^0 = 1$, so the sum is at least 1 and the log is safe.

The same shift gives the stable softmax, since $\operatorname{softmax}(\mathbf{z}) = \operatorname{softmax}(\mathbf{z} - m\mathbf{1})$ (softmax is shift invariant, as we saw from its Jacobian in §0.2). And it gives the stable log-softmax, $\log\operatorname{softmax}(\mathbf{z})_k = z_k - \operatorname{LSE}(\mathbf{z})$, which you should always use instead of taking the log of a softmax, because a probability that underflowed to 0 has a log of $-\infty$.

Two more facts make log-sum-exp a friendly function. It is a smooth maximum, since $\max_k z_k \le \operatorname{LSE}(\mathbf{z}) \le \max_k z_k + \log K$. And its gradient is the softmax, $\nabla_{\mathbf{z}}\operatorname{LSE}(\mathbf{z}) = \operatorname{softmax}(\mathbf{z})$, which is the one-line derivation of $\mathbf{p} - \mathbf{y}$ from §0.2.

The same idea, carried out in blocks so the maximum can be updated as new scores arrive, is the online softmax at the heart of FlashAttention (§12.2).

import numpy as np

def logsumexp(z, axis=-1, keepdims=False):
    raise NotImplementedError

def log_softmax(z, axis=-1):
    raise NotImplementedError

In PyTorch

torch.distributions packages densities, sampling, and both gradient estimators. sample() draws without a gradient path, rsample() uses reparameterization, and log_prob gives what the score-function estimator needs.

Read more

• Bishop (2006). Pattern Recognition and Machine Learning. Chapters 1 and 2 are the classic treatment of MLE, MAP, the Gaussian, and the polynomial example above.
• Murphy (2022). Probabilistic Machine Learning: An Introduction. Chapters 2 to 4 cover the probability, distributions, and estimation used here, freely available.
• Blitzstein & Hwang (2019). Introduction to Probability. The book behind Harvard's Stat 110, the best first course in probability, with free lectures and a free edition.
• Mohamed, Rosca, Figurnov & Mnih (2020). Monte Carlo Gradient Estimation in Machine Learning. The definitive survey of score-function and reparameterization estimators.
• Jang, Gu & Poole (2017). Categorical Reparameterization with Gumbel-Softmax. The Gumbel-max trick in practice and its differentiable relaxation.
• Blanchard, Higham & Higham (2021). Accurately computing the log-sum-exp and softmax functions. A careful numerical analysis of the shift trick.

Watch next

▶

▶

▶