Loss Functions

To train a model we need one number that says how wrong it is: the loss. Training then follows the gradient of the loss (section 1) to make it smaller. A language model is trained with exactly one loss, the cross-entropy of the next token. This page builds it step by step, makes it safe to compute, and adds the two details every LLM (large language model) training loop has: ignored positions and label smoothing.

The idea: minus the log of the right answer's probability

Most models do not give one answer: they give a probability to every possible answer. The natural way to score such a model is to look at the probability it gave to the answer that turned out to be right. High probability: good. Low probability: bad.

The loss of one example is minus the log of that probability, $-\log p$. The log turns probabilities into a scale that punishes confident mistakes hard:

For a whole dataset, average over the examples:

where:

• $\theta$ (theta) are the model's parameters;
• $N$ is the number of examples, and $(x_n, y_n)$ is example $n$: an input $x_n$ with its right answer $y_n$;
• $p_\theta(y_n \mid x_n)$, read "the probability of $y_n$ given $x_n$", is the probability the model gives to the right answer.

Example. On two examples the model gives the right answer probability $0.9$ and $0.5$. The losses are $0.105$ and $0.693$, and their average is $\mathcal{L} \approx 0.40$.

This is called the negative log-likelihood, and minimizing it is maximum likelihood: choosing the parameters under which the training data is most probable. Why a log and an average, rather than multiplying the probabilities directly? The probability of the whole dataset is the product of the probabilities of its examples, and the log turns a product into a sum, which is much easier to work with. The minus sign is there because optimizers minimize.

The familiar losses are all this one

Depending on what the model outputs, the same recipe gives the losses you may have seen under different names:

Here $\hat y$ ("y hat") is the model's prediction. How to read the rows:

• Predicting a number (a price, a reward): if you assume the true value is the prediction plus bell-shaped (Gaussian) noise, minus the log-probability is the squared error, up to constants. So the MSE is not an arbitrary choice.
• Predicting yes or no: the model gives a probability $q$ of "yes", and the loss is $-\log q$ when the answer is yes, $-\log(1 - q)$ when it is no. The formula in the table is just these two cases in one line.
• Choosing among $C$ classes: the model gives a probability $q_c$ to each class $c$, and the loss is $-\log q_y$ for the right class $y$. For a language model the classes are the tokens of the vocabulary, so this is the row that matters.

Cross-entropy, KL divergence and perplexity

Sometimes the right answer is not a single class but a distribution, for example 70% "cat", 20% "dog", 10% "sun". Call the true distribution $p$ and the model's distribution $q$, both over the same $C$ classes. The cross-entropy averages $-\log q_c$ with the true probabilities as weights:

where $p_c$ and $q_c$ are the true and the model's probabilities of class $c$. It splits into two parts:

where:

• $H(p) = -\sum_c p_c \log p_c$ is the entropy of the data: how unpredictable the answer is by itself. The model cannot change it;
• $D_{\mathrm{KL}}(p \,\|\, q)$ is the KL (Kullback–Leibler) divergence: how far the model's distribution is from the true one. It is 0 only when $q = p$.

Example. True $p = (0.7, 0.2, 0.1)$, model $q = (0.5, 0.3, 0.2)$.

• Cross-entropy: $-(0.7 \log 0.5 + 0.2 \log 0.3 + 0.1 \log 0.2) \approx 0.887$.
• Entropy: $-(0.7 \log 0.7 + 0.2 \log 0.2 + 0.1 \log 0.1) \approx 0.802$.
• KL divergence: the difference, $0.887 - 0.802 = 0.085$.

If the model predicted $q = p$ exactly, the cross-entropy would drop to $0.802$ and the KL divergence to 0.

So minimizing the cross-entropy is the same as minimizing the KL divergence: the model is pulled towards the true distribution. In ordinary training the target is one-hot: the right class $y$ has $p_y = 1$ and all others 0. Then the sum has a single term, $-\log q_y$, which is the last row of the table above.

Two numbers help read a loss value:

• The uniform baseline. A model that knows nothing gives every class $1/C$, so its loss is $-\log(1/C) = \log C$. For the 50,257-token vocabulary of GPT-2 (Generative Pre-trained Transformer 2) that is $\log 50257 \approx 10.8$ nats (the unit of a loss computed with the natural log). A freshly initialized language model should start near this value; if yours starts far above it, the initialization is too large.
• Perplexity is $e^{\mathcal{L}}$: the number of equally likely tokens the model is, on average, hesitating between. A loss of 2.3 nats is a perplexity of about 10: as uncertain as picking among 10 tokens.

Cross-entropy from logits

A network does not output probabilities directly: it outputs logits $z$, one raw score per class, and softmax turns them into probabilities $q$ (see the NumPy section). Computing log(softmax(z)) is unsafe: softmax can round a tiny probability to exactly 0, and then the log is $-\infty$. Instead, use the LSE (log-sum-exp) function, which stays finite:

where:

• $\ell$ is the loss of one example;
• $z_y$ is the logit of the right class $y$;
• $\mathrm{LSE}(z) = \log \sum_c e^{z_c}$, computed with the stable trick of the NumPy section.

Example. Logits $z = (2, 1, 0)$, right class $y = 0$.

1. $\mathrm{LSE}(z) = \log(e^2 + e^1 + e^0) \approx 2.408$.
2. $\ell = 2.408 - z_0 = 2.408 - 2 = 0.408$.

Check: softmax gives $q \approx (0.665, 0.245, 0.090)$, and $-\log 0.665 \approx 0.408$.

The gradient: probabilities minus the target

The gradient of this loss with respect to the logits is remarkably simple:

