Optimizing for ROC Curves on Class-Imbalanced Data by Training over a Family of Loss Functions

In this paper, we focus on binary classification problems with class imbalance. Specifically, let $D=\{(\mathbf{x}_{i},y_{i})\}_{i=1}^{n}$ be the training set consisting of $n$ i.i.d. samples from a distribution on $\mathcal{D}=\mathcal{X}\times\mathcal{Y}$ where $\mathcal{X}=\mathbb{R}^{d}$ and $\mathcal{Y}=\{0,1\}$. Then let $n_{c}$ be the number of samples with label $c$. Without loss of generality, we assume that $n_{1}<n_{0}$ (i.e., 0 is the majority class and 1 is the minority class) and measure the amount of imbalance in the dataset by $\beta=n_{0}/n_{1}>1$.

Let $f$ be a predictor with weights $\bm{\theta}$ and $\mathbf{z}=f(\mathbf{x},\bm{\theta})$ be $f$’s output for input $\mathbf{x}$. We assume $\mathbf{z}$ are logits (i.e., a vector $(z_{0},z_{1})$ of unnormalized scalars where $z_{c}$ is the logit corresponding to class $c$) and $\mathbf{p}$ are the outputs of the softmax function (i.e., a vector $(p_{0},p_{1})$ of normalized scalars such that $p_{0}+p_{1}=1$). Then the model’s prediction is

$\displaystyle\hat{y}=\begin{cases}1&\text{if }p_{1}>t,\\ 0&\text{otherwise},\end{cases}$

(1)

where $t\in[0,1]$ is a threshold and $t=0.5$ by default. To find the ROC curve of a classifier $f$, we compute the predictions over a range of $t$ values.

For all experiments, we use the Vector Scaling (VS) loss as defined by Kini et al. (2021). This loss is a modification of weighted Cross-entropy loss with two hyperparameters $\Delta,\iota$ that specify an affine transformation $\Delta z+\iota$ for each logit:

$\displaystyle\ell_{VS}(y,\mathbf{z})=-\omega_{y}\log\left(\frac{e^{\Delta_{y}z_{y}+\iota_{y}}}{\sum_{c\in\mathcal{Y}}e^{\Delta_{c}z_{c}+\iota_{c}}}\right).$

(2)

We follow parameterization by Kini et al. (2021) as follows:

$\displaystyle\iota_{c}=\tau\log\left(\frac{n_{c}}{n}\right)\quad\text{and}\quad\Delta_{c}=\left(\frac{n_{c}}{n_{0}}\right)^{\gamma}$

(3)

where $\tau\geq 0$ and $\gamma\geq 0$ are hyperparameters set by the user. We add an additional hyperparameter $\Omega$ to parameterize each class’ weight $\omega_{c}$ as follows

$\displaystyle\omega_{c}=\begin{cases}\Omega&\text{if }c=1,\\ 1-\Omega&\text{otherwise},\end{cases}$

(4)

where $\Omega\in[0,1]$. In the binary case, we can simplify the loss as follows (see Appendix E for details):

	$\displaystyle\ell_{VS}(0,\mathbf{z})$	$\displaystyle=(1-\Omega)\log\left(1+e^{z_{1}/\beta^{\gamma}-(z_{0}+\tau\log\beta)}\right)$		(5)
	$\displaystyle\ell_{VS}(1,\mathbf{z})$	$\displaystyle=\Omega\log\left(1+e^{(z_{0}+\tau\log\beta)-z_{1}/\beta^{\gamma})}\right).$		(6)

Note that VS loss is equivalent to cross entropy loss when $\Omega=0.5,\gamma=0$, and $\tau=0$. It is also equal to weighted cross entropy loss when $\Omega=\frac{n_{0}}{n},\gamma=0$, and $\tau=0$.

As discussed in Section 3.3, our parameterization of VS loss has three hyperparameters: $\Omega,\gamma$ and $\tau$. Each of these has theory associated with its effect on training under imbalance. $\Omega$ controls the overall balance between the magnitude of the gradients from the minority and majority class. $\gamma$ was originally proposed to compensate for the difference between the magnitude of the minority-class logits at training and testing time Ye et al. (2020); however, Kini et al. (2021) showed that it is essential for training beyond zero error and optimizing with $\gamma>0$ is equivalent to a solving Cost-Sensitive Support Vector Machine problem. $\tau$ enforces a larger margin on minority-class samples and has many of its own theoretical properties Cao et al. (2019): in particular, it has been shown to counteract some of the negative effects of training with $\gamma>0$ early in training Kini et al. (2021). Although there are intuitive and theoretical rationale for each of these hyperparameters Kini et al. (2021), it is not clear how sensitive a model’s performance is to their values, especially in terms of ROC curves on a binary dataset.

