AUC-mixup: Deep AUC Maximization with Mixup

In recent years, deep AUC maximization (DAM), focusing on developing deep learning models that directly optimize the area under the receiver operating characteristic curve (AUC), has gained growing importance. Different methods of DAM, such as optimizing the AUC margin loss [9] (AUCM) and compositional DAM [8], have been successfully applied to medical image classification to improve the AUC performance. These methods demonstrate superior performance in large-scale medical image classification tasks such as Chexpert [3] and Melanoma [5] compared to optimizing traditional loss functions, such as cross-entropy loss (CE) and focal loss.

Although DAM methods work well on large datasets, they still face the challenge of vulnerability to overfitting when training on imbalanced datasets with a small size. Typically, DAM losses place more emphasis on the positive class with fewer samples, and as the overall data volume decreases, the samples of this class become scarce and limited, which leads to the overfitting problem. In this context, mixup augmentation [11] is an effective solution by introducing soft labels into training and generating much more data using convex combinations of samples. With mixup augmentation, the focus of DAM would shift from the minor class to the combination of samples from different classes, which mitigates the overfitting issue.

However, existing DAM losses are developed on hard labels and thus are incompatible with mixup augmentation. Different from traditional loss functions that are defined over individual data, DAM losses are non-decomposable, making the incorporation of soft labels more complicated. To address the problem, we propose an AUC-mixup loss by replacing conditional means in the min-max AUC margin loss with soft means by using soft labels. Our goal is to utilize the AUC-mixup loss to improve DAM and compositional DAM methods in medical tasks with a small number of data by incorporating mixup augmentation. We validate our method on imbalanced benchmark and medical image datasets, including several 3D datasets, to demonstrate the superiority in generalization performance over two standard DAM baselines.

Let $\mathbb{I}(.)$ be an indicator function of a predicate, let $S=\{(\mathbf{x}_{1},y_{1}),...,(\mathbf{x}_{n},y_{n})\}$ denote a set of training examples (e.g., a 2D or 3D image), and let $y_{i}\in\{1,0\}$ denote its corresponding label. Let $\mathbf{w}\in\mathbb{R}^{d}$ denote the parameters of the deep neural network and let $h_{\mathbf{w}}(\mathbf{x})=h(\mathbf{w},\mathbf{x})$ denote the prediction of the neural network on input data $\mathbf{x}$. Yuan et al. [9] proposed DAM by minimizing the AUC margin loss that is equivalent to the following min-max optimization:

	$\displaystyle\min_{\mathbf{w}\in\mathbb{R}^{d},(a,b)\in\mathbb{R}^{2}}\max_{\alpha\geq 0}$	$\displaystyle\frac{\sum_{i=1}^{N}\left(h_{\mathbf{w}}(\mathbf{x}_{i})-a\right)^{2}\mathbb{I}(y_{i}=1)}{N_{+}}+\frac{\sum_{i=1}^{N}\left(h_{\mathbf{w}}(\mathbf{x}_{i})-b\right)^{2}\mathbb{I}(y_{i}=0)}{N_{-}}$		(1)
		$\displaystyle+2\alpha\left(m-\frac{\sum_{i=1}^{N}h_{\mathbf{w}}(\mathbf{x}_{i})\mathbb{I}(y_{i}=1)}{N}+\frac{\sum_{i=1}^{N}h_{\mathbf{w}}(\mathbf{x}_{i})\mathbb{I}(y_{i}=0)}{N_{-}}\right)-\alpha^{2},$

where $m$ is a hyperparameter to control the desired margin between optimal $a$ (i.e., the mean score of positive data) and optimal $b$ (i.e., the mean score of negative data), $N_{+}$ is the number of positive samples and $N_{-}$ is the number of negative samples. AUC margin loss is shown to enjoy better robustness than AUC square loss [9]. An improvement of minimizing the AUC margin loss from scratch is compositional DAM, which minimizes a compositional objective function, where the outer function corresponds to the AUC margin loss and the inner function represents a gradient descent step for minimizing a cross-entropy (CE) loss. However, the mixup technique cannot be directly applied to both methods, where an augmented example $\hat{\mathbf{x}}$ and its corresponding soft label $\hat{y}$ are generated from two randomly sampled training data and labels:

$\displaystyle\hat{\mathbf{x}}=\lambda\mathbf{x}_{i}+(1-\lambda)\mathbf{x}_{j},\quad\hat{y}=\lambda y_{i}+(1-\lambda)y_{j},$

(2)