where $q_c$ is the softmax probability of class $c$, and $[c = y]$ is 1 when $c$ is the right class and 0 otherwise. It comes from two facts: the derivative of $\mathrm{LSE}(z)$ with respect to $z_c$ is $q_c$, and the derivative of $z_y$ is 1 for $c = y$ and 0 for the others.

In words: the gradient is the probabilities minus the one-hot target.

Example. Same logits, $q \approx (0.665, 0.245, 0.090)$.

• Right class $y = 0$: gradient $(0.665 - 1,\ 0.245,\ 0.090) = (-0.335,\ 0.245,\ 0.090)$.
• Right class $y = 2$ instead: gradient $(0.665,\ 0.245,\ 0.090 - 1) = (0.665,\ 0.245,\ -0.910)$.

Training moves against the gradient, so the right class's logit goes up (its gradient is negative) and every other logit goes down, each in proportion to its probability. The worse the right answer's probability, the bigger the push.

The gradient of each example adds up to zero, and every entry stays between $-1$ and $1$, whatever the logits. This is why frameworks fuse softmax and cross-entropy into one operation: done separately, the softmax backward multiplies by tiny numbers and the log backward divides by them.

For a batch, the loss is the mean over the $N$ examples, so every example's gradient is also divided by $N$. Forgetting that factor makes the right learning rate silently depend on the batch size.

Binary cross-entropy from a logit

For a yes/no answer the network outputs one logit $x$, and the probability of "yes" is the sigmoid, which squashes any number into $(0, 1)$:

Example. $\sigma(2) \approx 0.881$, $\sigma(0) = 0.5$, $\sigma(-3) \approx 0.047$.

The binary cross-entropy, $-y \log q - (1 - y) \log(1 - q)$, breaks for large logits: $\sigma(40)$ rounds to exactly $1.0$, and then $\log(1 - q) = \log 0 = -\infty$. It can be rewritten without ever computing $q$:

where $y$ is the target (1 for yes, 0 for no) and $|x|$ is the absolute value of $x$.

Why this form is safe: the exponent $-|x|$ is never positive, so $e^{-|x|}$ is between 0 and 1 and nothing overflows; and $\log(1 + \text{something between 0 and 1})$ is never $\log 0$.

Example.

• $x = 2$, $y = 1$: $2 - 2 + \log(1 + e^{-2}) \approx 0.127$, the same as $-\log \sigma(2) = -\log 0.881$.
• $x = 40$, $y = 0$: the naive formula returns inf; the stable one gives $40 - 0 + \log(1 + e^{-40}) = 40$, the correct loss.

The gradient has the same shape as for cross-entropy, prediction minus target: $\partial \ell / \partial x = \sigma(x) - y$. The target $y$ may also be a soft target, any number between 0 and 1, such as $0.8$ for "probably yes".

The loss of a language model in practice

A language model predicts the next token at every position of every text in the batch: logits of shape $(B, T, V)$ (batch size, sequence length, vocabulary size) and targets of shape $(B, T)$. The loss is the mean cross-entropy over all positions, with two refinements.

Ignored positions

Some positions should not be learned: padding tokens, added only to make the texts of a batch the same length, and in fine-tuning the prompt tokens. Their target is set to a special value (PyTorch uses ignore_index=-100). These positions contribute nothing, and the mean is taken over the remaining positions only.

Example. A batch has 4 positions with losses $0.4$, $1.2$ and $0.8$, and one padding position. The loss is $(0.4 + 1.2 + 0.8) / 3 = 0.8$. Dividing by all 4 positions would give $0.6$: the loss would look better just because the batch had more padding.

Label smoothing

Label smoothing replaces the one-hot target with a softer one: $1 - \varepsilon$ on the right class, plus $\varepsilon$ spread evenly over all $C$ classes. The loss becomes

where $\varepsilon$ (epsilon), for example $0.1$, is the amount of smoothing. In words: mostly the usual loss, plus a small penalty for giving any class a very low probability.

Example. Same logits as above, $q \approx (0.665, 0.245, 0.090)$, $y = 0$, $\varepsilon = 0.1$, $C = 3$.

1. $-\log q_c$ for the three classes: $0.408$, $1.408$, $2.408$.
2. First part: $0.9 \cdot 0.408 \approx 0.367$.
3. Second part: $\tfrac{0.1}{3}(0.408 + 1.408 + 2.408) \approx 0.141$.
4. Total: $\ell \approx 0.508$, compared with $0.408$ without smoothing.

Label smoothing stops the model from pushing logits towards infinity to reach a probability of exactly 1, and it tends to improve calibration (Müller et al., below). Whether it helps depends on the task, but every framework exposes it, so it is worth knowing exactly what it computes.

References

• C. E. Shannon, A Mathematical Theory of Communication, Bell System Technical Journal 27(3), 379–423, 1948. Entropy and the information-theoretic view of coding.
• S. Kullback, R. A. Leibler, On Information and Sufficiency, Annals of Mathematical Statistics 22(1), 79–86, 1951. The divergence that cross-entropy minimizes.
• I. Goodfellow, Y. Bengio, A. Courville, Deep Learning, MIT Press, 2016. Section 3.13, information theory; section 5.5, maximum likelihood estimation; section 6.2.1.1, maximum likelihood and the cost functions of neural networks.
• C. Szegedy, V. Vanhoucke, S. Ioffe, J. Shlens, Z. Wojna, Rethinking the Inception Architecture for Computer Vision, CVPR, 2818–2826, 2016. Introduces label smoothing.
• R. Müller, S. Kornblith, G. Hinton, When Does Label Smoothing Help?, NeurIPS, 2019. Effects on calibration and on the learned representations.
• PyTorch documentation, torch.nn.CrossEntropyLoss: the exact semantics of ignore_index, label_smoothing and reduction that the last exercise reproduces.