To understand this better, we trained several models on the same dataset over different combination of hyperparameter values. Specifically, we trained 512 models on the same dataset (CIFAR10 cat vs. dog with $\beta=100$) and varied only their values of $\Omega,\gamma$ and $\tau$. ²²2We used all unique combinations of $\Omega\in\{$0.5, 0.6, 0.7, 0.8, 0.9, 0.95, 0.99, 0.999$\}$, $\gamma\in\{$0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7$\}$, $\tau\in\{$0, 0.5, 1, 1.5, 2, 2.5, 3, 3.5$\}$. Figure 2 shows that there is a wide variance in performance of these models Specifically, these models have AUC values which range from 0.58 to 0.74 with a standard deviation of 0.03. Additionally, Figure 1 shows that this variance becomes more severe as the imbalance ratio increases. This implies that models trained with the VS loss under severe class imbalance are especially sensitive and variable in their results.

We next consider how much of this variance is explained by the hyperparameter values (i.e., is the randomness coming from the hyperparameters themselves or other randomness in training?). If the variance is mostly a result of different hyperparameter choices, then we can reasonably expect the practitioner to search for hyperparameters that generalize well. Otherwise, tuning hyperparameters may not be an effective way to improve performance.

To test this, we fit polynomials between the three hyperparameters ($\Omega,\gamma,\tau$) and two scalar metrics: AUC and overall accuracy on a balanced test set when $t=0.5$. Although TPR, equivalently recall, is not a good indicator of a model’s overall performance by itself (because it is trivial to increase TPR at the expense of FPR), we also fit a polynomial to TPR to gain insights on the effect of the hyperparameters on an ROC curve. We show results, including the $R^{2}$ of these regressions, in the three plots in Figure 3. We also fit three (one for AUC, accuracy, and TPR) degree-2 polynomials using all three hyperparameters as features and report $R^{2}$ values of these polynomials in the table in Figure 3. We choose polynomials of degree two instead of linear regression because we expect that an optimum exists.

We see that no single hyperparameter is strongly correlated with AUC; however, 11% of the variance can be explained by the polynomial fit to all three hyperparameters. This suggests that some variance can be removed by choosing strong hyperparameters; however, much of the variance comes from other sources. Additionally, only 60% of the variance in overall accuracy can be explained by the polynomial with all three hyperparameters. While this is significantly better than AUC’s $R^{2}$, it is still far from perfect. Thus, VS Loss hyperparameter values do not affect performance metrics in a very predictable way. It would be preferable to reduce the variance of VS Loss over performing extensive hyperparameter tuning.

The third plot in Figure 3 shows that $\tau$ is strongly correlated with TPR when the threshold $t=0.5$ ($R^{2}=0.95$). To understand this, recall that $\tau$ affects the additive term on the logits in Equations 2 and 3. Specifically, increasing $\tau$ enforces a larger margin on the minority class. In general, this will lead to increasing the softmax scores of the minority class at inference time, which is similar to post-hoc VS calibration Zhao et al. (2020). Thus, when $t$ is constant, increasing $\tau$ will likely cause an increase in TPR.

We can also visualize the effect of $\tau$ by considering the set of “break-even” logits, which are defined as the set of points of $\mathbf{z}$ where the loss is equal whether the sample has label 0 or 1 (i.e., the points of $\mathbf{z}$ such that $\ell_{VS}(1,\mathbf{z})=\ell_{VS}(0,\mathbf{z})$). With regular Cross-entropy, this set is $z_{1}=z_{0}$. Assume instead that we keep $\Omega=0.5,\gamma=0$, but have arbitrary values for $\tau$, then the equation for the set of break-even points becomes $z_{1}=z_{0}+\tau\log\beta$. In other words, varying $\tau$ is equivalent to shifting the set of break-even points (and the rest of the loss landscape) by $\tau\log\beta$.³³3Varying $\tau$ also shifts the loss landscape in a similar way for any fixed $\Omega\in(0,1)$ and $\gamma>0$. See Appendix G for details.. Figure 4 visualizes this.

