Revisiting Class Imbalance: Connecting Weight Adjustment and Loss-Engineering Tricks

再谈类别不平衡问题：调节权重与魔改Loss的对比联系 苏剑林 · 2020-08-31 · Machine Learning optimization loss function smoothness

The class imbalance problem, also known as the long-tail distribution problem, has already come up several times on this blog, for instance in From Hard Thresholding, to Softening, to Focal Loss, Generalizing "Softmax + Cross-Entropy" to Multi-Label Classification, and Using Mutual Information to Alleviate Class Imbalance. The most basic way to mitigate class imbalance is to adjust sample weights, while the seemingly more "sophisticated" approach is to cook up various modified losses (such as Focal Loss, Dice Loss, Logit Adjustment, etc.). This post aims to build a somewhat more systematic understanding of how these two approaches are connected.

Long-tail distribution: a small number of classes have a huge number of samples, while a large number of classes have very few samples.

From Smoothed Accuracy to Cross-Entropy

The analysis here mainly focuses on binary classification with sigmoid, but most of the conclusions generalize in parallel to multi-class classification with softmax. Let $x$ be the input, $y\in\{0,1\}$ the target, and $p_{\theta}(x) \in [0, 1]$ the model. Ideally, whatever metric we evaluate on, that is exactly the metric we should optimize. For classification problems, the most natural metric is of course accuracy, but accuracy provides no useful gradient, so it can't be trained directly. more

To get around this, we need a smoothed version of the metric. From an earlier post, Musings on Function Smoothing: Differentiable Approximations of Non-Differentiable Functions, the smoothed approximation of accuracy is

\begin{equation}\text{ACC}_{\text{smooth}}=\mathbb{E}_{(x,y)\sim\mathcal{D}}\big[y p_{\theta}(x) + (1 - y)(1 - p_{\theta}(x))\big]\end{equation}

where $\mathcal{D}$ is the training set. So in principle, we should minimize $-\text{ACC}_{\text{smooth}}$ directly. But in practice, directly optimizing this objective doesn't work very well — it's better to optimize cross-entropy instead:

\begin{equation}\text{cross_entropy}=\mathbb{E}_{(x,y)\sim\mathcal{D}}\big[-y \log p_{\theta}(x) - (1 - y)\log(1 - p_{\theta}(x))\big]\end{equation}

This is a bit puzzling: $\text{ACC}_{\text{smooth}}$ is clearly closer to the evaluation metric we actually care about, so why does cross-entropy end up being more beneficial for that metric?

This can be explained through gradients. For $p_{\theta}(x)$, it is typically produced by a sigmoid activation, i.e. $p_{\theta}(x)=\sigma(z_{\theta}(x))$, where $\sigma(t)=\frac{1}{1+e^{-t}}$, and its derivative is $\sigma'(t)=\sigma(t)(1 - \sigma(t))$, with $z_{\theta}(x)$ being what we usually call the "logits."

Suppose $y$ is 1, so the corresponding $-\text{ACC}_{\text{smooth}}$ is $-p_{\theta}(x)=-\sigma(z_{\theta}(x))$, whose gradient is

\begin{equation}-\nabla_{\theta} p_{\theta}(x) = - p_{\theta}(x) (1 - p_{\theta}(x))\nabla_{\theta}z_{\theta}(x)\end{equation}

As just mentioned, $y$ is 1, so the training target is $p_{\theta}(x)\to 1$. Therefore we would expect that when $p_{\theta}(x)$ is close to 0 (large error), the gradient should be large, and when $p_{\theta}(x)$ is close to 1 (small error), the gradient should be small. But the $-\nabla_{\theta} p_{\theta}(x)$ above clearly doesn't behave that way — its modulating factor $p_{\theta}(x) (1 - p_{\theta}(x))$ attains its maximum at 0.5, while both 0 and 1 give the minimum. This means that when the error is large, the gradient actually ends up being small too, which hurts optimization efficiency and ultimately leads to poor overall performance. Cross-entropy, in contrast, gives

\begin{equation}-\nabla_{\theta} \log p_{\theta}(x) = - (1 - p_{\theta}(x))\nabla_{\theta}z_{\theta}(x)\end{equation}

