Cross-entropy and KL divergence: the language-model loss

What does it cost to describe data with the wrong probability model, and why is that cost the loss nearly every language model is trained on? Cross-entropy measures that cost, and KL divergence isolates the part that comes from the model being wrong.

Coding with the wrong distribution

Entropy told us that an outcome of probability $p(x)$ deserves a codeword of about $-\log p(x)$ bits. Now suppose the data really come from $p$, but we design our code from a model $q$. Each outcome $x$ then costs $-\log q(x)$, and the average cost is the cross-entropy:

Split it into two parts:

The first term is the unavoidable cost of the source's own uncertainty. The second is the Kullback–Leibler divergence, introduced by Solomon Kullback and Richard Leibler in 1951 as a measure of the information available for telling two distributions apart.^[1] In coding terms, it is the expected number of extra bits you pay per symbol for using $q$ in place of $p$.^[2]

Gibbs' inequality

$D_{\mathrm{KL}}(p\,\|\,q)\ge 0$, with equality if and only if $q=p$. One line of proof uses Jensen's inequality on the concave logarithm:

So $H(p,q)\ge H(p)$. No model can beat the true distribution at coding its own data, and the gap is exactly the KL divergence.^[2]

Worked example: KL is not symmetric

Let $p = (\tfrac12, \tfrac14, \tfrac14)$ and let $q$ be uniform, $(\tfrac13,\tfrac13,\tfrac13)$. In bits:

• $H(p) = \tfrac12(1)+\tfrac14(2)+\tfrac14(2) = 1.5$.
• $H(p,q) = \log_2 3 \approx 1.585$, because every outcome costs $\log_2 3$ under the uniform code.
• $D_{\mathrm{KL}}(p\|q) = 1.585 - 1.5 \approx 0.085$ bits.

Now reverse the roles:

The two directions give different numbers. KL divergence is not a distance metric: it is asymmetric and does not satisfy the triangle inequality. Which direction you minimise changes what you get, as discussed below.

Why this is the language-model loss

A language model defines conditional distributions $q_\theta(x_t\mid x_{<t})$ over a vocabulary. Training on a corpus minimises the average negative log-probability of each actual next token:

This is cross-entropy with a "one-hot" target: the empirical distribution puts all its mass on the token that actually appeared. Averaged over the data distribution $p$, it is

Since $H(p)$ does not depend on $\theta$, minimising cross-entropy is the same as minimising the forward KL divergence from the data to the model. It is also the same as maximising likelihood, because $-\log q_\theta$ summed over the data is the negative log-likelihood. The probabilistic view of this is in Maximum likelihood. Three standard ways of describing training, as maximum likelihood, KL minimisation and code-length minimisation, all describe the same objective.

Units and perplexity

Training code usually computes the loss in nats. Dividing by $\ln 2$ gives bits per token. Perplexity is the exponentiated loss, $\exp(\mathcal{L}_{\text{nats}})$. Suppose a model's validation loss is 2.0 nats per token. That is $2.0/0.693 \approx 2.89$ bits per token, and its perplexity is $e^{2.0}\approx 7.39$. On average the model is as uncertain as if it were choosing uniformly among about 7.4 tokens. Because the irreducible $H(p)$ is unknown, a loss value alone does not tell you how close a model is to the floor. Loss comparisons are meaningful only on the same data with the same tokenizer.

The gradient is "prediction minus target"

With logits $z$ and $q = \mathrm{softmax}(z)$, the cross-entropy against a target distribution $y$ (one-hot or soft) has a very simple gradient:

Each logit is pushed up by the probability the model failed to give the right answer and pushed down by the probability it gave the wrong ones. The gradient is bounded and well scaled, which is part of why softmax with cross-entropy is the default output layer for classifiers and language models. The full backward pass is covered in Backpropagation.

Forward versus reverse KL

The direction of the divergence decides the behaviour of the fitted model:

Two places in modern training make the distinction concrete:

• Distillation. A small "student" is trained to match a large "teacher" by minimising cross-entropy against the teacher's full output distribution ("soft targets"), not only the one-hot label.^[3] Since the teacher's entropy is fixed, this is forward-KL minimisation toward the teacher. The soft targets carry information about which wrong answers are nearly right.
• Fine-tuning with a KL leash. In reinforcement learning from human feedback, the policy is rewarded by a learned reward model but penalised by a per-token KL term against the supervised starting model. The penalty discourages it from drifting into outputs the reward model scores well but the original model considers very unlikely.^[4] Here the expectation is taken under the policy's own samples, which makes it the reverse direction.

Key takeaways

• Cross-entropy $H(p,q)$ is the average code length when data from $p$ are coded with model $q$. It equals $H(p) + D_{\mathrm{KL}}(p\|q)$.
• KL divergence is non-negative and zero only when $q=p$. It is asymmetric, and the direction you choose determines whether the fitted model is mass-covering or mode-seeking.
• The language-model loss is per-token cross-entropy. Minimising it is the same as maximum likelihood and the same as minimising forward KL from data to model.
• The softmax cross-entropy gradient is $q - y$: prediction minus target.
• Perplexity is $e^{\text{loss}}$. Loss numbers can be compared only on the same data and tokenizer, because the entropy floor $H(p)$ is unknown.

Sources