Generalizing "Softmax + Cross-Entropy" to Multi-Label Classification Problems

将“Softmax+交叉熵”推广到多标签分类问题 苏剑林 · 2020-04-25 · Machine Learning optimization loss function smoothness

(Note: the contents of this article have been organized into a paper, ZLPR: A Novel Loss for Multi-label Classification. If you wish to cite this work, please cite the English paper directly. Thanks.)

Generally, when dealing with a conventional multi-class classification problem, we use a fully-connected layer at the end of the model to output a score for each class, apply a softmax activation, and use cross-entropy as the loss function. In this article, we attempt to generalize the "Softmax + cross-entropy" scheme to the multi-label classification setting, hoping to obtain a loss suitable for multi-label classification tasks that does not require special tuning of class weights or thresholds.

Class imbalance

From Single-Label to Multi-Label

Generally speaking, a multi-class classification problem refers to a single-label classification problem, i.e., selecting $1$ target class(es) out of $n$ candidate classes. Suppose the scores of the various classes are $s_1,s_2,

\dots,s_n$，目标类为$t\in\{1,2,\dots,n\}$, then the loss used is

\begin{equation}-\log \frac{e^{s_t}}{\sum\limits_{i=1}^n e^{s_i}}= - s_t + \log \sum\limits_{i=1}^n e^{s_i}\label{eq:log-softmax}\end{equation}

The optimization direction of this loss is to make the score of the target class $s_t$ become the maximum among $s_1,s_2,\dots,s_t$. For more on softmax, you may also refer to articles such as In Search of a Smooth Maximum Function and Miscellaneous Notes on Function Smoothing: Differentiable Approximations to Non-Differentiable Functions. more

Now let's turn to the multi-label classification problem, i.e., selecting $k$ target classes out of $n$ candidate classes. In this case, one naive approach is to use sigmoid activation, turning the problem into $n$ binary classification problems, and use the sum of binary cross-entropy losses as the overall loss. Clearly, when $n\gg k$, this approach faces a serious class imbalance problem, which then requires balancing strategies such as manually adjusting the weights of positive and negative samples, focal loss, and so on. After training, one also needs to further determine the optimal threshold based on a validation set.

At this point, a very natural source of confusion arises: why does "select $k$ out of $n$" require so much more extra work than "select $1$ out of $n$"?

I think this is quite unreasonable. After all, intuitively, selecting $k$ out of $n$ should just be a natural extension of selecting $1$ out of $n$, so it shouldn't require so much more extra work than $n$. Even if selecting $k$ out of $n$ is somewhat more complex, the difficulty should increase gradually. But if we turn it into multiple binary classification problems, then selecting $1$ out of $n$ actually becomes the hardest case, because that's when the class imbalance is most severe. Formally speaking, single-label classification is easier than multi-label classification precisely because single-label classification can use "Softmax + cross-entropy," which does not suffer from class imbalance, whereas "sigmoid + cross-entropy" in multi-label classification does suffer from imbalance.

So, the ideal solution should be to generalize "Softmax + cross-entropy" to multi-label classification.

Searching High and Low

To work out this generalization, I made many attempts and rejected many results, eventually settling on a relatively elegant scheme: constructing a combinatorial form of softmax as a generalization of the single-label softmax. In this section, we will first assume that $k$ is a fixed constant, and then discuss a scheme for automatically determining $k$ in the general case, ultimately arriving at an effective generalized form.

Combinatorial Softmax

First, let's consider the scenario where $k$ is a fixed constant, meaning that at prediction time, we simply output the top $k$ highest-scoring classes. What about at training time? As a natural generalization of softmax, we can consider using the following expression as the loss:

\begin{equation}-\log \frac{e^{s_{t_1}+s_{t_2}+\dots+s_{t_k}}}{\sum\limits_{1\leq i_1 < i_2 < \cdots < i_k\leq n}e^{s_{i_1}+s_{i_2}+\dots+s_{i_k}}}=\log Z_k - (s_{t_1}+s_{t_2}+\dots+s_{t_k})\end{equation}