where $\hat{y}\in(0,1)$ and $\lambda\sim Beta(\alpha,\alpha)$, $\alpha\in(0,\infty)$. It is obvious that directly using the AUC margin loss will ignore all the augmented samples with $\lambda\in(0,1)$. To address this problem, we propose an AUC-mixup loss. The idea is to replace the conditional means in (1) by soft means, i.e.,

	$\displaystyle\min_{\mathbf{w}\in\mathbb{R}^{d},(a,b)\in\mathbb{R}^{2}}\max_{\alpha\geq 0}$	$\displaystyle\frac{\sum_{i=1}^{\hat{N}}\left(h_{\mathbf{w}}(\hat{\mathbf{x}}_{i})-a\right)^{2}\hat{y}_{i}}{\sum_{i=1}^{\hat{N}}\hat{y}_{i}}+\frac{\sum_{i=1}^{\hat{N}}\left(h_{\mathbf{w}}(\hat{\mathbf{x}}_{i})-b\right)^{2}(1-\hat{y}_{i})}{\sum_{i=1}^{\hat{N}}(1-\hat{y}_{i})}$		(3)
		$\displaystyle+2\alpha\left(m-\frac{\sum_{i=1}^{\hat{N}}h_{\mathbf{w}}(\hat{\mathbf{x}}_{i})\hat{y}_{i}}{\sum_{i=1}^{\hat{N}}\hat{y}_{i}}+\frac{\sum_{i=1}^{\hat{N}}h_{\mathbf{w}}(\hat{\mathbf{x}}_{i})(1-\hat{y}_{i})}{\sum_{i=1}^{\hat{N}}(1-\hat{y}_{i})}\right)-\alpha^{2},$

where $\{(\hat{\mathbf{x}}_{i},\hat{y}_{i})\}_{i=1}^{\hat{N}}$ denotes the mixup augmented dataset. It is not difficult to show that the optimal $a,b$ are soft mean scores of positive and negative data, respectively, i.e., $a=\frac{\sum_{i=1}^{\hat{N}}h_{\mathbf{w}}(\hat{\mathbf{x}}_{i})\hat{y}_{i}}{\sum_{i=1}^{N}\hat{y}_{i}},b=\frac{\sum_{i=1}^{\hat{N}}h_{\mathbf{w}}(\hat{\mathbf{x}}_{i})(1-\hat{y}_{i})}{\sum_{i=1}^{\hat{N}}(1-\hat{y}_{i})}$. The AUC-mixup loss can be easily integrated with compositional DAM.

Datasets. In this section, we perform extensive experiments to evaluate the proposed AUC-mixup approaches on diverse benchmark and medical image datasets. With respect to benchmark datasets, we choose the Cat&Dog, CIFAR-10, CIFAR-100, and STL-10 and construct a binary imbalanced version of these datasets by following the instruction of Yuan et al. [9]. We adopt a DenseNet121 [2] network pre-trained in ImageNet as the backbone for benchmark datasets, following Yuan et al. [9]. For medical image datasets, we choose six MEDMNIST [7] datasets with the format of 2D or 3D images, namely PneumoniaMNIST (PneumoniaM), BreastMNIST (BreastM), NoduleMNIST3D (NoduleM), AdrenalMNIST3D (AdrenalM), VesselMNIST3D (VesselM), SynapseMNIST3D (SynapseM)). These datasets are naturally imbalanced. We adopt ResNet18 [1] network for training on these 2D and 3D datasets (with an acsconv [6]). We split all datasets into training set, validation set, and test set to conduct cross-evaluation for tuning hyperparameters, and report the AUC score on the test set(the means and standard deviations of three runs).

Settings. In the experiments, we apply AUC mixup strategy to DAM for vanila training from scratch (AUC-mixup) and compositional training (CT-mixup). We choose four methods with different losses as baselines: cross-entropy loss (CE), focal loss (Focal), AUCM loss [9], and compositional AUC loss (CT-AUC) [8]. For all methods, we tune the learning rate in {0.1,0.01,0.001} on all datasets, and decrease it by a factor of 10 at 50% and 75% of total training time. The number of training epochs is 100 for all datasets except for breastmnist which is trained for 200 epochs, considering its relatively smaller size. For focal loss, its parameters $\hat{\alpha},\hat{\lambda}$ are fixed on 1 and 2. For AUCM and CT-AUC, the margin $m$ is set at 1.0 on all datasets. The ADAM optimizer [4] is used for optimizing CE and focal loss, while PESG [9] and PDSCA [8] are used for optimizing AUCM and CT-AUC losses, respectively, with weight decay at 0.0001 and epoch decay at 0.001. The beta parameters of PDSCA are set at 0.9 and the number of inner gradient steps $k$ is set at 1. We use a batch size of 64 and a Dualsampler [10] to guarantee positive data in a minibatch.

Results. We compare the testing AUC scores of different methods on all datasets in Table 1. We can observe that (i) the AUC-mixup helps achieve the highest AUC scores on all datasets; (ii) the AUC-mixup strategy usually yields an improvement of varying degrees compared to the corresponding DAM methods without using AUC-mixup; (iii) the AUC-mixup is competitive if not better than CT-mixup, which indicates that employing the AUC-mixup loss for training from scratch can eliminate the additional compositional training overhead without sacrificing the prediction performance. We further show the learned feature representations of DAM methods on the BreastMNIST training data in Figure 1, which illustrates that employing the AUC-mixup loss obtains better feature representations.