When AUC meets DRO: Optimizing Partial AUC for Deep Learning with Non-Convex Convergence Guarantee
Dixian Zhu, Gang Li, Bokun Wang, Xiaodong Wu, Tianbao Yang
Introduction
AUC, short for the area under the ROC curve, is a performance measure of a model, where the ROC curve is a curve of true positive rate (TPR) vs false positive rate (FPR) for all possible thresholds. AUC maximization in machine learning has a long history dating back to early 2000s (Herbrich et al. 1999). It has four ages in the twenty-years history, full-batch based methods in the first age, online methods in the second age, stochastic methods in the third age, and deep learning methods in the recent age. The first three ages focus on learning linear models or kernelized models. In each age, there have been seminal works in rigorous optimization algorithms that play important roles in the evolution of AUC maximization methods. Recent advances in non-convex optimization (in particular non-convex min-max optimization) (Liu et al. 2020) has driven large-scale deep AUC maximization to succeed in real-world tasks, e.g., medical image classification (Yuan et al. 2020) and molecular properties prediction (Wang et al. 2021).
Nevertheless, the research on efficient optimization algorithms for partial AUC (pAUC) lag behind. In many applications, there are large monetary costs due to high false positive rates (FPR) and low true positive rates (TPR), e.g., in medical diagnosis. Hence, a measure of primary interest is the region of the curve corresponding to low FPR and/or high TPR, i.e., pAUC. There are two commonly used versions of pAUC, namely one-way pAUC (OPAUC) (Dodd & Pepe 2003) and two-way pAUC (TPAUC) (Yang et al. 2019), where OPAUC puts a restriction on the range of FPR, i.e., FPR (Figure 1 middle) and TPAUC puts a restriction on the lower bound of TPR and the upper bound of FPR, i.e., TPR, FPR (Figure 1 right). Compared with standard AUC maximization, pAUC maximization is more challenging since its estimator based on training examples involves selection of examples whose prediction scores are in certain ranks.
To the best of our knowledge, there are few rigorous and efficient algorithms developed for pAUC maximization for deep learning. Some earlier works have focused on pAUC maximization for learning linear models. For example, Narasimhan & Agarwal 2017 have proposed a structured SVM approach for one-way pAUC maximization, which is guaranteed to converge for optimizing the surrogate objective of pAUC. However, their approach is not efficient for big data and is not applicable to deep learning, which needs to evaluate the prediction scores of all examples and sort them at each iteration. There are some heuristic approaches, e.g., updating the model parameters according to the gradient of surrogate pAUC computed based on a mini-batch data (Kar et al. 2014) or using an ad-hoc weighting function for each example for computing the stochastic gradient (Yang et al. 2021). However, such approaches are either not guaranteed to converge or could suffer a large approximation error.
In this paper, we propose more systematic and rigorous optimization algorithms for pAUC maximization with convergence guarantee, which are applicable to deep learning. We consider both OPAUC maximization and TPAUC maximization, where for OPAUC we focus on maximizing pAUC in the region where and for TPAUC we focus on maximizing pAUC in the region where and for some . In order to tackle the challenge of computing unbiased stochastic gradients of the surrogate objective of pAUC, we propose new formulations based on distributionally robust optimization (DRO), which allows us to formulate the problem into weakly convex optimization, and novel compositional optimization problems, and to develop efficient stochastic algorithms with convergence guarantee. We summarize our contributions below.
For OPAUC maximization, for each positive example, we define a loss over all negative examples based on DRO. We consider two special formulations of DRO, with one based on the conditional-value-at-risk (CVaR) function that yields an exact estimator of the surrogate objective of OPAUC, and another one based on Kullback–Leibler (KL) divergence regularized DRO that yields a soft estimator of the surrogate objective.
We propose efficient stochastic algorithms for optimizing both formulations of OPAUC and establish their convergence guarantee and complexities for finding a (nearly) stationary solution. We also demonstrate that the algorithm for optimizing the soft estimator based on the KL divergence regularized DRO can enjoy parallel speed-up.
For TPAUC maximization, we apply another level of DRO with respect to the positive examples on top of OPAUC formulations, yielding both exact and soft estimators for TPAUC. We also provide two rigorous stochastic algorithms with provable convergence for optimizing both the exact and soft estimator of TPAUC, with the latter problem formulated as a novel three-level compositional stochastic optimization problem.
We conduct extensive experiments for deep learning on image classification and graph classification tasks with imbalanced data. We compare with heuristic and ad-hoc approaches for pAUC maximization and multiple baseline methods, and observe superior performance of the proposed algorithms.
To the best our knowledge, this work is the first one that provides rigorous stochastic algorithms and convergence guarantee for pAUC maximization that are efficient and applicable to deep learning. We expect the proposed novel formulations for OPAUC and TPAUC will allow researchers to develop even faster algorithms than the proposed algorithms in this paper.
Related Work
In this section, we provide a brief overview of related work for pAUC maximization.
Earlier works have considered indirect methods for pAUC maximization (Rudin 2009; Agarwal 2011; Rakotomamonjy 2012; Li et al. 2014; Wu et al. 2008). They did not directly optimize the surrogate objective of pAUC but instead some objectives that have some relationship to the right corner of ROC curve, e.g., p-norm push (Rudin 2009), infinite-push (Agarwal 2011; Rakotomamonjy 2012; Li et al. 2014), and asymmetric SVM objective (Wu et al. 2008). Nevertheless, none of these studies propose algorithms that are scalable and applicable for deep learning.
In (Kar et al. 2014), the authors proposed mini-batch based stochastic methods for pAUC maximization. At each iteration, a gradient estimator is simply computed based on the pAUC surrogate function of the mini-batch data. However, this heuristic approach is not guaranteed to converge for minimizing the pAUC objective and its error scales as , where is the mini-batch size. Narasimhan & Agarwal 2013b; Narasimhan & Agarwal 2013a; Narasimhan & Agarwal 2017 developed rigorous algorithms for optimizing pAUC with FPR restricted in a range based on the structured SVM formulation. However, their algorithms are only applicable to learning linear models and are not efficient for big data due to per-iteration costs proportional to the size of training data. Recently, Yang et al. 2021 considered optimizing two-way partial AUC with FPR less than and TPR larger than . Their paper focuses on simplifying the optimization problem that involves selection of top ranked negative examples and bottom ranked positive examples. They use ad-hoc weight functions for each positive and negative examples to relax the objective function into decomposable over pairs. The weight function is designed such that the larger the scores of negative examples the higher are their weights, the smaller the scores of positive examples the higher are their weights. Nevertheless, their objective function might have a large approximation error for the pAUC estimator.
There are also some studies about partial AUC maximization without providing rigorous convergence guarantee on their methods, including greedy methods (Wang & Chang 2011; Ricamato & Tortorella 2011) and boosting methods (Komori & Eguchi 2010; Takenouchi et al. 2012). Some works also use pAUC maximization for learning non-linear neural networks (Ueda & Fujino 2018; Iwata et al. 2020). However, it is unclear how the optimization algorithms were designed as there were no discussion on the algorithm design and convergence analysis. Finally, it was brought to our attention that a recent work (Yao et al. 2022) also considered partial AUC maximization with a non-convex objective. The difference between this work and (Yao et al. 2022) is that: (i) they focus on optimizing one-way pAUC with FPR in a certain range where ; in contrast we consider optimizing both one-way pAUC and two-way pAUC, but for one-way pAUC we only consider FPR in a range of ; (ii) the second difference is that the proposed algorithms in this paper for one-way pAUC maximization has a better complexity than that established in (Yao et al. 2022).
Preliminaries
A function is weakly convex if there exists such that is a convex function. A function is -smooth if its gradient is Lipchitz continuous, i.e., .
where means that only negative examples whose prediction scores are in certain quantiles are considered. As a result, we have the following non-parametric estimator of OPAUC:
where . In this work, we will focus on optimizing for some .
Similarly, a non-parametric estimator of TPAUC with is given by
where .
By using the CVaR divergence for some such that is an integer, we have,
The estimator in (6) is also known as the estimator of conditional-value-at-risk (Rockafellar et al. 2000).
AUC meets DRO for OPAUC Maximization
The challenge of optimizing a surrogate objective of pAUC in (7) lies at tackling the selection of top ranked negative examples from , i.e., for some fixed . It is impossible to compute an unbiased stochastic gradient of the objective in (7) based on a mini-batch of examples that include only a part of negative examples.
To address the above challenge, we define new formulations for OPAUC maximization by leveraging the DRO. In particular, we define a robust loss for each positive data by
Then we define the following objective for OPAUC maximization:
When , we refer to the above estimator (i.e., the objective function) as CVaR-based OPAUC estimator; and when , we refer to the above estimator as KLDRO-based OPAUC estimator. Below, we present two theorems to state the equivalent form of the objective, and the relationship between the two estimators and the surrogate objective in (7) of OPAUC.
Remark: The above theorem indicates that CVaR-based OPAUC estimator is an exact estimator of OPAUC, which is consistent for OPAUC maximization. The variable can be considered as the threshold variable to select the top-ranked negative examples for each positive data.
By choosing , then the problem (8) becomes
Remark: Theorem 2 indicates that KLDRO-based OPAUC estimator is a soft estimator, which interpolates between OPAUC() and OPAUC() by varying .
2 Optimization Algorithms and Convergence Results
In this subsection, we present the optimization algorithms for solving both (9) and (10), and then present their convergence results for finding a nearly stationary solution. The key to our development is to formulate the two optimization problems into known non-convex optimization problems that have been studied in the literature, and then to develop stochastic algorithms by borrowing the existing techniques.
Optimizing CVaR-based estimator. We first consider optimizing the CVaR-based estimator, which is equivalent to (9). A benefit for solving (9) is that an unbiased stochastic subgradient can be computed in terms of . However, this problem is still challenging because the objective function is non-smooth non-convex. In order to develop a stochastic algorithm with convergence guarantee, we prove that is weakly convex in terms of , which allows us to borrow the techniques of optimizing weakly convex function (Davis & Drusvyatskiy 2018) for solving our problem and to establish the convergence. We first establish the weak convexity of .
If is a -smooth function for any , then is -weakly convex with .
Optimizing KLDRO-based estimator of OPAUC. Next, we consider optimizing the KLDRO-based estimator, which is equivalent to (10). A nice property of the objective function is that it is smooth under a proper condition as stated in Assumption 2. However, the challenge for solving (10) is that an unbiased stochastic gradient is not readily computed. To highlight the issue, the problem (10) can be written as
There are two key differences between SOPA-s and SOPA. First, the pairwise weights in SOPA-s (step 5) are soft weights between 0 and 1, in contrast to the hard weights in SOPA. Second, the update for is a momentum-based update where . We can also use an Adam-style update, which shares similar convergence as the momentum-based update (Guo et al. 2021).
3 Convergence Analysis
A bounded smooth score function is ensured if the activation function of the neural network is smooth and the output layer uses a bounded and smooth activation function. For example, let denote the plain output of the neural network, then the score function is bounded and smooth. The Lipschitz continuity of can be guaranteed if is bounded.
We first consider the analysis of SOPA. Since is non-smooth, for presenting the convergence result, we need to introduce a convergence measure based on the Moreau envelope of given below for some :
It is guaranteed that is a smooth function (Drusvyatskiy & Paquette 2019). A point is called an -nearly stationary solution to if for some , where is the weak convexity parameter of . This convergence measure has been widely used for weakly convex optimization problems (Davis & Drusvyatskiy 2018; Rafique et al. 2020; Chen et al. 2019). Then we establish the following convergence guarantee for SOPA.
Next, we establish the convergence of SOPA-s. Under Assumption 2, we can show that in (11) is smooth. Hence, we use the standard convergence measure in terms of gradient norm of .
Remark: The convergence analysis of Algorithm 2 follows directly from that in (Wang & Yang 2022). Compared with that in Theorem 3 for SOPA, the convergence of SOPA-s is stronger than that of SOPA in several aspects: (i) the convergence measure of SOPA-s is stronger than that of SOPA due to that Theorem 4 guarantees the convergence in terms of gradient norm of the objective, while Theorem 3 guarantees the convergence on a weaker convergence measure namely the gradient norm of the Moreau envelope of the objective; (ii) the complexity of SOPA-s enjoys a parallel speed-up by using a mini-batch of data. However, it is also notable that the complexity of SOPA does not depend on the number of positive data as that of SOPA-s.
AUC meets DRO for TPAUC Maximization
In this section, we propose estimators for the surrogate objective of TPAUC and stochastic algorithms with convergence guarantee for optimizing the estimators. To this end, we apply another level of DRO on top of and define the following estimator of TPAUC:
Next, we focus on optimizing the soft estimator of TPAUC defined by using . First, we have the following form for the estimator.
When , we have
For minimizing this function, we formulate the problem as a novel three-level compositional stochastic optimization:
Under Assumption 2, Algorithm 3 with , , , ensures that after iterations we can find an nearly stationary solution of , where and .
Remark: It is notable that SOTA-s has an iteration complexity in the same order of SOPA-s for OPAUC maximization.
is a consistent surrogate function of TPAUC for and in view of the estimator given in (3).
Similar to Theorem 1, we can show that is equivalent to:
Like in (9), we can prove the inner function is weakly convex in terms of . However, computing an unbiased stochastic gradient in terms of and is also impossible due to that is inside a hinge function. To solve this problem, we can use the conjugate form of the hinge function to convert the minimization of into a weakly-convex concave min-max problem (Rafique et al. 2020) and we can develop a stochastic algorithm but only with iteration complexity for finding a nearly stationary solution. We present the algorithm and analysis in the supplement for interested readers.
Experiments
Datasets. We consider binary classification tasks on two types of datasets, namely image datasets and molecular datasets. For image datasets, we use CIFAR-10, CIFAR-100, Melanoma for experiments. For CIFAR-10 and CIFAR-100 (Krizhevsky et al. 2009), we construct imbalanced versions of the datasets by randomly removing some positive samples following (Yuan et al. 2020). Specifically, we take first half of classes as the negative class and last half of classes as the positive class, and then remove 80% samples from the positive class to make it imbalanced. The Melanoma dataset is a naturally imbalanced medical dataset which is released on Kaggle (Rotemberg et al. 2021). For molecular datasets, we use ogbg-moltox21 (the No.0 target), ogbg-molmuv (the No.1 target) and ogbg-molpbca (the No.0 target) for experiments, which are from the Stanford Open Graph Benchmark (OGB) website (Hu et al. 2020). The task on these molecular datasets is to predict certain property of molecules. The statistics for the datasets are presented in Table 5 in the supplement.
Deep Models. For image datasets, we learn convolutional neural network (CNN) and use ResNet18 (He et al. 2016) for CIFAR-10, CIFAR100 and Melanoma. For molecular datasets, we learn graph neural network (GNN) and use Graph Isomorphism Network (GIN) as the backbone model on all datasets (Xu et al. 2018), which has 5 mean-pooling layers with 64 number of hidden units and dropout rate 0.5.
Target Measures. For OPAUC maximization, we evaluate OPAUC with two FPR upper bounds, i.e., FPR and FPR separately. For TPAUC maximization, we evaluate TPAUC with two settings, i.e, FPR and TPR, and FPR and TPR.
Parameter Tuning. The learning rate of all methods is tuned in {1e-3, 1e-4, 1e-5}, except for PESG which is tuned at {1e-1, 1e-2, 1e-3} because it favors a larger learning rate. Weight decay is fixed as 2e-4. Each method is run 60 epochs in total and learning rate decays 10-fold after every 20 epochs. The mini-batch size is 64. For AUC-M, we tune the hyperparameter that controls consecutive epoch-regularization in {100, 500, 1000}. For P-push, we tune the polynomial power hyper-parameter in {2, 4, 6}. For MB that optimizes OPAUC, we tune the top proportion of negative samples in , and for MB that optimizes TPAUC we tune the top proportion of negative samples in , and tune the bottom proportion of positive samples in the range . For AW-poly, we follow (Yang et al. 2021) and tune its parameter in {101, 34, 11}. For SOPA, we tune the truncated FPR i.e. in {0.1, 0.3, 0.5}. For SOPA-s, we fix and tune the KL-regularization parameter in {0.1, 1.0, 10}, and for SOTA-s, we fix , and tune both and in {0.1, 1.0, 10}. The momentum parameter for updating in SOPA-s (i.e., ) and SOTA-s (i.e., ) is set to the default value as in the Adam optimizer, i.e., 0.1. For comparison of training convergence, the parameters are tuned according to the training performance. For comparison of testing performance, the parameters are tuned according to the validation performance. For each experiment, we repeat multiple times with different train/validation splits and random seeds, then report average and standard deviation over multiple runs.
Results. We show the plots of training convergence in Figure 2 on two image datasets (CIFAR-10, -100) and on two molecular datasets (moltox21, molpcba). From the results, we can see that SOPA-s (SOTA-s) converge always faster than MB and AW-poly for OPAUC (TPAUC) maximization. For OPAUC maximization, SOPA-s is usually faster than SOPA. More results are included in the supplement on other datasets with similar observations. The testing performance on all six datasets are shown in Table 2, 2, 4 and 4. In most cases, the proposed methods are better than the baselines. In particular, dramatic improvements have been observed on Melanoma and ogbg-molmuv datasets, which are two datasets with the highest imbalance ratios. In addition, we see that AUC maximization methods (AUC-M, AUC-SH) are not necessarily good for pAUC maximization.
Accuracy of KLDRO-based estimator. Of independent interest, we conduct simple experiments to verify the accuracy of KLDRO-based estimator of OPAUC. To this end, we compute the relative error (RE) of KLDRO-based estimator compared with the exact estimator (i.e., CVaR-based estimator). For a given upper bound of FRP we vary for 100 independently randomly generated model parameters , and the results are shown in the following figure on moltox21-t0 data (please refer to the experiments section for more information of the dataset), which demonstrates that for a given FPR there exists such that KLDRO estimator is close to the exact estimator.
Ablation Study. We also conduct some ablation study to understand the proposed algorithm SOPA-s and SOTA-s. In particular for both algorithms, we verify that tuning in SOPA-s and in SOTA-s can help further improve the performance. The results are included in the supplement.
Conclusions
In this paper, we have proposed new formulations for partial AUC maximization by using distributionally robust optimization. We propose two formulations for both one-way and two-way partial AUC, and develop stochastic algorithms with convergence guarantee for solving the two formulations, respectively. Extensive experiments on image and molecular graph datasets verify the effectiveness of the proposed algorithms.
Acknowledgements
This work is partially supported by NSF Grant 2110545, NSF Career Award 1844403, and NSF Grant 1933212. D. Zhu and X. Wu was partially supported by NSF grant CCF-1733742. We also thank anonymous reviewers for constructive comments.
References
Appendix A More Experimental Results
We present more training convergence plots on Melanoma dataset and molmuv dataset at Figure 4. For OPAUC maximization, We can observe that both our proposed SOPA-s and SOPA converge much better than AW-poly and MB method under different settings, i.e., FPR and FPR. And our proposed SOTA-s converge higher by a noticeable margin than AW-poly and MB method for TPAUC maximization all the time.
A.2 Ablation study for γ0\gamma_{0} in SOPA-s and γ0,γ1\gamma_{0},\gamma_{1} in SOTA-s
We conduct extensive ablation study for understanding the extra hyper-parameters in SOPA-s and in SOTA-s algorithms. We fix it as 0.9 for all of our experiments in the main content. But in practice, the performance would be further improved if we tune those hyper-parameters as well.
For image datasets, we conduct experiments on CIFAR-10 and CIFAR-100; for molecule datasets, we conduct experiments on ogbg-moltox21 and ogbg-molmuv. For each dataset, we first fix the best learning rate and other hyper-parameters based on our previous results in the paper. Then, for SOPA-s, we investigate at {1.0, 0.9, 0.7, 0.5, 0.3, 0.1}; for SOTA-s, we investigate both and at {1.0, 0.9, 0.7, 0.5, 0.3, 0.1}
For training aspect, we include the comparisons for SOPA-s at Figure 5; we include the comparisons for SOTA-s at Figure 6. From Figure 5, we can see that better training performance could be achieved by tuning the parameter in SOPA-s, compared with fixing it as 0.9; the similar result for SOTA-s can be also observed from Figure 6.
For testing aspect, we include the testing pAUC results at Table 6 for SOPA-s; Table 7 for SOTA-s. Both verify that tuning these parameters can further improve the performance.
Appendix B Proofs
We next present several lemmas. The first lemma is straightforward.
B.2 Proof of Lemma 7
B.3 Proof of Theorem 1
B.4 Proof of Theorem 2
When choosing , the in problem (8) becomes
With Karush-Kuhn-Tucker(KKT) conditions, it is not difficult to have . By plugging in this back we obtain the claimed objective (10). When , the above becomes the maximal one among for each . Then the object is , which is the surrogate of pAUC with FPR. When , the above becomes the average of , which gives the standard surrogate of AUC. ∎
B.5 Proof of Lemma 2
First note that , and . We prove that is weakly convex in terms of , i.e. there exists such that is jointly convex in terms of . To this end, let which is convex and Lipchitz continuous, and , which is -smooth function with respect to due to that is -smooth function. Then for any we have
where we use . The above inequality implies that is -weakly convex in terms of (Davis & Drusvyatskiy 2018), i.e., is convex. As a result is jointly convex in . Then is jointly convex in terms of .
B.6 Proof of Theorem 3
Let denote the objective function, and let . Define for some and the minimizer is denoted by . Let . Define .
Summing the above inequality over , we have
Taking expectation and re-arrange, we have
Let . As a result, we have
B.7 Proof of Theorem 4
Note that the SOPA-s algorithm is just a special case of the SOX algorithm (Wang & Yang 2022) for the more general problem and the convergence proof just follows the proof of Theorem 1 in Wang & Yang 2022.
B.8 Proof of Theorem 5
If is -Lipschitz, -smooth and are -Lipschitz, -smooth, in (12) is -smooth and .
We can show that . Thus,
Consider a sequence and the -smooth function and the step size .
where
For the gradient estimator in SOTA-s and ,
where we denote , and .
Based on the update rule , we have
Based on the Young’s inequality for products, we have for .
where we denote and . ∎
For the function value estimator in SOTA-s and ,
According to the update of in SOTA-s, we have
Denoting , we have
Re-arranging the terms and telescoping (19) from to leads to
Based on (14), the term \fontsize{7pt}{0}\fontfamily{phv}\selectfonte⃝ can be upper bounded as
Based on Lemma 2 in Wang & Yang 2022, the term \fontsize{7pt}{0}\fontfamily{phv}\selectfontg⃝ can be upper bounded as
With , based on (18), the term \fontsize{7pt}{0}\fontfamily{phv}\selectfontf⃝ can be bounded as
Plug the upper bounds of \fontsize{7pt}{0}\fontfamily{phv}\selectfontf⃝ and \fontsize{7pt}{0}\fontfamily{phv}\selectfontg⃝ into (20).
If we choose , we have
Besides, Lemma 2 in Wang & Yang 2022 and Lemma 11 imply that
Set , , , and
Appendix C Optimization of CVaR-estimator of TPAUC
It is not difficult to show that the above estimator is equivalent to
The reason is that . Using the conjugate of , we have
Define such that . Based on a minibatch , we can estimate by and . We consider the function , which can be proved to be weakly convex w.r.t. and concave w.r.t. . Hence, we can use the stagewise proximal point method to solve the problem (Rafique et al. 2020). At the -th stage, we solve the following problem approximately:
Assume . Then, we have
Let . We first show that is weakly convex in terms of for any .
Under Assumption 2, then is -weakly convex in terms of for any , where is the smoothness constant of w.r.t. .
Following similar analysis of Lemma 2, we can show that is jointly convex in terms for any . Then is -weakly convex in terms of for any . ∎
Due to the update of , we have
Assume there exists such that at every stage. Let . SOTA ensures that after iterations we can find an -nearly stationary solution for .
Let , , and . Apply Lemma 13 to .
Take expectation on both sides conditioned on the randomness that occurred before the -th iteration in the -th stage.
Similarly, apply Lemma 13 to and and take the conditional expectations.
Note that . If , we have is convex w.r.t. such that . where . Note that . Adding the above inequalities for together we have
where , , . Applying the same analysis to the update of , we have
We do not assume is independent of the randomness in the updates of our algorithm. As a result,
According to the initialization in Algorithm 4, for any and we have
Note that and . It remains to apply the analysis in (Rafique et al. 2020, Theorem 4.1) to derive the convergence for the Moreau envelope of with a complexity in the order of by setting and and the total number of stages . ∎