which precisely removes the problematic factor $p_{\theta}(x)$ from the gradient, making optimization more efficient and ultimately yielding better results. The analysis above is for $y=1$; if $y=0$, the conclusion is exactly the same.

From Smoothed F1 to Weighted Cross-Entropy

From this process, we get the sense that all these various loss modifications are, at their core, just ways of adjusting the gradient — by obtaining a more sensible gradient, we can optimize more effectively and end up with a better model. Let's push this reasoning further: originally, the gradient of our approximate objective was $-\nabla_{\theta}p_{\theta}(x)$, and it turned out that $-\nabla_{\theta}\log p_{\theta}(x)$ worked better. What if, without carefully digging into the underlying reason, we simply take $p\to \log p$ as an "axiom" to be applied elsewhere? Would that hold up, and would it lead to anything interesting?

Consider the following example. When negative samples vastly outnumber positive ones, our evaluation metric is usually no longer accuracy (otherwise just always predicting 0 would already give high accuracy). We typically care about the F1 score of the positive class instead, and directly optimizing F1 isn't easy either, so we again need a smoothed version. The post Musings on Function Smoothing: Differentiable Approximations of Non-Differentiable Functions also gives us the result:

\begin{equation}\text{F1}_{\text{smooth}}=\frac{2 \mathbb{E}_{(x,y)\sim\mathcal{D}}\big[y p_{\theta}(x)\big]}{\mathbb{E}_{(x,y)\sim\mathcal{D}}\big[y + p_{\theta}(x)\big]}\end{equation}

So our minimization objective was originally $-\text{F1}_{\text{smooth}}$. Following the "axiom" above, let's first take the gradient of $-\text{F1}_{\text{smooth}}$ directly:

\begin{equation}\begin{aligned}&-\nabla_{\theta}\frac{2 \mathbb{E}_{(x,y)\sim\mathcal{D}}\big[y p_{\theta}(x)\big]}{\mathbb{E}_{(x,y)\sim\mathcal{D}}\big[y + p_{\theta}(x)\big]}\\ =&-2\frac{\mathbb{E}_{(x,y)\sim\mathcal{D}}\big[y \nabla_{\theta}p_{\theta}(x)\big]}{\mathbb{E}_{(x,y)\sim\mathcal{D}}\big[y + p_{\theta}(x)\big]} + 2\frac{\mathbb{E}_{(x,y)\sim\mathcal{D}}\big[y p_{\theta}(x)\big]\mathbb{E}_{(x,y)\sim\mathcal{D}}\big[\nabla_{\theta}p_{\theta}(x)\big]}{\left(\mathbb{E}_{(x,y)\sim\mathcal{D}}\big[y + p_{\theta}(x)\big]\right)^2}\\ =&-\frac{2\mathbb{E}_{(x,y)\sim\mathcal{D}}\big[\big(y-\text{F1}_{\text{smooth}}/2\big)\nabla_{\theta}p_{\theta}(x)\big]}{\mathbb{E}_{(x,y)\sim\mathcal{D}}\big[y + p_{\theta}(x)\big]} \end{aligned}\end{equation}

where $\frac{2}{\mathbb{E}_{(x,y)\sim\mathcal{D}}[y + p_{\theta}(x)]}$ is an overall scaling factor; what we really care about is the gradient contribution per sample, so the result is

\begin{equation}-\mathbb{E}_{(x,y)\sim\mathcal{D}}\big[\big(y-\text{F1}_{\text{smooth}}/2\big)\nabla_{\theta}p_{\theta}(x)\big]\end{equation}

Applying the "$p\to \log p$" axiom (and $-p\to\log(1-p)$ for negative samples), we get the final gradient:

\begin{equation}-\mathbb{E}_{(x,y)\sim\mathcal{D}}\big[y\cdot\big(1-\text{F1}_{\text{smooth}}/2\big)\cdot\nabla_{\theta}\log p_{\theta}(x) + (1 - y)\cdot\text{F1}_{\text{smooth}}/2\cdot\nabla_{\theta}\log (1-p_{\theta}(x))\big]\end{equation}

This is equivalent to the gradient of the objective

\begin{equation}-\mathbb{E}_{(x,y)\sim\mathcal{D}}\big[y\cdot\big(1-\text{F1}_{\text{smooth}}/2\big)\cdot\log p_{\theta}(x) + (1 - y)\cdot\text{F1}_{\text{smooth}}/2\cdot\log (1-p_{\theta}(x))\big]\end{equation}