In addition, $\Omega$ and $\gamma$ also have well-understood effects on the logit outputs. Appendices G and H include a general equation for the set of break-even points and more loss contours. Just as we saw in the example with $\tau$, different combinations of hyperparameters have different effects on how $\mathbf{z}$ is scaled, which in turn affects their softmax scores. Thus, training with different values of VS Loss hyperparameters corresponds to optimizing for different tradeoffs of TPR and FPR.

Dosovitskiy & Djolonga (2020) observed the computational redundancy involved in training separate models with slightly different loss functions, such as neural image compression models with different compression rates. They proposed loss conditional training (LCT) as a method to train one model to work over a family of losses.

Let $\bm{\lambda}$ be a vector which parameterizes a loss function. For example, we can parameterize VS loss by $\bm{\lambda}=(\Omega,\gamma,\tau)$. Then let $\mathcal{L}(\cdot,\cdot,\bm{\lambda})$ be the family of loss functions parameterized by $\bm{\lambda}\in\bm{\Lambda}\subseteq\mathbb{R}^{d_{\bm{\lambda}}}$ where $d_{\bm{\lambda}}$ is the length of $\bm{\lambda}$. In our example, $\mathcal{L}(\cdot,\cdot,\bm{\lambda})$ is the set of all possible VS loss functions obtained by different parameter values for $\Omega,\gamma,\tau$.

Normally, training finds the weights $\bm{\theta}$ on a model $f$ which minimize a single loss function $\mathcal{L}(\cdot,\cdot,\bm{\lambda})$ (i.e., we optimize for a single combination of $\bm{\lambda}_{0}=(\Omega_{0},\gamma_{0},\tau_{0})$ values in VS loss) as shown below.

$\displaystyle\bm{\theta}^{*}_{\lambda_{0}}$

$\displaystyle=\operatorname*{arg\,min}_{\bm{\theta}}\mathbb{E}_{(\mathbf{x},y)\sim D}\mathcal{L}(y,f(\mathbf{x},\bm{\theta}),\bm{\lambda}_{0})$

(7)

LCT instead optimizes over a distribution of $\bm{\lambda}$ values $P_{\bm{\Lambda}}$ as shown below.

$\displaystyle\bm{\theta}^{*}$

$\displaystyle=\operatorname*{arg\,min}_{\bm{\theta}}\mathbb{E}_{\bm{\lambda}\sim P_{\bm{\Lambda}}}\mathbb{E}_{(\mathbf{x},y)\sim D}\mathcal{L}(y,f(\mathbf{x},\bm{\theta},\bm{\lambda}),\bm{\lambda})$

(8)

LCT is implemented on Deep Neural Network (DNN) predictors by augmenting the network to take an additional input vector $\bm{\lambda}$ along with each data sample $\mathbf{x}$. During training, $\bm{\lambda}$ is sampled from $\bm{\Lambda}$ with each sample and is used in two ways: 1) as an additional input to the network and 2) in the loss function for the sample. During inference, the model takes as input a $\bm{\lambda}$ along with the sample and outputs the appropriate prediction.

In order to condition the model on $\bm{\lambda}$, the DNN is augmented with Feature-wise Linear Modulation (FiLM) layers Perez et al. (2018). These are small neural networks that take the conditioning parameter $\bm{\lambda}$ as input and output a $\bm{\mu}$ and $\bm{\sigma}$, used to modulate the activations channel-wise based on the value of $\bm{\lambda}$. Specifically, suppose a layer in the DNN has an activation map of size $H\times W\times C$. In LCT, we transform each activation $\mathbf{f}\in\mathbb{R}^{C}$ by $\bm{\mu}$ and $\bm{\sigma}$ as follows: $\mathbf{\tilde{f}}=\bm{\sigma}*\mathbf{f}+\bm{\mu}$ where both $\bm{\mu}$ and $\bm{\sigma}$ are vectors of size $C$, and “$*$” stands for element-wise multiplication.

We define our family of loss functions $\bm{\lambda}$ using VS loss. We try several choices for $\bm{\lambda}$, including $\bm{\lambda}=\Omega,\bm{\lambda}=\gamma,\bm{\lambda}=\tau$ and $\bm{\lambda}=(\Omega,\gamma,\tau)$. In the first three cases, we set the hyperparameters excluded from $\bm{\lambda}$ to constants (as is done in regular VS loss). We find that $\bm{\lambda}=\tau$ is especially strong at improving the average performance and reducing the variance in performance (Figure 8), which is supported by our insights in Figure 3 and Section 4.3. Thus, for most of our experiments, we use $\bm{\lambda}=\tau$.