where $t_1,t_2,\dots,t_k$ is the set of $k$ target labels, and $Z_k = \sum\limits_{1\leq i_1 < i_2 < \cdots < i_k\leq n}e^{s_{i_1}+s_{i_2}+\dots+s_{i_k}}$ is the partition function. Clearly, the above expression constructs a softmax with the total score of any $k$ classes $s_{i_1}+s_{i_2}+\dots+s_{i_k}$ as the basic unit, so it is a reasonable generalization of the single-label softmax. Alternatively, it can be understood as still being a single-label classification problem, except that it is a "select $1$ out of $C_n^k$" problem.

The tricky part of this scheme is computing $Z_k$, which is a sum of exponentials over $C_n^k$ total scores. However, we can use Newton's identities to help us compute this recursively. Let $S_k = \sum\limits_{i=1}^n e^{k s_i}$, then

\begin{equation}\begin{aligned} Z_1 =&\, S_1\\ 2Z_2 =&\, Z_1 S_1 - S_2\\ 3Z_3 = &\, Z_2 S_1 - Z_1 S_2 + S_3\\ \vdots\\ k Z_k = &\, Z_{k-1} S_1 - Z_{k-2} S_2 + \dots + (-1)^{k-2} Z_1 S_{k-1} + (-1)^{k-1} S_k \end{aligned}\end{equation}

So, to compute $Z_k$, we only need to recursively compute $k$ steps, which can be done in a reasonable amount of time. At prediction time, we simply output the $k$ highest-scoring classes.

Automatically Determining the Threshold

The discussion above concerns multi-label classification with a fixed number of outputs, but in general multi-label classification, the number of target labels is not fixed. To handle this, we set a maximum number of target labels $K\geq k$, and add a $0$ label as a padding label, at which point the loss becomes

\begin{equation}\log \overline{Z}_K - (s_{t_1}+s_{t_2}+\dots+s_{t_k}+\underbrace{s_0+\dots+s_0}_{K-k\text{count}})\end{equation}

and

\begin{equation}\begin{aligned} \overline{Z}_K =&\, \sum\limits_{1\leq i_1 < i_2 < \cdots < i_K\leq n}e^{s_{i_1}+s_{i_2}+\dots+s_{i_K}} + \sum\limits_{0 = i_1 = \dots = i_j < i_{j+1} < \cdots < i_K\leq n}e^{s_{i_1}+s_{i_2}+\dots+s_{i_K}}\\ =&\, Z_K + e^{s_0} \overline{Z}_{K-1} \end{aligned}\end{equation}

This looks complicated, but is actually quite simple: it still uses the total score of $K$ classes as the basic unit, but allows — and only allows — the $0$ class to repeat. At prediction time, we still output the $K$ highest-scoring classes, but allow the $0$ class to be output repeatedly; the equivalent effect is to use $s_0$ as a threshold, outputting only classes whose scores exceed $s_0$. The final expression shows that $\overline{Z}_K$ can also be computed recursively, so there's no difficulty in implementation.

Turning Back, There It Was, All Along

It seems that after "searching high and low," we finally have a result: the theory is in place, and the implementation isn't difficult either — so surely the next step is to run experiments and see how it performs? If it works well, maybe we could even consider writing a paper about it? A bright future seems to lie ahead! However...

Fortunately, or perhaps unfortunately, while verifying the effectiveness of this scheme, I consulted a few senior colleagues, and on their suggestion I went back and took a closer look at Circle Loss, which I hadn't studied carefully before. There, I saw its unified loss form (Equation (1) in the original paper), and realized that this unified form implies a much more concise generalization scheme.

So the unfortunate part is that, since such a ready-made and more concise scheme already exists, all this "searching high and low" turns out to have been of little use. The fortunate part, however, is that I did manage to find this better scheme in time — otherwise, if I had happily written up the earlier scheme as an article and published it, only to find that it was less simple and effective than an existing scheme, that would have been quite embarrassing indeed~