(where $\text{F1}_{\text{smooth}}$ is treated as not requiring gradient), which is nothing but cross-entropy for positive samples weighted by $1-\text{F1}_{\text{smooth}}/2$, and cross-entropy for negative samples weighted by $\text{F1}_{\text{smooth}}/2$.

From Margin Expansion to Logit Adjustment

Regardless of what the evaluation metric actually is, what we ultimately want is for every single sample to be predicted correctly. The trouble is that classes with fewer samples aren't learned thoroughly enough, so their generalization tends to be poor.

Let's think about this geometrically. Ideally, in the embedding space, each class occupies its own "territory," and the territories of different classes don't overlap. The poor generalization of minority classes essentially manifests as their territory being too small — and often further squeezed by the "pressure" from classes with many more samples. Under such pressure, mere "survival" is already a challenge, let alone being able to accommodate new samples never seen during training.

How do we fix this? The intuition is actually quite vivid: if every sample in a minority class is a "heavyweight" — capable of fighting off ten at once — then even with few samples, the class won't lose out in the "territorial dispute." Consider a $n$-class classification problem, where a given sample's embedding vector is $f_{\theta}(x)$ and the class vector is $u_y$. The similarity between the sample and a class vector is usually measured by the inner product $\langle f_{\theta}(x), u_y\rangle$. Suppose each sample can claim a territory of radius $r_y$ — meaning any $z$ satisfying $\Vert z - f_{\theta}(x)\Vert \leq r_y$ is considered a valid embedding for that sample. This means that for any $z$ satisfying this condition, its similarity to $u_y$ should be greater than its similarity to any other class.

Now consider

\begin{equation}\langle z, u_y\rangle = \langle f_{\theta}(x), u_y\rangle + \langle z - f_{\theta}(x), u_y\rangle\end{equation}

Since $\Vert z - f_{\theta}(x)\Vert \leq r_y$, we clearly have

\begin{equation}\langle f_{\theta}(x), u_y\rangle - r_y\Vert u_y\Vert\leq\langle z, u_y\rangle \leq \langle f_{\theta}(x), u_y\rangle + r_y\Vert u_y\Vert\end{equation}

So, in order to guarantee that "$z$'s similarity to $u_y$ should always exceed its similarity to any other class," it suffices that "the minimum similarity between $z$ and $u_y$ should exceed the maximum similarity between it and any other class." Hence our optimization objective becomes

\begin{equation}-\log\frac{e^{\langle f_{\theta}(x), u_y\rangle - r_y\Vert u_y\Vert}}{e^{\langle f_{\theta}(x), u_y\rangle - r_y\Vert u_y\Vert}+\sum\limits_{i\neq y} e^{\langle f_{\theta}(x), u_i\rangle + r_y\Vert u_i\Vert}}\end{equation}

As you can see, this is essentially the same as margin-based softmax variants like AM-Softmax and Circle Loss — the specific form doesn't really matter. What matters is simply assigning a larger margin to minority classes (since each sample in a small class needs to be more "combat-capable"). So how should we set the margin for each class? An earlier post, Using Mutual Information to Alleviate Class Imbalance, offered one solution: $m_y=-\tau\log p(y)$, where $p(y)$ is the prior distribution. This gives us

\begin{equation}-\log\frac{e^{\langle f_{\theta}(x), u_y\rangle + \tau \log p(y)}}{\sum\limits_{i} e^{\langle f_{\theta}(x), u_i\rangle + \tau \log p(i)}}\end{equation}

which connects us directly to the logit adjustment loss — or, put another way, gives logit adjustment loss a geometric interpretation. Fundamentally, logit adjustment is also a form of weight adjustment; the difference is that ordinary weight adjustment is applied after the $\log$ in the loss function, whereas logit adjustment is applied before the $\log$.

A Rough Summary

This post has reflected on the phenomenon of class imbalance and its countermeasures, with the main goal of using relatively intuitive reasoning to reveal the ideas behind various loss-engineering tricks. What we find is that, fundamentally, all of these schemes amount to adjusting sample weights or class weights. The line of reasoning here has been fairly loose and free-form — basically a record of the author's own brainstorming — so please forgive, and do point out, any errors or omissions.