Compositional Training for End-to-End Deep AUC Maximization

1 INTRODUCTION . Deep AUC maximization ( DAM ) represents a new learning paradigm for deep learning , which maximizes the area under ROC curve ( AUC ) on a training dataset for learning a deep neural network . It has received increasing attention recently due to the advancement in large-scale non-convex optimization algorithms for AUC maximization ( Liu et al. , 2019a ; Yuan et al. , 2021 ; Guo et al. , 2020a ; b ) . Recently , DAM has been successfully applied to different domains with imbalanced data ( Yuan et al. , 2020 ; Wang et al. , 2021b ) . For example , Yuan et al . ( 2020 ) has employed DAM for a variety of medical image classification tasks , e.g. , classification of X-ray images , skin lesion images , mammograms , and microscopic images , and they observed great improvements with about 1 % ∼5 % AUC increase over traditional deep learning approaches by optimizing a standard loss function , e.g. , the cross-entropy ( CE ) loss . These pioneering studies on DAM open a new direction for deep learning in the presence of imbalanced data but also raise many questions yet to be solved . A particular question relates to how the network is trained by DAM . Existing studies employ a two-stage approach for DAM , in which the first stage pretrains the network on the training data by optimizing a traditional loss function ( e.g . the CE loss ) and the second stage finetunes the network by optimizing an AUC loss . It was conjectured in ( Yuan et al. , 2020 ) the feature extraction layers learned by directly optimizing AUC from scratch are not as good as optimizing the standard CE loss , similarly as optimizing a class-weighted loss for deep learning with imbalanced data ( Cao et al. , 2019 ) . Although the ad-hoc two-stage method of DAM has achieved some success , this approach leads to several undesirable consequences increasing the engineering costs in practice : ( i ) which layers should we finetune in the second stage ? Fine-tuning all layers increases the training costs but not necessarily improves the final performance ( Jamal et al. , 2020 ; Qi et al. , 2020 ) ; ( ii ) when do we stop the training for the first-stage ? A long training time for the first-stage increases the overall training costs but not necessarily increases the final prediction performance , while a short training time for the first-stage could harm the prediction performance ( Kang et al. , 2019 ) . Hence , the literature has suggested different tricks for the two-stage approach , including the decoupling method that simply optimizes the classifier layer in the second-stage ( Kang et al. , 2019 ) and deferred re-weighting that only dedicates the iterations with the largest step size to the CE loss ( Cao et al. , 2019 ) , which could be borrowed to DAM as well . However , an important question remains open regarding DAM : How can we conduct end-to-end training for deep AUC maximization ? To answer this question , we have examined the learned feature representations by optimizing an AUC loss directly from scratch and confirmed the conjecture in ( Yuan et al. , 2020 ) that the learned feature representations exhibit no advantage over optimizing the CE loss directly . In Figure 1 we visualize the feature representations on an imbalanced training set for the CatvsDog classification learned by different methods and visualized by t-SNE ( van der Maaten & Hinton , 2008 ) . We can see that optimizing an AUC loss from scratch ( 2nd column ) does not yield a cleaner feature representations for the two classes of data than optimizing the CE loss . What makes end-to-end deep learning successful is its superb feature learning capability , i.e. , the lower layers capture the low-level features and higher layers capture the high-level features . In terms of feature learning , different examples roughly have equal weights regardless which classes they belong to . From this perspective , we could understand why optimizing AUC loss alone bears worse feature learning capability . The AUC loss assigns different weights to different examples from different classes for more robust classifier learning . In particular , the data in the positive class has a higher weight than data in the negative class . These non-equal weights are important for learning a robust classifier given that feature representations have been learned well but are not readily helpful for learning feature representations in an end-to-end fashion . Can we achieve both effects in a unified and end-to-end learning framework , i.e. , optimizing the CE loss with equal weights for robust feature learning and optimizing an AUC loss with uneven weights for robust classifier learning ? A naive approach is to simply optimize a linear combination of the CE loss and an AUC loss . However , this approach has a trade-off , meaning that AUC is not necessarily maximized due to the presence of the CE loss in the objective and the learned feature representations could be degraded by the AUC loss ( Figure 1 , 3rd column ) . In this paper , we propose a better and novel end-to-end training method that not only achieves both benefits of minimizing the CE loss for robust feature learning and minimizing an AUC loss for robust classifier learning , but also achieves the effect of “ 1+1 > 2 '' , i.e. , achieves better performance than the naive linear combination approach . The novel synthesis lies in how we composite the two training steps corresponding to the CE loss and the AUC loss . The central idea is to minimize a two-level compositional objective , where the outer function is an AUC loss , and the inner function is a gradient descent step towards minimizing the CE loss , which represents a quick adaptation to the solution to minimizing the CE loss . We propose a novel efficient stochastic algorithm with provable convergence for minimizing the compositional objective , which performs alternating gradient-based updates that are first based on the gradient of CE loss and then based on the gradient of an AUC loss at the point obtained in the first step . We summarize our contributions below . • We propose a novel training framework for end-to-end deep AUC maximization , namely compositional DAM . The novel compositional objective enables not only the robust feature learning of lower layers from minimizing the standard loss function but also the robust learning of a classifier from minimizing an AUC loss . • Theoretically , we propose an efficient stochastic optimization algorithm for solving compositional DAM , and establish the same convergence rate as standard SGD for optimizing a standard averaged loss . Empirically , we conduct extensive studies on benchmark and medical image datasets , and observe that the proposed method not only improves the baseline methods including the naive linear combination approach but also improves ad-hoc two-stage approaches for DAM . The learned feature representations of compositional DAM ( e.g. , Figure 1 right ) are much better than those learned by minimizing the CE loss or an AUC loss alone and their combination . 2 RELATED WORK . Deep AUC Maximization . AUC maximization has a history of two decades . Most of the existing studies revolve around the design of efficient optimization algorithms . Earlier papers have focused on full batch methods ( Herbrich et al. , 1999 ; Yan et al. , 2003 ; Ferri et al. , 2002 ; Freund et al. , 2003 ; Joachims , 2005 ; Herschtal & Raskutti , 2004 ; Rakotomamonjy , 2004 ; Zhang et al. , 2012 ) and online optimization methods ( Zhao et al. , 2011 ; Kar et al. , 2014 ; 2013 ; Gao et al. , 2013 ) . Recently , stochastic optimization for AUC maximization has become the dominating approach ( Ying et al. , 2016 ; Liu et al. , 2018 ; Natole et al. , 2018 ; 2019 ) . Ying et al . ( 2016 ) propose a milestone work for stochastic optimization of AUC . They consider optimizing the pairwise square loss and propose an equivalent min-max formulation that transforms the original non-decomposable objective into a decomposable one , which enables the design of efficient stochastic methods based on mini-batch of data without explicitly constructing the pairs . The min-max formulation also serves as the basis for recent works on DAM ( Liu et al. , 2019a ; Yuan et al. , 2021 ; Guo et al. , 2020a ; b ) . ( Liu et al. , 2019a ) is the first work that explicitly considers DAM and develops the first practical and provable stochastic algorithms for DAM based on the min-max formulation of the pairwise square loss function . However , this work only focuses on optimization and experiments are done on simple benchmark datasets . Later , Yuan et al . ( 2020 ) propose a new robust loss in the min-max form for DAM and evaluates the performance of DAM on various medical image classification tasks , which demonstrates great success of DAM . However , none of these works have addressed the problem of end-to-end training for DAM . Deep Learning with Imbalanced data . Deep learning in the presence of imbalanced data has recently attracted tremendous attention ( Cui et al. , 2019 ; Johnson & Khoshgoftaar , 2019 ; Masko & Hensman , 2015 ; Lee et al. , 2016 ; Khan et al. , 2017 ; Dablain et al. , 2021 ; Ren et al. , 2018 ; Jamal et al. , 2020 ; Qi et al. , 2020 ; Lin et al. , 2017 ; Cao et al. , 2019 ; Kang et al. , 2019 ; Liu et al. , 2019b ; Zhu & Yang , 2020 ; Zhou et al. , 2020 ; Xiang et al. , 2020 ; Wang et al. , 2021a ; 2020 ; Menon et al. , 2021 ) . Among these studies that are closely related to our work include ( Lin et al. , 2017 ; Cui et al. , 2019 ; Cao et al. , 2019 ; Qi et al. , 2020 ; Kang et al. , 2019 ; Jamal et al. , 2020 ) , which focus on optimizing different objectives from the standard CE loss , including class-weighted loss , focal loss , individually weighted loss functions , etc . Nevertheless , these works are not directly comparable to our method for maximizing AUC . Two-stage Approaches . However , directly optimizing a weighted loss for training a deep neural network from scratch does not work well ( Cao et al. , 2019 ; Kang et al. , 2019 ; Jamal et al. , 2020 ; Qi et al. , 2020 ; Yuan et al. , 2020 ) . This phenomenon was first observed in ( Cao et al. , 2019 ) , which is attributed to the degraded feature representations . To tackle this issue , Cao et al . ( 2019 ) propose a deferred re-weighting/re-sampling trick . It minimizes the standard average loss for the first stage and switches to minimizing the class-weighted loss or re-sampling method in the second stage . In their paper , the first stage is defined as the training period from the beginning to the iteration that the step size was reduced for the first time in SGD or momentum methods . Kang et al . ( 2019 ) investigate a decoupling approach , where the first stage learns the feature representations ( i.e. , a feature extraction network ) by optimizing a standard loss with a large number of iterations , and the second stage learns a robust classifier ( i.e . the classifier layer ) . The authors show that the decoupling approach can achieve better performance than the deferred re-weighting trick in ( Cao et al. , 2019 ) . However , the decoupling approaches do not necessarily yield the best performance . In particular , some studies have found that fine-tuning some higher layers besides the classifier layer in the second stage is beneficial ( Jamal et al. , 2020 ; Qi et al. , 2020 ) . Recently , Yang & Xu ( 2020 ) propose to use self-supervised learning for learning the feature representations for the first stage and to switch to re-weighting method in the second stage . Different from these studies , our work is to design an elegant end-to-end training framework for deep learning with an AUC loss .

Summary: This paper presents a new loss function for end-to-end training of deep networks for AUC maximization.  The proposed compositional loss combines the standard cross-entropy (CE) loss within the AUC loss such that the AUC is maximized for good classification while the CE portion of the loss helps to learn robust features. An optimization method is presented and is shown to have the same convergence rate as standard SGD. The proposed method is tested on 4 benchmark and 4 medical image datasets and compared against related losses, showing consistent improvement in AUC.