For each mini-batch, we draw one $\bm{\lambda}$. This is done by independently sampling each hyperparameter in $\bm{\lambda}$ from a distribution that has a linear density over range $[a,b]$. In this distribution, the user specifies $a,b$, and the height of the probability density function (pdf) at $b$, $h_{b}$. The height at $a$, $h_{a}$, is then found to ensure the area under the pdf is 1. Unlike the triangular distribution, this general linear distribution does not require the pdf to be 0 at either endpoint and, unlike the uniform function, it does not require the slope of the pdf to be 0.⁴⁴4Appendix D contains more details about the linear distribution. We use this general distribution as a way to experiment with different distributions.

Although the value of $\bm{\lambda}$ at inference time affects the TPR and FPR when $t=0.5$, we observe that the ROC curves are almost identical for all values of $\bm{\lambda}$ in the range the model was trained on (see Appendix C). Thus, we evaluate LCT models at one $\bm{\lambda}$ and find their ROC curve from this output.

Methods. For each dataset and imbalance ratio $\beta$, we train 48 models with the regular VS loss, varying their hyperparameter values. We also train 48 models with LCT applied to VS loss where $\bm{\lambda}=\tau$ (VS Loss + LCT). We use $\Omega=0.5,0.7,0.9,0.99$, $\gamma=0.0,0.2,0.4$ for both methods (note that we set $\Omega,\gamma$ to constants when $\bm{\lambda}=\tau$ in LCT). We use $\tau=0,1,2,3$ for VS loss, and $[a,b]=[0,3]$ with $h_{b}=0.0,0.15,0.33,0.66$ for VS loss + LCT. We evaluate with $\bm{\lambda}=\tau=3$. To apply LCT, we augment the networks with one FiLM block after the final convolutional layer and before the linear layer. This block is comprised of two linear layers with 128 hidden units. Specifically, the first layer takes one input $\bm{\lambda}$ and outputs 128 hidden values and the second layer outputs 64 values which are used to modulate the 64 channels of convolutional activations (total of $1*128+128*64=8320$ parameters if $\bm{\lambda}$ is a scalar).

Datasets. We experiment on both toy datasets derived from CIFAR10/100 and more realistic datasets derived from Kaggle competitions. CIFAR10 cat/dog consists of the cat and dog classes from the CIFAR10 dataset. We also show results for all pairs of CIFAR10 classes in Section 6.3, but find the cat/dog pair is particularly well-suited for additional experiments since it is a challenging classification problem. CIFAR100 household electronics vs. furniture was proposed as a binary dataset with imbalance by Wang et al. (2016). Each class contains 5 classes from CIFAR100: electronics contains clock, computer keyboard, lamp, telephone and television, and furniture contains bed, chair, couch, table and wardrobe. For all CIFAR experiments, we split the data according to their train and test splits. Kaggle Dogs vs. Cats contains 25,000 images of dogs and cats from the Petfinder website Cukierski (2013). We split the data so that each class has 1,000 validation samples (a 92/8 train/validation split). Finally, SIIM-ISIC Melanoma is a classification dataset, which was designed by the Society for Imaging Informatics in Medicine (SIIM) and the International Skin Imaging Collaboration (ISIC) for a Kaggle competition Zawacki et al. (2020). This is a binary dataset where 8.8% of the samples belong to the positive (melanoma) class (i.e., $\beta=12$). We follow the procedure of Fang et al. (2023) and combine the 33,126 and 25,331 images from 2020 and 2019 respectively and split the data into an 80/20 train/validation split. Table 9 outlines the sizes of these. We subsample the minority class to obtain various imbalance ratios $\beta$. Specifically, we test the CIFAR datasets at $\beta=10,50,100,200$, the Kaggle datasets at $\beta=100,200$.

