A Simple Scheme for Mitigating the Overconfidence of Cross-Entropy

缓解交叉熵过度自信的一个简明方案 苏剑林 · 2023-03-14 · Machine Learning optimization loss function smoothness

As is well known, the standard evaluation metric for classification problems is accuracy, while the standard loss function is cross-entropy. Cross-entropy has the advantage of fast convergence, but it is not a smooth approximation of accuracy, which creates an inconsistency between training and prediction. On the other hand, when the predicted probability of a training sample is very low, cross-entropy assigns it an enormous loss (tending toward $-\log 0^{+}=\infty$), which means cross-entropy pays disproportionate attention to low-probability samples — even if such a sample might just be "dirty data." As a result, models trained with cross-entropy tend to become overconfident, i.e., they assign high predicted probability to nearly every sample, which brings two side effects: first, degraded performance due to overfitting on dirty data, and second, the predicted probabilities no longer serve as a good indicator of uncertainty.

The academic community has been continuously producing improvements around cross-entropy, and this line of research is currently in a state of "each immortal crossing the sea by their own means" — that is, there's no single standard answer. In this post, we'll look at yet another simple candidate solution to this problem, proposed in the paper Tailoring Language Generation Models under Total Variation Distance. more

Overview of the Result

As the name suggests, the modification in the original paper targets text generation tasks, with its theoretical basis in the Total Variation distance (see Designing GANs: Yet Another GAN Production Workshop). But in fact, after a series of relaxations and simplifications in the original paper, the final result no longer has any obvious connection to the Total Variation distance, and in theory it isn't limited to text generation tasks either. So in this post, we'll treat it as a loss function for general classification tasks.

For a data pair $(x,y)$, the loss function given by cross-entropy is

\begin{equation}-\log p_{\theta}(y|x)\end{equation}

The modification in the original paper is simple — it becomes

\begin{equation}-\frac{\log \big[\gamma + (1 - \gamma)p_{\theta}(y|x)\big]}{1-\gamma}\label{eq:gamma-ce}\end{equation}

where $\gamma\in[0,1]$. When $\gamma=0$, this reduces to ordinary cross-entropy; when $\gamma=1$, taking the limit, the result is $-p_{\theta}(y|x)$.

In the experiments of the original paper, the choice of $\gamma$ varies quite a bit across different tasks: for the language modeling task it's set to $\gamma=10^{-7}$, for machine translation it's $\gamma=0.1$, and for text summarization it's $\gamma=0.8$. One rule of thumb worth referencing is: if training from scratch, it's better to choose a $\gamma$ closer to 0; if fine-tuning, a relatively larger $\gamma$ can be considered. There's also a more intuitive approach, which is to treat $\gamma$ as a dynamic parameter, starting from $\gamma=0$ and gradually shifting toward $\gamma=1$ as training progresses — though this adds an extra schedule to tune.

In terms of performance, since there's now an extra tunable parameter $\gamma$, and ordinary cross-entropy is included as a special case, as long as you tune carefully there's generally a good chance of getting better results than plain cross-entropy — so this isn't something to worry too much about.

My Own Derivation

How should we understand equation $\eqref{eq:gamma-ce}$? In the "accuracy" section of Miscellaneous Notes on Function Smoothing: Differentiable Approximations of Non-Differentiable Functions, we derived that the smooth approximation of accuracy is

\begin{equation}\mathbb{E}_{(x,y)\sim \mathcal{D}}[p_{\theta}(y|x)]\end{equation}

So, if our evaluation metric is accuracy, then intuitively we should use $-p_{\theta}(y|x)$ as the loss function, since in that case the loss function's behavior would track the accuracy more closely. However, in practice cross-entropy tends to perform better. But the starting point of cross-entropy is merely "training more effectively," so sometimes it can "overshoot," leading to overfitting. So a natural idea is to see whether we can "interpolate" between the two results, so as to combine the strengths of both.

To this end, let's consider the gradients of both [where "accuracy" refers to its negative smooth approximation $-p_{\theta}(y|x)$]:

\begin{equation}\begin{aligned} \text{accuracy:}&\,\quad-\nabla_{\theta} p_{\theta}(y|x) \\ \text{cross entropy:}&\,\quad-\frac{1}{p_{\theta}(y|x)}\nabla_{\theta} p_{\theta}(y|x) \end{aligned}\end{equation}

The two differ only by a factor of $\frac{1}{p_{\theta}(y|x)}$. How do we turn $\frac{1}{p_{\theta}(y|x)}$ into 1? The scheme in the original paper is:

\begin{equation}\frac{1}{\gamma + (1 - \gamma)p_{\theta}(y|x)}\end{equation}

Of course this construction isn't unique, but the one chosen in the original paper preserves the gradient characteristics of cross-entropy as much as possible, and thereby preserves as much as possible cross-entropy's fast-convergence property. Based on this construction, we then want the gradient of the new loss function to be

\begin{equation}-\frac{\nabla_{\theta}p_{\theta}(y|x)}{\gamma + (1 - \gamma)p_{\theta}(y|x)} = \nabla_{\theta}\left(-\frac{\log \big[\gamma + (1 - \gamma)p_{\theta}(y|x)\big]}{1-\gamma}\right)\label{eq:gamma-ce-g}\end{equation}

This is how the loss function $\eqref{eq:gamma-ce}$ is found — in this process, we first designed the new gradient, and then found the corresponding loss function by integrating to recover the original function.

A Few More Remarks

Why design loss functions starting from the gradient? There are roughly two reasons.

First, many loss functions simplify considerably once their gradient is taken, so designing in gradient space often gives more inspiration and more degrees of freedom. For instance, in the example in this post, designing the transition function $\frac{1}{\gamma + (1 - \gamma)p_{\theta}(y|x)}$ between $\frac{1}{p_{\theta}(y|x)}$ and $1$ in gradient space isn't too complicated, but designing the transition function $\frac{\log \big[\gamma + (1 - \gamma)p_{\theta}(y|x)\big]}{1-\gamma}$ between $p_{\theta}(y|x)$ and $\log p_{\theta}(y|x)$ directly in the space of loss functions would be much more complicated.

Second, the optimizers we currently use are all gradient-based, so oftentimes it's enough to just design the gradient well — we don't even need to find the original function. The original result in the paper actually only gives the gradient:

\begin{equation}-\max\left(b, \frac{p_{\theta}(y|x)}{\gamma + (1 - \gamma)p_{\theta}(y|x)}\right)\nabla_{\theta}\log p_{\theta}(y|x)\end{equation}

When $b=0$, this is equivalent to equation $\eqref{eq:gamma-ce}$. In other words, the original paper also added a threshold when designing the gradient, at which point it becomes hard to write down a simple original function. But the expression above isn't difficult to implement — we just need to consider the loss function

\begin{equation}-\max\left(b, \frac{p_{\theta}(y|x)}{\gamma + (1 - \gamma)p_{\theta}(y|x)}\right)_{\text{stop_grad}}\log p_{\theta}(y|x)\end{equation}

Here $\text{stop_grad}$ simply means directly cutting off the gradient of this part of the result, which corresponds to the tf.stop_gradient operator in TensorFlow.

Summary

This post mainly introduced a simple scheme for mitigating the overconfidence of cross-entropy.