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: The paper addresses the problem of how to perform end-to-end training for deep AUC ("area under ROC curve") maximisation. This is a challenging problem as this yields a harder optimisation problem than the standard CE minimisation. Existing approaches use two stage optimisation approaches where the first stage optimises a traditional loss function (e.g. the CE loss) and the second stage fine-tunes the network by optimising the AUC loss. The proposed approach is termed compositional deep AUC maximisation which can be trained end-to-end. The paper furthermore proposes an efficient stochastic optimisation algorithm solving compositional deep AUC maximisation problem. It has been evaluated on standard benchmark computer vision and medical imaging classification datasets.