Besides Cross Entropy, What Other Choices Are There for LM Loss?

By 苏剑林 | August 09, 2026

For a long time, Cross Entropy has been the standard loss function for LLM pre-training and fine-tuning. Can this "standard" be changed? If one wants to change it, what are the options? And what impact would such a change bring?

Many readers may have never seriously questioned these issues. On one hand, cross entropy is simple and effective, backed by information theory interpretations, making it feel so "natural" that we accept it readily. On the other hand, changing the loss function is a "move one hair and the whole body shakes" kind of affair—changing it means all loss-based comparisons are no longer valid, and we can only compare downstream task performance, which involves a massive amount of engineering work.

However, "natural" does not mean "only choice." Thinking through the underlying principles not only helps us better understand the model's optimization process but may also provide new perspectives for improving results.

Analysis

More accurately, cross entropy is the standard loss function for classification problems. Since LLM training appears to be a token-by-token classification problem, cross entropy loss was adopted. In that case, does it mean we can just swap it for any other classification loss function?

Unfortunately, that’s not quite right. Unlike conventional classification problems, the patterns of natural language are one-to-many. For example, "white-cut" (白切) can be followed not only by "chicken" (鸡) but also by "duck" (鸭), "dog" (狗), "lamb" (羊), etc. Therefore, we need to model the complete distribution rather than just predicting a single label. In other words, we need to estimate the probabilities of "chicken," "duck," "dog," and "lamb" following "white-cut," rather than just providing one correct answer—as the correct answer is not unique.

The core difficulty here is that the training corpus is fed into the model "in fragments." A "white-cut chicken" comes this time, a "white-cut duck" might come next time, and a "white-cut chicken" might appear again later. We cannot calculate the complete frequency distribution beforehand. This requires the loss function to have the ability to be transformed into a sampled form. In mathematical terms, this requires the loss function to be linear with respect to the target distribution $\boldsymbol{p}$, i.e.,

where $\boldsymbol{p}=(p_1,p_2,\cdots,p_n)$ and $\boldsymbol{q}=(q_1,q_2,\cdots,q_n)$ represent the target and predicted distributions, respectively. The condition after $\text{s.t.}$ indicates that for a fixed $\boldsymbol{p}$, the minimum of $L(\boldsymbol{p}, \boldsymbol{q})$ occurs at $\boldsymbol{q}^*=\boldsymbol{p}$, which is a basic requirement for a loss function. The linearity constraint excludes many common probability metrics, such as Total Variation:

To estimate this via sampling, we would need to calculate $|1 - q_i/p_i|$. However, $p_i$ is unknown to us beforehand, so Total Variation cannot be adapted for LLM training.

Derivation

Now let's solve for the objective $\eqref{eq:obj}$. According to the basic requirement of a loss function, for a fixed $\boldsymbol{p}$, the minimum of $L(\boldsymbol{p}, \boldsymbol{q})$ is $H(\boldsymbol{p})\triangleq L(\boldsymbol{p}, \boldsymbol{p})$. Thus we can write:

Now we focus our attention on the variable $\boldsymbol{p}$. Clearly, $L(\boldsymbol{p}, \boldsymbol{q})$ is linear with respect to $\boldsymbol{p}$. So, when we fix $\boldsymbol{q}$, the left side describes a hyperplane, and the right side describes a curved surface; they intersect at $\boldsymbol{p}=\boldsymbol{q}$. If we further assume, or require, that the minimum point is unique, then the intersection is actually a tangency, meaning $L(\boldsymbol{p}, \boldsymbol{q})$ is the tangent plane of $H(\boldsymbol{p})$ at $\boldsymbol{p}=\boldsymbol{q}$!

Since the above inequality holds universally, 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 determine that $H(\boldsymbol{p})$ is a concave function. Conversely, choosing any concave function $H(\boldsymbol{p})$ with respect to $\boldsymbol{p}$, its tangent plane at $\boldsymbol{p}=\boldsymbol{q}$ is:

where the equality uses the constraint $\boldsymbol{p},\boldsymbol{q}\in\Delta^{n-1}$ (the sum of components is 1), "$\cdot$" denotes the dot product, and vector-scalar additions are treated element-wise. Based on the previous derivation, the above expression is the $L(\boldsymbol{p}, \boldsymbol{q})$ we expect. Thus, we can directly read out:

