What Alternatives to Cross-Entropy Exist for LM Loss?
For a long time, cross-entropy has been the standard loss function for LLM pretraining and fine-tuning. Can this "standard" be changed? If so, what alternatives are available? And what would the consequences of switching be?
Many readers may never have carefully thought through these questions. On one hand, cross-entropy is simple and effective, and backed by an information-theoretic interpretation, which makes it feel so "obviously correct" that we accept it without much thought. On the other hand, changing the loss function is a matter that "pulls on one thread and moves the whole fabric"—once you change it, all loss-based comparisons become invalid, and you're forced to compare downstream task performance instead, which is a huge amount of engineering work.
However, "obviously correct" is not the same as "the only choice." Thinking through the underlying principles carefully not only helps us better understand the model's optimization process, but may also offer a new perspective for improving performance.
Analysis
To be more precise, cross-entropy is the standard loss function for classification problems, and LLM training looks like a token-by-token classification problem, which is why cross-entropy has been carried over. Given this, does it seem like we could just swap in any other classification loss function?
Unfortunately, no. Unlike conventional classification problems, the patterns of natural language are one-to-many. For instance, after "white-cut" (a Chinese cooking style), the next word could be "chicken," "duck," "dog," "sheep," and so on. What we need to model is the full distribution, not just predict a single label. In other words, we need to estimate the probabilities of "chicken," "duck," "dog," "sheep," etc. all following "white-cut," rather than simply giving one correct answer—there isn't a unique correct answer.
The core difficulty here is that training corpora are fed into the model in a "scattered" fashion: this time we get "white-cut chicken," next time maybe "white-cut duck," and later maybe "white-cut chicken" again. We cannot compute the full frequency distribution in advance, which requires the loss function to be expressible in a form amenable to sampling estimation. Mathematically, this requires the loss function to be linear in the target distribution $\boldsymbol{p}$, i.e.,
\begin{equation}\newcommand{argmin}{\mathop{\text{argmin}}} L(\boldsymbol{p}, \boldsymbol{q}) = \sum_{i=1}^n p_i S(\boldsymbol{q}, i) = \mathbb{E}_{i\sim\boldsymbol{p}}[S(\boldsymbol{q}, i)]\qquad\text{s.t.}\qquad \boldsymbol{p} = \argmin_{\boldsymbol{q}\in\Delta^{n-1}} L(\boldsymbol{p}, \boldsymbol{q})\label{eq:obj}\end{equation}
where $\boldsymbol{p}=(p_1,p_2,\cdots,p_n),\boldsymbol{q}=(q_1,q_2,\cdots,q_n)$ denote the target distribution and the predicted distribution respectively, and the condition after $\text{s.t.}$ indicates that once $\boldsymbol{p}$ is fixed, the minimum of $L(\boldsymbol{p}, \boldsymbol{q})$ is attained at $\boldsymbol{q}^*=\boldsymbol{p}$—this is the basic requirement for a loss function. The linearity constraint rules out many common probability measures, such as Total Variation:
\begin{equation}TV(\boldsymbol{p}, \boldsymbol{q}) = \sum_{i=1}^n |p_i - q_i| = \sum_{i=1}^n p_i\left|1 - \frac{q_i}{p_i}\right| = \mathbb{E}_{i\sim \boldsymbol{p}}\left[\left|1 - \frac{q_i}{p_i}\right|\right]\end{equation}
To estimate it via sampling, we would need to compute $|1 - q_i/p_i|$, but $p_i$ is something we cannot know in advance, so Total Variation cannot be adapted into a form usable for LLM training.
Derivation
Now let's solve for the objective $\eqref{eq:obj}$. Based on the basic requirement for loss functions, once $\boldsymbol{p}$ is fixed, the minimum of $L(\boldsymbol{p}, \boldsymbol{q})$ is $H(\boldsymbol{p})\triangleq L(\boldsymbol{p}, \boldsymbol{p})$, so we can write
\begin{equation}L(\boldsymbol{p}, \boldsymbol{q})\geq H(\boldsymbol{p})\end{equation}
Now let's focus our attention on the variable $\boldsymbol{p}$. Clearly $L(\boldsymbol{p}, \boldsymbol{q})$ is linear in $\boldsymbol{p}$, so when we fix $\boldsymbol{q}$, the left-hand side describes a hyperplane while the right-hand side describes a hypersurface, and they intersect at $\boldsymbol{p}=\boldsymbol{q}$. If we further assume, or require, that the minimum point is unique, then this intersection is actually tangency—that is, $L(\boldsymbol{p}, \boldsymbol{q})$ is the tangent plane of $H(\boldsymbol{p})$ at $\boldsymbol{p}=\boldsymbol{q}$!
Since the inequality above holds identically, this is equivalent to saying that $H(\boldsymbol{p})$ always lies below its tangent plane, which is exactly the definition of a concave function. So we can conclude that $H(\boldsymbol{p})$ is a concave function. Conversely, if we pick any concave function $H(\boldsymbol{p})$ of $\boldsymbol{p}$, its tangent plane at $\boldsymbol{p}=\boldsymbol{q}$ is
\begin{equation}H(\boldsymbol{q}) + (\boldsymbol{p}-\boldsymbol{q})\cdot\nabla_{\boldsymbol{q}} H(\boldsymbol{q}) = \boldsymbol{p}\cdot\big[H(\boldsymbol{q}) + \nabla_{\boldsymbol{q}} H(\boldsymbol{q}) - \boldsymbol{q}\cdot\nabla_{\boldsymbol{q}} H(\boldsymbol{q})\big]\end{equation}
where the equality uses the constraint $\boldsymbol{p},\boldsymbol{q}\in\Delta^{n-1}$ (that the components sum to 1), "$\cdot$" denotes the inner product, and adding a vector and a scalar is understood element-wise. Following the earlier derivation, the expression above is exactly the $L(\boldsymbol{p}, \boldsymbol{q})$ we're looking for, so we can directly read off
\begin{align}S(\boldsymbol{q},i) =&\, H(\boldsymbol{q}) + \partial_i H(\boldsymbol{q}) - \boldsymbol{q}\cdot\nabla_{\boldsymbol{q}} H(\boldsymbol{q}) \label{eq:S-q-i-1} \\[5pt] =&\, H(\boldsymbol{q}) + (\boldsymbol{e}_i - \boldsymbol{q})\cdot\nabla_{\boldsymbol{q}} H(\boldsymbol{q})\label{eq:S-q-i-2}\end{align}
This is the general form of $S(\boldsymbol{q},i)$ (allowing addition/subtraction of a constant and multiplication by a positive constant), where $\partial_i H(\boldsymbol{q})$ denotes the $i$-th component of $\nabla_{\boldsymbol{q}} H(\boldsymbol{q})$, and $\boldsymbol{e}_i$ is the one-hot vector with a 1 in position $i$. It's easy to see that $S(\boldsymbol{q},i)$ is linear in $H(\boldsymbol{q})$, and since a linear interpolation of two concave functions is still concave, a linear interpolation of two scoring functions is likewise still a scoring function.
Scoring
The result above actually has a dedicated name: Proper Scoring Rules. We won't trace the origin of this name, but here are a few classic examples:
$$\newcommand{\rs}{\rule[-1.2ex]{0pt}{3.5ex}} \begin{array}{c|c|c} \hline \rs\text{name} & H(\boldsymbol{p}) & S(\boldsymbol{q},i) \\ \hline \rs\text{log score (cross entropy)} & -\sum_i p_i\log p_i & -\log q_i \\ \hline \rs\text{Brier score (squared loss)} & 1-\sum_i p_i^2 & \|\boldsymbol{q}-\boldsymbol{e}_i\|^2 \\ \hline \rs\text{Tsallis score}(\alpha > 0) & \frac{1-\sum_i p_i^\alpha}{\alpha-1} & \sum_j q_j^\alpha - \frac{\alpha}{\alpha-1}q_i^{\alpha-1} + \frac{1}{\alpha-1} \\ \hline \rs\text{spherical score}(\alpha > 0) & \frac{1-\Vert\boldsymbol{p}\Vert_\alpha}{\alpha-1} & \frac{1}{\alpha-1}\left(1 -\frac{q_i^{\alpha-1}}{\Vert\boldsymbol{q}\Vert_\alpha^{\alpha-1}}\right) \\ \hline \rs\text{Rényi score}(0 < \alpha < 1) & \frac{1}{1-\alpha}\log\sum_i p_i^\alpha & \frac{1}{1-\alpha}\left(\log\sum_j q_j^\alpha+\alpha \frac{q_i^{\alpha-1}}{\sum_j q_j^\alpha}-\alpha\right) \\ \hline \end{array}$$
It's worth pointing out that the last three scoring functions all degenerate to the logarithmic score when $\alpha\to 1$, meaning they are all generalizations of cross-entropy in some sense. In particular, if we further require that $S(\boldsymbol{q},i)$ depends only on $q_i$, i.e., $S(\boldsymbol{q},i)=S(q_i)$, then cross-entropy $-\log q_i$ is the only choice. This isn't hard to prove: in that case $H(\boldsymbol{q}) = \sum_i q_i S(q_i)$, and substituting into $\eqref{eq:S-q-i-1}$ gives
\begin{equation}\require{cancel}\cancel{S(q_i)} = \cancel{S(q_i)} + q_i S'(q_i) - \sum_j q_j^2 S'(q_j)\end{equation}
where the term $\sum_j q_j^2 S'(q_j)$ is also effectively constant for a single $q_i$, so this equation is equivalent to $q_i S'(q_i) = -c$, which is easily solved to give $S(q_i) = - c \log q_i$—yielding the logarithmic score.
Note: The argument here is actually a bit less than fully rigorous. Because of the constraint $\sum_i q_i = 1$, "depending only on $q_i$" is not as straightforward a notion as it seems—for instance $S(q_n) = S(1 - q_1 - \cdots - q_{n-1})$—so we can't simply claim that $S(q_n)$ depends only on $q_n$.
A similar subtlety arises when arguing whether $\sum_j q_j^2 S'(q_j)$ depends on $q_i$. The more rigorous statement here is: we only require that the first $n-1$ instances of $q_i$ satisfy $q_i S'(q_i) = \sum_j q_j^2 S'(q_j)$, with $q_n$ eliminated via $q_n=1 - q_1 - \cdots - q_{n-1}$, leaving only $n-1$ relatively independent variables. $q_i S'(q_i) = \sum_j q_j^2 S'(q_j)$ then implies that the first $n-1$ instances of $q_i S'(q_i)$ all equal the same expression involving $q_1,\cdots,q_{n-1}$, but $q_i S'(q_i)$ can depend on at most $q_i$, so this expression must be a constant, giving us $q_i S'(q_i) = -c$.
All these scoring functions can be extended to continuous distributions, simply by replacing the discrete variable $i$ with a continuous variable $\boldsymbol{x}$ and replacing the sum with an integral. But the difficulty with continuous distributions is usually the intractability of the normalizing constant—these scoring functions all require an explicit probability density, so they're often not so "convenient" in that setting. This is where scoring functions that don't rely on the normalizing constant become more useful (typically relying on gradients, which is a feature unique to the continuous case), but we won't get into that here.
Gradients
We just noted that the last three scoring functions are generalizations of cross-entropy, so intuitively, could fine-tuning $\alpha$ give us an edge on downstream tasks? Unfortunately, it's not that simple.
In general, the model can only predict an unbounded logits vector $\boldsymbol{z}\in\mathbb{R}^n$, which we need to pass through an activation function to project it into a probability distribution $\boldsymbol{q}$, and the usual choice of activation function is Softmax. Since we can only use gradient-based optimizers, the convexity and gradient properties of the loss function with respect to $\boldsymbol{z}$ become especially important. Taking the logarithmic score and the Brier score as examples, under the Softmax activation, the gradients of the loss with respect to $\boldsymbol{z}$ are respectively
\begin{align}\newcommand{diag}{\mathop{\text{diag}}} \text{log score (cross entropy):}&\qquad \nabla_{\boldsymbol{z}} S(\boldsymbol{q}, i) = \boldsymbol{q} - \boldsymbol{e}_i \\[5pt] \text{Brier score (squared loss):}&\qquad \nabla_{\boldsymbol{z}} S(\boldsymbol{q}, i) = 2(\diag(\boldsymbol{q})-\boldsymbol{q}\boldsymbol{q}^\top)(\boldsymbol{q}-\boldsymbol{e}_i) \\ \end{align}
Clearly, cross-entropy's gradient looks "cleaner": it vanishes if and only if $\boldsymbol{q}=\boldsymbol{e}_i$, which means that as long as the target hasn't been reached, it will keep providing a useful gradient, and the farther away we are from the target, the larger the gradient. The squared loss has an extra factor $\diag(\boldsymbol{q})-\boldsymbol{q}\boldsymbol{q}^\top$, and when $\boldsymbol{q}=\boldsymbol{e}_j\neq \boldsymbol{e}_i$, this term also vanishes, meaning that when the model is "confidently wrong," its gradient can vanish too.
This property cuts both ways: in the early-to-middle stages of training, most of the model's predictions are inaccurate, meaning learning with squared loss would be extremely inefficient; but in later stages, the model has largely stabilized, and the cases where it's still "confidently wrong" are likely to be extremely hard or mislabeled samples—skipping them may actually be more beneficial to overall performance. So cross-entropy has higher learning efficiency and should serve as the primary loss, but squared loss has better robustness to noise, and could be worth trying in later stages of training.
These properties can also be understood via the convexity of $S(\boldsymbol{q},i)$ with respect to $\boldsymbol{z}$. One can show that, under Softmax activation: cross-entropy is convex in $\boldsymbol{z}$, which means the optimum is unique, and the gradient at any point points toward the target $\boldsymbol{e}_i$; but squared loss is non-convex in $\boldsymbol{z}$, so these nice properties are no longer guaranteed, and the model may fall into a saturated trap of being "both wrong and unable to escape."
Of course, even if the loss is convex in $\boldsymbol{z}$, $\boldsymbol{z}$ is itself a function of the model parameters, and in deep models the loss is generally non-convex with respect to the parameters regardless. Requiring convexity in $\boldsymbol{z}$ is more about not adding unnecessary "trouble" for the model on top of everything else—deep learning optimization is already hard enough without introducing extra obstacles right at the last layer.
Working Backwards
The results in the previous section all had one precondition: Softmax activation. But what if the activation isn't Softmax? Or conversely, given a scoring function, can we derive the "optimal" activation function that pairs with it?
The first question to consider is: how should "optimal" be defined? Looking at the "Softmax + cross-entropy" combination, it has two nice properties: first, convexity in $\boldsymbol{z}$, and second, a relatively "clean" gradient $\nabla_{\boldsymbol{z}} S(\boldsymbol{q}, i) = \boldsymbol{q} - \boldsymbol{e}_i$. In practice, the second property is stronger and more useful, so let's start from there. That is, we want to find a transformation $\boldsymbol{q} = \sigma(\boldsymbol{z})\in\Delta^{n-1}$ such that, for a given $S(\boldsymbol{q}, i)$,
\begin{equation}\nabla_{\boldsymbol{z}} S(\boldsymbol{q}, i) = \boldsymbol{q} - \boldsymbol{e}_i\end{equation}
holds. Note that $\boldsymbol{e}_i= \nabla_{\boldsymbol{z}} z_i$, so the equation above can also be written as $\nabla_{\boldsymbol{z}} (S(\boldsymbol{q}, i) + z_i) = \boldsymbol{q}$, which means there exists some scalar function $\Phi(\boldsymbol{z})$, independent of $i$, such that
\begin{equation}\Phi(\boldsymbol{z}) = S(\boldsymbol{q}, i) + z_i,\qquad \boldsymbol{q} = \nabla_{\boldsymbol{z}}\Phi(\boldsymbol{z})\end{equation}
Rearranging gives $\Phi(\boldsymbol{z}) - z_i = S(\boldsymbol{q}, i)$; multiplying both sides by $p_i$ and summing gives
\begin{equation}\Phi(\boldsymbol{z}) - \boldsymbol{p}\cdot\boldsymbol{z} = L(\boldsymbol{p},\boldsymbol{q})\geq H(\boldsymbol{p})\end{equation}
Rearranging further gives $\Phi(\boldsymbol{z}) \geq \boldsymbol{p}\cdot\boldsymbol{z} + H(\boldsymbol{p})$, which holds for any $\boldsymbol{p}$, so $\Phi(\boldsymbol{z})$ is an upper bound over all $\boldsymbol{p}\cdot\boldsymbol{z} + H(\boldsymbol{p})$. Then substituting $\boldsymbol{p}=\boldsymbol{q}$ into the above gives $\Phi(\boldsymbol{z}) - \boldsymbol{q}\cdot\boldsymbol{z} = L(\boldsymbol{q},\boldsymbol{q}) = H(\boldsymbol{q})$, i.e., equality is attained at $\boldsymbol{p}=\boldsymbol{q}$, so $\Phi(\boldsymbol{z})$ is the "supremum" over all $\boldsymbol{p}\cdot\boldsymbol{z} + H(\boldsymbol{p})$, i.e.,
\begin{equation}\newcommand{argmax}{\mathop{\text{argmax}}}\Phi(\boldsymbol{z}) = \max_{\boldsymbol{p}\in\Delta^{n-1}} \boldsymbol{p}\cdot \boldsymbol{z} + H(\boldsymbol{p}),\qquad \boldsymbol{q} = \argmax_{\boldsymbol{p}\in\Delta^{n-1}} \boldsymbol{p}\cdot \boldsymbol{z} + H(\boldsymbol{p})\end{equation}
This is exactly the convex conjugate of the convex function $-H(\boldsymbol{p})$, and together, the content of the last few sections is precisely the classic Fenchel-Young Losses framework.
Activations
In this section, we'll similarly work out the optimal activation functions corresponding to the scoring functions mentioned earlier. It's not hard to see that they all share the same structure for $H(\boldsymbol{q})$, namely $g(\sum_i q_i^{\alpha})$, so we can solve for them all at once. Let $t = \sum_i q_i^{\alpha}$, so that $H(\boldsymbol{q})=g(t)$, where
\begin{equation}g(t)=\frac{1-t}{\alpha-1}\ (\text{Tsallis/Brier}),\quad g(t)=\frac{1-t^{1/\alpha}}{\alpha-1}\ (\text{sphere}),\quad g(t)=\frac{\log t}{1-\alpha}\ (\text{Rényi}) \end{equation}
Introducing the Lagrangian $\boldsymbol{q}\cdot \boldsymbol{z} + H(\boldsymbol{q}) - \lambda(\sum_i q_i - 1)$, differentiating with respect to $q_i$ and setting the result to zero gives
\begin{equation}z_i + \alpha g'(t) q_i^{\alpha-1} = \lambda \qquad\Rightarrow\qquad q_i^{\alpha-1} = \frac{\lambda - z_i}{\alpha g'(t)}\end{equation}
For a single component, $\lambda$ and $\alpha g'(t)$ are both shared "constants," and what we need to do is tune these two constants so that $\boldsymbol{q}$ becomes a valid distribution. Note that $q_i$ can only take one of two forms, $q_i > 0$ or $q_i = 0$, the latter being trivial, so we only need to analyze the former. Also note that when $\alpha > 1$, we have $g'(t) < 0$, and when $\alpha < 1$, $g'(t) > 0$, so we can write
\begin{equation}q_i = \left\{\begin{aligned} &\, e^{z_i - \lambda},&\, \alpha \to 1 \\ &\,\left[\frac{z_i - \lambda}{-\alpha g'(t)}\right]_+^{\frac{1}{\alpha-1}},&\, \alpha \neq 1 \\ \end{aligned}\right.\end{equation}
where $[x]_+ = \max(x, 0)$, but when $\alpha < 1$ this truncation is redundant, because in that case the exponent is negative and a zero raised to a negative power is undefined, so the truncation never actually activates—this also shows that $\alpha > 1$ corresponds to a sparse distribution, while $\alpha < 1$ corresponds to a dense distribution; as for $\lambda$, it's determined by the equation $\sum_i q_i = 1$. Clearly, when $\alpha\to 1$ this is exactly the classic Softmax. When $\alpha\neq 1$, for the Tsallis score we have $g'(t)=1/(1-\alpha)$, so
\begin{equation}q_i = \left[\frac{\alpha-1}{\alpha}(z_i - \lambda)\right]_+^{\frac{1}{\alpha-1}}\end{equation}
$\lambda$ can be solved via bisection using the condition $\sum_i q_i = 1$. When $\alpha=2$, the result is exactly Sparsemax; the other cases are called Entmax-$\alpha$. When $\alpha=2$ and $\alpha=1.5$, $\lambda$ admits more efficient exact solutions than bisection, which we've also covered in The Road to Probability Distributions: A Survey of Softmax and Its Alternatives. As for the remaining scoring functions, the results are a bit more involved, and we'll leave those for readers to try themselves.
Summary
This post started from two angles—"learning a distribution" and "allowing sampling"—to derive a general construction for LM loss functions. We then combined this with the activation functions used for the predicted distribution, computed the gradients and convexity of these losses, and used this to briefly assess their relative merits. Finally, we tried working backward from a given loss to derive its matching optimal activation function, and found that the optimal activation function for cross-entropy is precisely Softmax—which explains why the two almost always appear together.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.