Model architectures and training procedure. We use two model architectures and training procedures, following the literature for each type of dataset. For CIFAR10/100 data, we follow Kini et al. (2021) and use a ResNet-32 model architecture. We find that LCT takes longer to train, so we train for 500 epochs instead of 200 (note: training the baseline models longer does not significantly alter performance and we do this for consistency). We train these models from scratch using an initial learning rate of 0.1 and decreasing this to $1e^{3}$ and $1e^{5}$ at 400 and 450 epochs respectively. For the Kaggle datasets, we follow Shwartz-Ziv et al. (2023) and use ResNext50-32x4d with weights pre-trained on ImageNet. We train the Melanoma dataset for 10 epochs, following the learning rates and weight decays of Fang et al. (2023) and train the Kaggle Dogs vs. cats models for 30 epochs, following the learning rates and weight decays of other finetuning tasks in Fang et al. (2023). For all experiments we use a batch size of 128 and gradient clipping with maximum norm equal to 0.5. Additionally, with the exception of Figure 8, we use Stochastic Gradient Descent (SGD) optimization with momentum=0.9. We run all experiments on A5000 GPUs.

Figures 5 and 6 compare the performance of the VS loss and VS loss + LCT methods on all four datasets, each at multiple imbalance ratios. For each method, we find the mean, min, max, and standard deviation of the 48 models trained with that method. We calculate these values across false positive rates (FPRs). We also show the distribution of AUCs on the SSIM-ISIC Melanoma dataset in Figure 1.

Figure 5 shows that at moderate imbalance ratios ($\beta=10$), there is very little variation in the ROC curves obtained by models trained with different hyperparameters. Here the two methods are essentially tied in their performance. However, at high imbalance ratios $\beta=100,200$ (Figures 5 and 6), the variance in ROC curves becomes much more drastic. Here LCT consistently outperforms the Baseline method. Specifically, LCT improves or ties the max and mean of the ROC curves in all cases, with a large gap in improvement for Melanoma at $\beta=100,200$.

For a more comprehensive evaluation, we compare the results of VS loss with and without LCT on all 45 pairs of classes chosen from the 10 CIFAR10 classes. We train each method on each of these binary datasets using 48 different hyperparameter values. We then analyze the aggregated performance of these models by AUC. For each binary dataset, we find the mean, maximum, minimum, and standard deviation of the AUCs over the hyperparameter values. We then compare Baseline (VS loss) and LCT (VS loss + LCT) for each of these metrics. For example, let $a^{max}_{k,\text{Baseline}}$ and $a^{max}_{k,\text{LCT}}$ be the maximum AUCs achieved on dataset $k$ by LCT and Baseline respectively. Then we obtain 45 of these values for each method (one for each dataset $k$). We then analyze these using a paired t-test and report results in Table 1. The table shows that LCT improves performance, increasing the min and mean AUC in nearly all cases, as well as the max AUC in most cases. It also reduces the standard deviation of AUC values. This shows that LCT significantly outperforms Baseline on a comprehensive set of datasets.

Previous work has shown that Sharpness Aware Minimization (SAM) is effective at optimizing imbalanced datasets in the multi-class setting Foret et al. (2021); Rangwani et al. (2022). SAM searches for parameters that lie in neighborhoods of uniformly low loss and usees a hyperparameter $\rho$ to define the size of these neighborhoods. Rangwani et al. (2022) shows that SAM can effectively escape saddle points in imbalanced settings by using large values for $\rho$ (i.e., $\rho=0.5,0.8$). We, however, find that using SAM in this way does not translate to improved performance on a binary dataset. In Figure 8, we compare the performance of the models across several hyperparameter values $\rho$. We find that, unlike the multi-class case, in the binary case the best $\rho$ is quite small (i.e., $\rho=0.1$). While training with SAM and $\rho=0.1$ can produce some very strong results, this also exacerbates the amount of variance in the performance. LCT with SGD gives models which are better on average and require less tuning. Note that larger $\rho$ values are absolutely detrimental to training in this setting and lead to models which are effectively naive classifiers. We also tried combining LCT and SAM with $\rho=0.1$ and found this to be less effective than training with LCT and SGD.

In this section, we compare different options for which hyperparameters in VS loss to use as $\bm{\lambda}$ in LCT. We consider setting $\bm{\lambda}$ to each of the three hyperparameters (using constant values for the other two). We also consider setting $\bm{\lambda}=(\Omega,\gamma,\tau)$. Figure 8 compares these options and shows that $\bm{\lambda}=\tau$ achieves a good compromise between improving the AUC of the models and reducing the variance of results. Furthermore, this choice of $\lambda$ follows our insights from Section 4.