This is the general form of $S(\boldsymbol{q},i)$ (allowing for 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 a one-hot vector with 1 at the $i$-th position. It is easy to see that $S(\boldsymbol{q},i)$ is linear with respect to $H(\boldsymbol{q})$, and since the linear interpolation of two concave functions is still a concave function, the linear interpolation of two scoring functions remains a scoring function.

Scoring

The results above actually have a specific name, called Proper Scoring Rules. We won't trace the origin of the name here, but below are several classic examples:

It is worth noting that the latter three scoring functions all degenerate into the logarithmic score as $\alpha\to 1$, meaning they are all generalizations of cross entropy. Specifically, 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 is not difficult to prove. In this case, $H(\boldsymbol{q}) = \sum_i q_i S(q_i)$. Substituting into Equation $\eqref{eq:S-q-i-1}$ gives:

The term $\sum_j q_j^2 S'(q_j)$ is equivalent to a constant for a single $q_i$. Thus, this equation is equivalent to $q_i S'(q_i) = -c$. Solving this yields $S(q_i) = - c \log q_i$, which gives the logarithmic score.

Note: The argument here is actually slightly lacking in rigor. Due to the constraint $\sum_i q_i = 1$, the idea of "depending only on $q_i$" is not so simple—for example, $S(q_n) = S(1 - q_1 - \cdots - q_{n-1})$. We cannot simply claim $S(q_n)$ only depends on $q_n$.

Similar confusion arises in the argument of "whether $\sum_j q_j^2 S'(q_j)$ depends on $q_i$." A more rigorous statement here is to require the first $n-1$ elements of $q_i$ to satisfy $q_i S'(q_i) = \sum_j q_j^2 S'(q_j)$. $q_n$ is eliminated via $q_n=1 - q_1 - \cdots - q_{n-1}$, leaving $n-1$ relatively independent variables. $q_i S'(q_i) = \sum_j q_j^2 S'(q_j)$ implies that the first $n-1$ values of $q_i S'(q_i)$ equal the same expression containing $q_1,\cdots,q_{n-1}$. But since $q_i S'(q_i)$ depends at most on $q_i$, this expression must be a constant, resulting in $q_i S'(q_i) = -c$.

These scoring functions can all be generalized to continuous distributions by replacing the discrete variable $i$ with a continuous variable $\boldsymbol{x}$ and the summation with an integral. However, the difficulty with continuous distributions is often the inability to calculate the normalization factor. These scoring functions require explicit probability densities, so they are usually not very "handy." In such scenarios, scoring functions that do not depend on the normalization factor are needed (usually involving gradients, a unique feature of continuous distributions), but we won't expand on that here.

Gradient

As just mentioned, the latter three scoring functions are generalizations of cross entropy. Intuitively, it seems that by fine-tuning $\alpha$, one could achieve improvements in downstream tasks. However, things are not that simple.

Typically, models can only predict unbounded Logits vectors $\boldsymbol{z}\in\mathbb{R}^n$. We need to apply an activation function to project this into a probability distribution $\boldsymbol{q}$, with Softmax being the standard choice. Since we can only use gradient-based optimizers, the convexity and gradient properties of the loss function with respect to $\boldsymbol{z}$ are especially important. Taking logarithmic and Brier scores as examples, the gradients of the loss with respect to $\boldsymbol{z}$ under the Softmax activation function are:

Clearly, the gradient of cross entropy appears much "cleaner." The gradient is zero if and only if $\boldsymbol{q}=\boldsymbol{e}_i$, which means as long as the target hasn't been reached, it provides an effective gradient. Furthermore, the further it is from the target, the larger the gradient. Squared loss has an additional transformation $\diag(\boldsymbol{q})-\boldsymbol{q}\boldsymbol{q}^\top$. When $\boldsymbol{q}=\boldsymbol{e}_j\neq \boldsymbol{e}_i$, this term also becomes zero, meaning the model suffers from gradient vanishing when it is "confidently wrong."

This characteristic is double-edged: in the early and middle stages, most of the model's predictions are inaccurate, meaning learning efficiency is extremely low with squared loss. However, in the late stages when the model is basically stable, samples that are still "confidently wrong" might be extremely difficult or erroneous samples. Skipping them might be more beneficial for the overall effect. Therefore, cross entropy has higher learning efficiency and should serve as the primary loss, but squared loss has better noise resistance and can be attempted in later stages.

These can also be understood through the convexity of $S(\boldsymbol{q},i)$ with respect to $\boldsymbol{z}$. It can be proven that under Softmax activation: cross entropy is convex with respect to $\boldsymbol{z}$, meaning the optimal point is unique, and gradients anywhere point toward the target point $\boldsymbol{e}_i$. However, squared loss is non-convex with respect to $\boldsymbol{z}$. These good properties are not guaranteed to exist, and the model might fall into a saturation dilemma where it is "both wrong and unable to escape."

Of course, even if it is convex with respect to $\boldsymbol{z}$, $\boldsymbol{z}$ still has parameters. In deep models, the loss function is generally non-convex with respect to parameters. Our requirement for convexity with respect to $\boldsymbol{z}$ is more about not adding "extra trouble"—optimizing deep learning is already hard enough; there's no need to add obstacles at the final layer.

Back-derivation

The results of the previous section all held a premise—Softmax activation. But what if it's not Softmax activation? Or conversely, given a specific scoring function, can we derive the "optimal" activation function suited for that scoring function?

The first question to consider is: how do we define "optimal"? Referring to the "Softmax + Cross Entropy" combination, it has two advantages: one is convexity with respect to $\boldsymbol{z}$, and the other is that the gradient $\nabla_{\boldsymbol{z}} S(\boldsymbol{q}, i) = \boldsymbol{q} - \boldsymbol{e}_i$ is "clean." In fact, the second point is stronger and more practical, so we start from there. That is, we hope to find a transformation $\boldsymbol{q} = \sigma(\boldsymbol{z})\in\Delta^{n-1}$ such that for a given $S(\boldsymbol{q}, i)$, it holds that:

Note that $\boldsymbol{e}_i= \nabla_{\boldsymbol{z}} z_i$, so the above expression can also be written as $\nabla_{\boldsymbol{z}} (S(\boldsymbol{q}, i) + z_i) = \boldsymbol{q}$. This suggests that there exists a scalar function $\Phi(\boldsymbol{z})$ independent of $i$ such that:

Rearranging gives $\Phi(\boldsymbol{z}) - z_i = S(\boldsymbol{q}, i)$. Multiplying both sides by $p_i$ and summing gives:

Further rearranging gives $\Phi(\boldsymbol{z}) \geq \boldsymbol{p}\cdot\boldsymbol{z} + H(\boldsymbol{p})$. Since this holds for any $\boldsymbol{p}$, $\Phi(\boldsymbol{z})$ is an upper bound for all $\boldsymbol{p}\cdot\boldsymbol{z} + H(\boldsymbol{p})$. Substituting $\boldsymbol{p}=\boldsymbol{q}$ into the equation gives $\Phi(\boldsymbol{z}) - \boldsymbol{q}\cdot\boldsymbol{z} = L(\boldsymbol{q},\boldsymbol{q}) = H(\boldsymbol{q})$. That is, equality is achieved when $\boldsymbol{p}=\boldsymbol{q}$. Thus, $\Phi(\boldsymbol{z})$ is the "supremum" of all $\boldsymbol{p}\cdot\boldsymbol{z} + H(\boldsymbol{p})$, namely:

This is exactly the convex conjugate of the convex function $-H(\boldsymbol{p})$. The contents of these sections combined happen to be the classic Fenchel-Young Losses framework.

Activation

In this section, we also derive the optimal activation functions corresponding to the scoring functions mentioned earlier. It’s easy to see that their corresponding $H(\boldsymbol{q})$ all share the same structure $g(\sum_i q_i^{\alpha})$, so we can solve them uniformly. Let $t = \sum_i q_i^{\alpha}$, then $H(\boldsymbol{q})=g(t)$, where

Introducing the Lagrangian $\boldsymbol{q}\cdot \boldsymbol{z} + H(\boldsymbol{q}) - \lambda(\sum_i q_i - 1)$, taking the derivative with respect to $q_i$ and setting it to 0 gives:

For a single component, $\lambda$ and $\alpha g'(t)$ are shared "constants." We simply adjust these two constants so that $\boldsymbol{q}$ becomes a suitable distribution. Note that $q_i$ can only be $q_i > 0$ or $q_i = 0$. The latter is trivial, so we only need to analyze the former. Notice that for $\alpha > 1$, $g'(t) < 0$, and for $\alpha < 1$, $g'(t) > 0$. Therefore, we can write:

where $[x]_+ = \max(x, 0)$. However, for $\alpha < 1$, this truncation is redundant because the exponent is less than 0, and a negative power of zero is undefined, so the truncation will never be triggered. This also shows that $\alpha > 1$ corresponds to sparse distributions, while $\alpha < 1$ leads to dense distributions. $\lambda$ is determined by the condition $\sum_i q_i = 1$. Clearly, when $\alpha\to 1$, it is the classic Softmax. When $\alpha\neq 1$, for the Tsallis score, $g'(t)=1/(1-\alpha)$, thus:

$\lambda$ can be solved using the bisection method through the condition $\sum_i q_i = 1$. When $\alpha=2$, the result is exactly Sparsemax. Other cases are called Entmax-$\alpha$. When $\alpha=2$ and $\alpha=1.5$, $\lambda$ has exact solutions more efficient than bisection, which we previously introduced in "Path to Probability Distributions: Softmax and Its Alternatives". As for the results of other scoring rules, they are slightly more complex and left for the reader to try.

Summary

This article derived the general construction of LM Loss from the perspectives of "learning distributions" and "allowing sampling." Subsequently, combined with the activation functions used for predictive distributions, we calculated the gradients and convexity of these losses to simply judge their pros and cons. Finally, we attempted to start from a given loss and back-derive its optimal matching activation function. The optimal activation function for cross entropy is exactly Softmax—explaining why these two almost always appear together.