A Unified Loss Form

Let's look at the single-label classification cross-entropy $\eqref{eq:log-softmax}$ in a different form:

\begin{equation}-\log \frac{e^{s_t}}{\sum\limits_{i=1}^n e^{s_i}}=-\log \frac{1}{\sum\limits_{i=1}^n e^{s_i-s_t}}=\log \sum\limits_{i=1}^n e^{s_i-s_t}=\log \left(1 + \sum\limits_{i=1,i\neq t}^n e^{s_i-s_t}\right)\end{equation}

Why does this loss work? From the articles In Search of a Smooth Maximum Function and Miscellaneous Notes on Function Smoothing: Differentiable Approximations to Non-Differentiable Functions, we know that $\text{logsumexp}$ is in fact a smooth approximation of $\max$, so we have:

\begin{equation}\log \left(1 + \sum\limits_{i=1,i\neq t}^n e^{s_i-s_t}\right)\approx \max\begin{pmatrix}0 \\ s_1 - s_t \\ \vdots \\ s_{t-1} - s_t \\ s_{t+1} - s_t \\ \vdots \\ s_n - s_t\end{pmatrix}\end{equation}

The characteristic of this loss is that all non-target class scores $\{s_1,\cdots,s_{t-1},s_{t+1},\cdots,s_n\}$ are pairwise compared with the target class score $\{s_t\}$ via subtraction, and the maximum of these differences must be kept as close to zero (from below) as possible — thereby achieving the effect that "the target class score is greater than every non-target class score."

So, in a multi-label classification scenario with multiple target classes, we likewise want "every target class score to be no less than every non-target class score," and thus the following form of loss naturally suggests itself:

\begin{equation}\log \left(1 + \sum\limits_{i\in\Omega_{neg},j\in\Omega_{pos}} e^{s_i-s_j}\right)=\log \left(1 + \sum\limits_{i\in\Omega_{neg}} e^{s_i}\sum\limits_{j\in\Omega_{pos}} e^{-s_j}\right)\label{eq:unified}\end{equation}

where $\Omega_{pos},\Omega_{neg}$ are the sets of positive and negative classes of the sample, respectively. This form of loss is easy to understand: whenever we want $s_i < s_j$, we simply add a term of the form $e^{s_i - s_j}$ to $\log$. If we further add a scaling factor $\gamma$ and a margin $m$, we obtain the unified form found in the Circle Loss paper:

\begin{equation}\log \left(1 + \sum\limits_{i\in\Omega_{neg},j\in\Omega_{pos}} e^{\gamma(s_i-s_j + m)}\right)=\log \left(1 + \sum\limits_{i\in\Omega_{neg}} e^{\gamma (s_i + m)}\sum\limits_{j\in\Omega_{pos}} e^{-\gamma s_j}\right)\end{equation}

As a side note, the above expression is Equation (1) in the Circle Loss paper, but Equation (1) in the original paper is not itself called "Circle Loss" — it's Equation (4) that is called Circle Loss. So, strictly speaking, the expression above should not be called Circle Loss. Still, I personally think Equation (1) is the most interesting part of the entire paper.

Application to Multi-Label Classification

$\gamma$ and $m$ are generally only considered in metric learning, so here we'll focus solely on Equation $\eqref{eq:unified}$. If, in a multi-label classification setting selecting $k$ out of $n$, $k$ is fixed, then we can directly use Equation $\eqref{eq:unified}$ as the loss, and at prediction time simply output the $k$ highest-scoring classes.

For multi-label classification where $k$ is not fixed, we need a threshold to determine which classes to output. To this end, we likewise introduce an additional $0$ class, hoping that the scores of target classes all exceed $s_0$ and the scores of non-target classes are all below $s_0$. As mentioned earlier, "whenever we want $s_i < s_j$, we add a term $e^{s_i - s_j}$ to $\log$," so Equation $\eqref{eq:unified}$ now becomes:

\begin{equation}\begin{aligned} &\log \left(1 + \sum\limits_{i\in\Omega_{neg},j\in\Omega_{pos}} e^{s_i-s_j}+\sum\limits_{i\in\Omega_{neg}} e^{s_i-s_0}+\sum\limits_{j\in\Omega_{pos}} e^{s_0-s_j}\right)\\ =&\log \left(e^{s_0} + \sum\limits_{i\in\Omega_{neg}} e^{s_i}\right) + \log \left(e^{-s_0} + \sum\limits_{j\in\Omega_{pos}} e^{-s_j}\right)\\ \end{aligned}\end{equation}

If the threshold is specified to be 0, this simplifies to

\begin{equation}\log \left(1 + \sum\limits_{i\in\Omega_{neg}} e^{s_i}\right) + \log \left(1 + \sum\limits_{j\in\Omega_{pos}} e^{-s_j}\right)\label{eq:final}\end{equation}

This is the final loss form we've arrived at — a natural and concise generalization of "softmax + cross-entropy" to multi-label classification tasks. It has no class imbalance issue, because instead of turning multi-label classification into multiple binary classification problems, it turns the problem into pairwise comparisons between target class scores and non-target class scores, and thanks to the favorable properties of $\text{logsumexp}$, the weight of each term is automatically balanced.

Here is a reference implementation in Keras:

def multilabel_categorical_crossentropy(y_true, y_pred):
    """多标签分类的交叉熵
    说明：y_true和y_pred的shape一致，y_true的元素非0即1，
         1表示对应的类为目标类，0表示对应的类为非目标类。
    警告：请保证y_pred的值域是全体实数，换言之一般情况下y_pred
         不用加激活函数，尤其是不能加sigmoid或者softmax！预测
         阶段则输出y_pred大于0的类。如有疑问，请仔细阅读并理解
         本文。
    """
    y_pred = (1 - 2 * y_true) * y_pred
    y_pred_neg = y_pred - y_true * 1e12
    y_pred_pos = y_pred - (1 - y_true) * 1e12
    zeros = K.zeros_like(y_pred[..., :1])
    y_pred_neg = K.concatenate([y_pred_neg, zeros], axis=-1)
    y_pred_pos = K.concatenate([y_pred_pos, zeros], axis=-1)
    neg_loss = K.logsumexp(y_pred_neg, axis=-1)
    pos_loss = K.logsumexp(y_pred_pos, axis=-1)
    return neg_loss + pos_loss

So, the Conclusion Is

So, the final conclusion is Equation $\eqref{eq:final}$ — this is the unified loss for multi-label classification that this article set out to find. I welcome everyone to test it and report their results. I have also experimented with it on several multi-label classification tasks myself, and it matches the performance of carefully weight-tuned binary classification approaches.

It's worth pointing out that, besides standard multi-label classification problems, there are several other common task formulations that can also be regarded as multi-label classification — for example, sequence labeling based on 0/1 tagging, a typical example being my own "half-pointer, half-tagging" tagging scheme. So, from this perspective, there are actually many tasks that can be treated as multi-label classification for testing Equation $\eqref{eq:final}$. Indeed, I have tried it in an earlier triple-extraction example, task_relation_extraction.py $\eqref{eq:final}$, and ultimately achieved results consistent with those here.

Of course, finally, I should note that although in theory Equation $\eqref{eq:final}$, as a loss function for multi-label classification, can automatically resolve many issues, there is ultimately no such thing as a perfect solution that guarantees improvement. So when you replace your existing multi-label classification scheme with this one, there's no guarantee it will necessarily improve results — especially if your existing scheme has already handled class imbalance well through careful weight tuning or similar techniques, in which case the benefit of Equation $\eqref{eq:final}$ will be quite limited. After all, the original intent of Equation $\eqref{eq:final}$ was simply to let us achieve most of the benefit without having to tune too many hyperparameters.