Information Theory and Loss Functions

Nearly every classifier and every language model is trained by minimizing one number: the cross-entropy between the data and the model's predicted distribution. Language models are compared by perplexity; distillation and RLHF use KL divergence; sampling is controlled by temperature on logits. All of these come from information theory, the mathematics of uncertainty and surprise that Claude Shannon introduced in 1948. This lesson derives them from one idea, connects them to maximum likelihood, and shows how to compute them without numerical overflow.

Intuition: surprise is measured in bits

If someone tells you the sun rose this morning, you learn almost nothing; if they tell you it snowed in Dallas in July, you learn a lot. Information is surprise, and surprise is high for unlikely events. A good predictor is one that is rarely surprised. Cross-entropy is just "the average surprise your model experiences on real data", and training lowers it.

To make this quantitative, imagine sending messages with a code. Frequent symbols should get short codewords and rare ones long codewords (Morse code gives "E" a single dot). Shannon showed that the best possible average codeword length is set by the distribution's entropy, and that using a code designed for the wrong distribution costs extra bits, exactly the KL divergence.

Self-information and entropy

The self-information (surprisal) of an outcome with probability $p$ is

With $\log_2$ it is measured in bits; with the natural log, in nats ($1 \text{ nat} = 1/\ln 2 \approx 1.443$ bits). Deep-learning frameworks use natural logs, so losses are in nats. A certain event ($p = 1$) carries 0 information; a fair coin flip ($p = 1/2$) carries 1 bit.

The entropy of a distribution $p$ is the expected surprisal:

with the convention $0 \log 0 = 0$. For $K$ outcomes, $0 \le H(p) \le \log K$; the minimum is a deterministic distribution, the maximum the uniform one.

Worked example. A fair coin: $H = -2 \times \tfrac12 \log_2 \tfrac12 = 1$ bit. A biased coin with $p = (0.9, 0.1)$:

The biased coin is more predictable, so it has lower entropy. Where entropy appears in AI: decision trees split on information gain (reduction in entropy); reinforcement-learning algorithms add an entropy bonus to keep policies exploring; and the entropy of a model's predictive distribution is a simple uncertainty signal (high entropy over next tokens often accompanies an LLM being unsure).

Cross-entropy

If data truly come from $p$ but you predict with a model $q$, your average surprise is the cross-entropy

In classification, the "true distribution" for one example is a one-hot vector: all mass on the correct class $y$. The sum then collapses to a single term:

So cross-entropy loss is simply the negative log of the probability the model assigned to the correct answer. Worked example: three classes, true class 0, prediction $q = (0.7, 0.2, 0.1)$: $\ell = -\ln 0.7 \approx 0.357$ nats. If instead $q(y) = 0.01$, $\ell = -\ln 0.01 \approx 4.605$: confident wrong answers are punished heavily, and as $q(y) \to 0$ the loss goes to infinity. Over a dataset you average: $\frac{1}{N}\sum_i -\log q(y_i \mid x_i)$, which is exactly the negative log-likelihood from the probability lesson.

KL divergence

The Kullback–Leibler divergence measures the extra surprise from using $q$ when the truth is $p$:

Properties: $D_{\mathrm{KL}}(p\|q) \ge 0$ with equality only when $p = q$ (Gibbs' inequality), and it is not symmetric, so it is not a true distance.

Worked example. $p = (0.5, 0.5)$, $q = (0.9, 0.1)$, natural logs:

Because $H(p)$ does not depend on the model, minimizing cross-entropy $H(p, q)$ over $q$ is the same as minimizing $D_{\mathrm{KL}}(p\|q)$. That is the deep reason cross-entropy is the right training objective: it pulls the model distribution toward the data distribution.

Where KL appears: knowledge distillation trains a student to match a teacher's softened output distribution; RLHF and related methods penalize the KL divergence between the fine-tuned policy and the reference model so the model does not drift too far; variational autoencoders include a KL term pulling the latent distribution toward a prior. The direction matters: $D_{\mathrm{KL}}(p\|q)$ (forward) heavily penalizes $q$ for missing places where $p$ has mass, so it is "mass-covering"; $D_{\mathrm{KL}}(q\|p)$ (reverse) penalizes $q$ for putting mass where $p$ has little, so it is "mode-seeking".

Logits and softmax

Networks output unconstrained real numbers called logits $\mathbf{z} = (z_1, \dots, z_K)$. The softmax turns them into a probability distribution:

Every output is positive and they sum to 1. Softmax is invariant to adding a constant to all logits ($e^{z_k + c} / \sum_j e^{z_j + c}$, the $e^c$ cancels), so only differences between logits matter. For two classes, softmax reduces to the sigmoid of the logit difference.

Worked example. $\mathbf{z} = (2, 1, 0.1)$: $e^{2} = 7.389$, $e^{1} = 2.718$, $e^{0.1} = 1.105$; sum $= 11.213$; probabilities $\approx (0.659, 0.242, 0.099)$.

Temperature

Dividing logits by a temperature $T > 0$ before the softmax reshapes the distribution:

As $T \to 0$ the distribution approaches the argmax (greedy decoding); as $T \to \infty$ it approaches uniform. When sampling text, low temperature gives focused, repetitive output and high temperature gives diverse, error-prone output. Distillation uses $T > 1$ on both teacher and student so the student learns from the relative probabilities of wrong classes ("dark knowledge"), as Hinton, Vinyals and Dean described in 2015. Temperature scaling is also a standard post-hoc calibration method: fit a single $T$ on a validation set to make probabilities honest without changing the argmax.

Softmax plus cross-entropy: the gradient

Combine softmax and cross-entropy for one example with logits $\mathbf{z}$ and one-hot target $\mathbf{y}$:

Differentiate with respect to $z_k$: the first term gives $-1$ if $k = y$, and the derivative of the log-sum-exp is $e^{z_k}/\sum_j e^{z_j} = q_k$. Therefore

The gradient is "prediction minus target", bounded between $-1$ and $1$ per logit, and it never vanishes while the model is wrong. For the example above with true class 0, $\nabla = (0.659 - 1, 0.242, 0.099) = (-0.341, 0.242, 0.099)$: raise the correct logit, lower the others. Squared error on top of a softmax or sigmoid lacks this property; its gradient includes the softmax/sigmoid derivative, which goes to zero when the output saturates, so a confidently wrong model learns slowly.

Why cross-entropy for classification and MSE for regression

Both losses fall out of maximum likelihood with different noise models:

• Classification. Model the label as categorical with probabilities $\mathbf{q} = \mathrm{softmax}(\mathbf{z})$. The NLL of the observed label is $-\log q_y$: cross-entropy. (Binary case: Bernoulli likelihood gives binary cross-entropy $-[y \log \hat p + (1-y)\log(1 - \hat p)]$.)
• Regression. Model the target as Gaussian around the prediction: $y \sim \mathcal{N}(\hat y, \sigma^2)$. The NLL is $\frac{(y - \hat y)^2}{2\sigma^2} + \frac12 \log(2\pi\sigma^2)$; with $\sigma$ fixed, minimizing it is minimizing squared error. If you assume Laplace noise instead, you get mean absolute error, which is more robust to outliers.

Choosing a loss is therefore choosing an assumption about how the data are distributed around the model's prediction.

Perplexity and language models

A language model factorizes the probability of a token sequence $x_1, \dots, x_N$ by the chain rule of probability, $p(x_{1:N}) = \prod_{t=1}^N p(x_t \mid x_{<t})$, and is trained to minimize the average per-token cross-entropy

Perplexity is its exponential:

the geometric mean of inverse probabilities. Intuitively, a perplexity of $k$ means the model is, on average, as uncertain as if it were choosing uniformly among $k$ tokens. A uniform model over a 50,000-token vocabulary has perplexity 50,000; a perfect model has perplexity 1.

Worked example. A model assigns probabilities $0.5, 0.25, 0.125, 0.5$ to four actual tokens. Per-token losses: $0.693, 1.386, 2.079, 0.693$ nats; mean $\approx 1.213$; perplexity $= e^{1.213} \approx 3.36$. Check via the geometric mean: $(2 \times 4 \times 8 \times 2)^{1/4} = 128^{1/4} \approx 3.36$.

Caveats: perplexity depends on the tokenizer (a model with a larger vocabulary predicts fewer, harder tokens), so it is only comparable between models sharing a tokenizer and test text. For tokenizer-independent comparison people use bits per byte or bits per character. And lower perplexity does not guarantee better instruction following or reasoning; it measures fit to a text distribution.

Numerically stable log-softmax

Naively computing $e^{z}$ overflows float32 for $z$ above about 88, and taking $\log$ of a probability that underflowed to 0 gives $-\infty$. The fix uses the shift invariance: with $m = \max_j z_j$,

Now the largest exponent is $e^0 = 1$, so nothing overflows, and the sum is at least 1, so the log is safe. This is the log-sum-exp trick.

import numpy as np

def log_softmax(z, axis=-1):
    m = z.max(axis=axis, keepdims=True)
    return z - m - np.log(np.exp(z - m).sum(axis=axis, keepdims=True))

def cross_entropy(logits, labels):                # logits (B, K), labels (B,)
    lp = log_softmax(logits)
    return -lp[np.arange(len(labels)), labels].mean()

def entropy(p):  return -(p * np.log(np.where(p > 0, p, 1))).sum(-1)
def kl(p, q):    return (p * (np.log(p) - np.log(q))).sum(-1)

z = np.array([[2.0, 1.0, 0.1]])
for T in (0.5, 1.0, 2.0):
    print(T, np.exp(log_softmax(z / T)).round(3))

big = np.array([[1000.0, 1001.0, 1002.0]])
print(np.exp(log_softmax(big)).round(3))          # [[0.09 0.245 0.665]], no overflow

print(round(cross_entropy(z, np.array([0])), 4))  # 0.417 = -log 0.659
print(kl(np.array([.5, .5]), np.array([.9, .1])).round(3),
      kl(np.array([.9, .1]), np.array([.5, .5])).round(3))   # 0.511 0.368

# gradient of CE wrt logits equals softmax - onehot (numerical check)
q = np.exp(log_softmax(z))[0]; y = np.array([1., 0., 0.])
eps = 1e-6; num = np.zeros(3)
for k in range(3):
    zp, zm = z.copy(), z.copy(); zp[0, k] += eps; zm[0, k] -= eps
    num[k] = (cross_entropy(zp, np.array([0])) - cross_entropy(zm, np.array([0]))) / (2 * eps)
print(num.round(4), (q - y).round(4))

probs = np.array([0.5, 0.25, 0.125, 0.5])
print("perplexity", round(np.exp(-np.log(probs).mean()), 3))   # 3.364

Note the loss on the logits $(2, 1, 0.1)$ is $-\ln 0.659 \approx 0.417$; the earlier $0.357$ example used the probability $0.7$ directly. Frameworks fuse softmax and log for exactly this stability reason: PyTorch's nn.CrossEntropyLoss and F.cross_entropy take raw logits and apply a stable log-softmax internally, and F.log_softmax exists for when you need log-probabilities.

Common pitfalls

• Applying softmax before a loss that expects logits. Passing probabilities into CrossEntropyLoss applies softmax twice; the model still trains, but poorly.
• Computing log(softmax(x)) in two steps. Use a fused log-softmax or cross-entropy to avoid -inf and nan.
• Comparing perplexities across tokenizers or datasets. The numbers are not on the same scale.
• Mixing bits and nats. Framework losses are in nats; divide by $\ln 2$ for bits.
• KL direction errors. $D_{\mathrm{KL}}(p\|q)$ and $D_{\mathrm{KL}}(q\|p)$ give different gradients and behaviours; check which one a method specifies, and note that torch.nn.functional.kl_div expects its first argument as log-probabilities of the model.
• Taking log(0). Add label smoothing, clip probabilities, or work in log space.

In practice

• Train classifiers on logits with fused cross-entropy; use label smoothing (mixing the one-hot target with a little uniform mass) when overconfidence is a problem.
• Report validation loss in nats per token and perplexity for language models, and also bits per byte when comparing different tokenizers.
• Calibrate with a single temperature fitted on held-out data when probabilities drive decisions.
• Tune sampling temperature per task: low values (or greedy decoding) for extraction and code, moderate values for open-ended writing.
• Watch the KL to the reference model during preference fine-tuning; a rapidly growing KL often signals reward hacking.