Loss Function

Cross-Entropy Loss

Cross-entropy loss measures the difference between two probability distributions, and is the standard loss for Classification problems. LLMs use this loss to train on predicting next tokens.

Resources:

• https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#torch.nn.CrossEntropyLoss

The cross-entropy $\mathbb{H}(p,q)$ between a “true” distribution $p$ and an estimated distribution $q$ is defined as:

$$\begin{array}{rl}\mathbb{H}(p,q)& =-{\mathbb{E}}_{x\sim p}[\log q(x)]\\ & =-\sum _{x}p(x)\log q(x)\end{array}$$

Intuition

Cross-entropy penalizes confident wrong predictions catastrophically. As the predicted probability of the true class approaches zero, $-\log q(y)\to \infty$. A mildly-wrong hedge ($q(y)=0.3$) costs $\approx 1.2$, a confidently-wrong prediction ($q(y)=0.01$) costs $\approx 4.6$, and being sure the true class has probability zero is infinitely bad. This is what forces models to be calibrated: they cannot cheaply bluff.

Where does this log come from?

This comes from Shannon Entropy.

He proved that under certain very reasonable assumptions, log is the only possible choice.

For our ML classification problems, let $y$ be the correct class. We can simplify the cross-entropy equation $-\sum _{x}p(x)\log q(x)\to -\log q(y)$

• $p(x)$ vanishes into a One-Hot distribution (Kronecker Delta), since we have $p(x)=\{\begin{array}{ll}1& x=y\\ 0& x\ne y\end{array}$ So we have $-\sum _{x}p(x)\log q(x)=-p(y)\log q(y)=-\log q(y)$

PyTorch Cross-Entropy Loss

https://docs.pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html $\ell (x,y)=L=\{{\ell }_1,\dots ,{\ell }_N{\}}^{\top }$

• $N$ spans the minibatch dimension (i.e. there are $N$ samples in a batch). If this is already confusing to you, see a more basic example in L1 Loss.

There are 2 different formulations depending on how the classes are predicted.

where:

• $n$ is the $n$-th example
• $N$ is the number dimensions
• $y_n$ is the correct class for the $n$-th example, $y_n\in [0,C)$
• $l_n$ is the loss for the $n$-th example
• $w_{y_n}$ is an optional weight matrix for the class $y_n$
• $x_{n,c}$ is the value of the logit at index $c$ for the $n$-th example

If you use indices as the target: where

$${\ell }_n=-w_{y_n}\cdot \log \left(\frac{\exp (x_{n,y_n})}{{\sum }_{c=1}^C\exp (x_{n,c})}\right)\cdot 1\{y_n\ne \text{ignore\_index}\}$$

$\text{softmax}(x_n)=\frac{\exp (x_{n,y_n})}{{\sum }_{c=1}^C\exp (x_{n,c})}$

${\ell }_n=-w_{y_n}\log \left(\text{softmax}(x_{n,y_n})\right)$

• The Negative Log Likelihood of the correct output

Summary

Model outputs logits → softmax(logits) → log(softmax(logits))

The cross entropy loss is calculated over the softmax.

Exercise:

• Help explain how we derive the Behavior Cloning loss, and how we end up using the Cross-Entropy Loss, which is where the log comes from

Related

• KL Divergence