Maximum Likelihood, Cross-Entropy & KL Divergence

Why It Matters for Modern Models

Maximum likelihood chooses the parameters of a probability model that make the observed data most probable. Language-model pretraining is maximum likelihood: it minimizes next-token cross-entropy on large collections of text and code.
Diffusion models such as Stable Diffusion are trained with a noise-prediction mean-squared error, which is a reweighted form of the variational bound on log-likelihood (Ho et al., 2020; Rombach et al., 2022).
Reward models for RLHF are trained with a cross-entropy (logistic) loss on human comparisons between pairs of responses (Ouyang et al., 2022).

Questions to Think About

How does a model trained with cross-entropy behave on inputs unlike its training data, where it must still put its probability somewhere and may produce fluent but wrong text?

Interactive Demo

Core Math

If you prefer intuition first, try the demo above, then come back for the math.

Key Equation

\min_\theta H(\hat p, p_\theta) = -\mathbb E_{x\sim \hat p} \log p_\theta(x)

Language models, and many other generative models, are fitted by maximum likelihood (exactly, or through a bound on it):

\max_\theta \sum_{i=1}^n \log p_\theta(x^{(i)})

Equivalently, minimize the empirical cross-entropy between the data distribution $\hat p$ and the model $p_\theta$ :

\min_\theta H(\hat p, p_\theta) = \min_\theta \left[ -\mathbb E_{x\sim \hat p} \log p_\theta(x) \right]

For discrete data, this is the same as minimizing the KL divergence, because $\mathrm{KL}(\hat p \,\|\, p_\theta) = H(\hat p, p_\theta) - H(\hat p)$ and the entropy $H(\hat p)$ does not depend on $\theta$ :

\mathrm{KL}(\hat p \,\|\, p_\theta) = \mathbb E_{\hat p} \log \frac{\hat p(x)}{p_\theta(x)}

For autoregressive language models, the factorization comes from the chain rule of probability:

p_\theta(x_1,\dots,x_T) = \prod_{t=1}^T p_\theta(x_t \mid x_{<t})

Key Papers and Books

A Neural Probabilistic Language Model

Bengio et al.2003JMLR

Read paper →

Deep Learning, Section 5.5: Maximum Likelihood Estimation

Goodfellow, Bengio & Courville2016MIT Press (textbook)

View book →

Connections

Where to go next

A newer notebook covers Maximum Likelihood, with an interactive demo.

Continue in the newer notebook

Other connections (12)

Which limitations does it address, or leave to a later idea?

Stochastic MLE: a limitation of Maximum Likelihood that Adam addresses. Open Adam
Training by maximum likelihood means minimizing an average cross-entropy over a large dataset, in practice with noisy minibatch gradients. Adam was designed for such stochastic objectives (Kingma & Ba, 2015): momentum smooths the noisy gradients, and per-parameter step sizes help with sparse gradients, such as those of rare-word embeddings.
Fix text degeneration: a limitation of Maximum Likelihood that Decoding & Sampling addresses. Open Decoding & Sampling
Greedy decoding from a model trained by maximum likelihood often falls into repetitive loops. Nucleus sampling was introduced as a decoding-time fix that avoids those loops while cutting off the unreliable low-probability tail, without retraining (Holtzman et al., 2020).

Where does it break down?

Non-overlapping supports: where GANs breaks down, taken up on this page. Open GANs
When the generator's distribution and the data distribution do not overlap, the Jensen–Shannon divergence that the original GAN minimizes is constant at log 2, so a near-optimal discriminator gives the generator almost no gradient (Arjovsky & Bottou, 2017). WGAN replaced it with the Wasserstein distance, which keeps changing smoothly as the two distributions move closer (Arjovsky et al., 2017).
Unseen words: where Tokenization breaks down, taken up on this page. Open Tokenization
A word-level vocabulary has no token for an unseen word, so the model cannot assign it any probability. Subword methods such as byte-pair encoding (Sennrich et al., 2016), especially with a fallback to bytes, remove that problem, but rare strings are split into many pieces whose probabilities are estimated from little data.
Repetition loops: where Decoding & Sampling breaks down, taken up on this page. Open Decoding & Sampling
Greedy or beam-search decoding from a model trained by maximum likelihood often falls into repetitive loops (Holtzman et al., 2020). Sampling with temperature, top-k or top-p adds randomness that avoids them, at the cost of no longer choosing the single most likely continuation.

Where else does the same idea appear?

Cross-entropy: the same technique appears in RLHF. Open RLHF
RLHF reward models are trained with cross-entropy on preference pairs: p(preferred > rejected). Same loss function, different label source.
KL divergence: the same technique appears in VAEs. Open VAEs
VAE ELBO contains KL(q(z|x) || p(z))—same KL divergence from MLE, but now regularizing latents toward prior.
Viterbi segmentation: the same technique appears in Tokenization. Open Tokenization
Unigram tokenization chooses the max-prob segmentation by maximizing ∑ log p(tᵢ) under concat constraints. That is MAP/MLE-style inference solved by dynamic programming (a Viterbi-like trick).
Likelihood ratios: the same technique appears in DPO. Open DPO
DPO is logistic regression on log-probability ratios (log πθ − log πref). That's classic maximum-likelihood machinery: fit a classifier via cross-entropy on preference-labeled data.
MLE-trained classifier: the same technique appears in GANs. Open GANs
GAN discriminator maximizes likelihood of correct real/fake labels via logistic cross-entropy. MLE classification repurposed as training signal for generator.

Which ideas are its mathematical duals?

Likelihood ↔ MDL: Learning Theory is its mathematical dual. Open Learning Theory
The negative log-likelihood −log₂ p_θ(x) is the length in bits of x under the code built from the model p_θ, so maximizing likelihood minimizes the length of the encoded data. MDL adds the length of the description of the model itself, which is what penalizes complex models.
Denoising ↔ ELBO: Diffusion is its mathematical dual. Open Diffusion
With the right weight on each noise level, the noise-prediction loss equals the variational bound on −log p(x) under the fixed noising process, up to constants and a final decoding term; the unweighted "simple" loss used in practice is a reweighted version of that bound (Ho et al., 2020). It is denoising on the surface and approximate likelihood training underneath.