Stochastic Optimization of Areas Under Precision-Recall Curves with Provable Convergence

1 Introduction . Although deep learning ( DL ) has achieved tremendous success in various domains , the standard DL methods have reached a plateau as the traditional objective functions in DL are no longer sufficient to model all requirements in new applications , which slows down the democratization of AI . For instance , in healthcare applications , data is often highly imbalanced , e.g. , patients suffering from rare diseases are much less than those suffering from common diseases . In these applications , accuracy ( the proportion of correctly predicted examples ) is deemed as an inappropriate metric for evaluating the performance of a classifier . Instead , area under the curve ( AUC ) , including area under ROC curve ( AUROC ) and area under the Precision-Recall curve ( AUPRC ) , is widely used for assessing the performance of a model . However , optimizing accuracy on training data does not necessarily lead to a satisfactory solution to maximizing AUC [ 12 ] . To break the bottleneck for further advancement , DL must be empowered with the capability of efficiently handling novel objectives such as AUC . Recent studies have demonstrated great success along this direction by maximizing AUROC [ 60 ] . For example , Yuan et al . [ 60 ] proposed a robust deep AUROC maximization method with provable convergence and achieved great success for classification of medical image data . However , to the best of our knowledge , novel DL by maximizing AUPRC has not yet been studied thoroughly . Previous studies [ 14 , 20 ] have found that when dealing ∗Contribute Equally . Correspondence to qi-qi @ uiowa.edu , tianbao-yang @ uiowa.edu 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) , virtual . with highly skewed datasets , Precision-Recall ( PR ) curves could give a more informative picture of an algorithm ’ s performance , which entails the development of efficient stochastic optimization algorithms for DL by maximizing AUPRC . Compared with maximizing AUROC , maximizing AUPRC is more challenging . The challenges for optimization of AUPRC are two-fold . First , the analytical form of AUPRC by definition involves a complicated integral that is not readily estimated from model predictions of training examples . In practice , AUPRC is usually computed based on some point estimators , e.g. , trapezoidal estimators and interpolation estimators of empirical curves , non-parametric average precision estimator , and parametric binomial estimator [ 3 ] . Among these estimators , non-parametric average precision ( AP ) is an unbiased estimate in the limit and can be directly computed based on the prediction scores of samples , which lends itself well to the task of model parameters optimization . Second , a surrogate function for AP is highly complicated and non-convex . In particular , an unbiased stochastic gradient is not readily computed , which makes existing stochastic algorithms such as SGD provide no convergence guarantee . Most existing works for maximizing AP-like function focus on how to compute an ( approximate ) gradient of the objective function [ 4 , 6 , 8 , 11 , 24 , 38 , 40 , 43 , 47 , 48 ] , which leave stochastic optimization of AP with provable convergence as an open question . Can we design direct stochastic optimization algorithms both in SGD-style and Adam-style for maximizing AP with provable convergence guarantee ? In this paper , we propose a systematic and principled solution for addressing this question towards maximizing AUPRC for DL . By using a surrogate loss in lieu of the indicator function in the definition of AP , we cast the objective into a sum of non-convex compositional functions , which resembles a two-level stochastic compositional optimization problem studied in the literature [ 52 , 53 ] . However , different from existing two-level stochastic compositional functions , the inner functions in our problem are dependent on the random variable of the outer level , which requires us developing a tailored stochastic update for computing an error-controlled stochastic gradient estimator . Specifically , a key feature of the proposed method is to maintain and update two scalar quantities associated with each positive example for estimating the stochastic gradient of the individual precision score at the threshold specified by its prediction score . By leveraging recent advances in stochastic compositional optimization , we propose both adaptive ( Adam-style ) and non-adaptive ( SGD-style ) algorithms , and establish their convergence under mild conditions . We conduct comprehensive empirical studies on class imbalanced graph and image datasets for learning graph neural networks and deep convolutional neural networks , respectively . We demonstrate that the proposed method can consistently outperform prior approaches in terms of AUPRC . In addition , we show that our method achieves better results when the sample distribution is highly imbalanced between classes and is insensitive to mini-batch size . 2 Related Work . AUROC Optimization . AUROC optimization 2 has attracted significant attention in the literature . Recent success of DL by optimizing AUROC on large-scale medical image data has demonstrated the importance of large-scale stochastic optimization algorithms and the necessity of accurate surrogate function [ 60 ] . Earlier papers [ 25 , 28 ] focus on learning a linear model based on the pairwise surrogate loss and could suffer from a high computational cost , which could be as high as quadratic of the size of training data . To address the computational challenge , online and stochastic optimization algorithms have been proposed [ 18 , 35 , 42 , 58 , 63 ] . Recently , [ 21 , 22 , 36 , 57 ] proposed stochastic deep AUC maximization algorithms by formulating the problem as non-convex strongly-concave minmax optimization problem , and derived fast convergence rate under PL condition , and in federated learning setting as well [ 21 ] . More recently , Yuan et al . [ 60 ] demonstrated the success of their methods on medical image classification tasks , e.g. , X-ray image classification , melanoma classification based on skin images . However , an algorithm that maximizes the AUROC might not necessarily maximize AUPRC , which entails the development of efficient algorithms for DL by maximizing AUPRC . AUPRC Optimization . AUPRC optimization is much more challenging than AUROC optimization since the objective is even not decomposable over pairs of examples . Although AUPRC optimization has been considered in the literature ( cf . [ 15 , 47 , 41 ] and references therein ) , efficient scalable algorithms for DL with provable convergence guarantee is still lacking . Some earlier works tackled 2In the literature , AUROC optimization is simply referred to as AUC optimization . this problem by using traditional optimization techniques , e.g. , hill climbing search [ 37 ] , cuttingplane method [ 61 ] , dynamic programming [ 50 ] , and by developing acceleration techniques in the framework of SVM [ 39 ] . These approaches are not scalable to big data for DL . There is a long list of studies in information retrieval [ 5 , 11 , 38 , 47 ] and computer vision [ 4 , 6 , 8 , 9 , 24 , 40 , 48 , 43 ] , which have made efforts towards maximizing the AP score . However , most of them focus on how to compute an approximate gradient of the AP function or its smooth approximation , and provide no convergence guarantee for stochastic optimization based on mini-batch averaging . Due to lack of principled design , these previous methods when applied to deep learning are sensitive to the mini-batch size [ 6 , 47 , 48 ] and usually require a large mini-batch size in order to achieve good performance . In contrast , our stochastic algorithms are designed in a principled way to guarantee convergence without requiring a large mini-batch size as confirmed by our studies as well . Recently , [ 15 ] formulates the objective function as a constrained optimization problem using a surrogate function , and then casts it into a min-max saddle-point problem , which facilitates the use of stochastic min-max algorithms . However , they do not provide any convergence analysis for AUPRC maximization . In contrast , this is the first work that directly optimizes a surrogate function of AP ( an unbaised estimator of AUPRC in the limit ) and provides theoretical convergence guarantee for the proposed stochastic algorithms . Stochastic Compositional Optimization . Optimization of a two-level compositional function in the form of Eξ [ f ( Eζ [ g ( w ; ζ ) ] ; ξ ) ] where ξ and ζ are independent random variables , or its finite-sum variant has been studied extensively in the literature [ 1 , 10 , 52 , 27 , 30 , 31 , 33 , 34 , 46 , 53 , 59 , 62 , 45 ] . In this paper , we formulate the surrogate function of AP into a similar but more complicated two-level compositional function of the form Eξ [ f ( Eζg ( w ; ζ , ξ ) ) ] , where ξ and ζ are independent and ξ has a finite support . The key difference between our formulated compositional function and the ones considered in previous work is that the inner function g ( w ; ζ , ξ ) also depends on the random variable ξ of the outer level . Such subtle difference will complicate the algorithm design and the convergence analysis as well . Nevertheless , the proposed algorithm and its convergence analysis are built on previous studies of stochastic two-level compositional optimization . 3 The Proposed Method . Notations . We consider binary classification problem . Denote by ( x , y ) a data pair , where x ∈ Rd denotes the input data and y ∈ { 1 , −1 } denotes its class label . Let h ( x ) = hw ( x ) denote the predictive function parameterized by a parameter vector w ∈ RD ( e.g. , a deep neural network ) . Denote by I ( · ) an indicator function that outputs 1 if the argument is true and zero otherwise . To facilitate the presentation , denote by X a random data , by Y its label and by F = h ( X ) its prediction score . Let D = { ( x1 , y1 ) , . . . , ( xn , yn ) } denote the set of all training examples and D+ = { xi : yi = 1 } denote the set of all positive examples . Let n+ = |D+| denote the number of positive examples . xi ∼ D means that xi is randomly sampled from D . 3.1 Background on AUPRC and its estimator AP . Following the work of Bamber [ 2 ] , AUPRC is an average of the precision weighted by the probability of a given threshold , which can be expressed as A = ∫ ∞ −∞ Pr ( Y = 1|F ≥ c ) dPr ( F ≤ c|Y = 1 ) , where Pr ( Y = 1|F ≥ c ) is the precision at the threshold value of c. The above integral is an importance-sampled Monte Carlo integral , by which we may interpret AUPRC as the fraction of positive examples among those examples whose output values exceed a randomly selected threshold c ∼ F ( X ) |Y = 1 . For a finite set of examples D = { ( xi , yi ) , i = 1 , . . . , n } with the prediction score for each example xi given by hw ( xi ) , we consider to use AP to approximate AUPRC , which is given by AP = 1 n+ n∑ i=1 I ( yi = 1 ) n∑ s=1 I ( ys = 1 ) I ( hw ( xs ) ≥ hw ( xi ) ) n∑ s=1 I ( hw ( xs ) ≥ hw ( xi ) ) , ( 1 ) where n+ denotes the number of positive examples . It can be shown that AP is an unbiased estimator in the limit n→∞ [ 3 ] . However , the non-continuous indicator function I ( hw ( xs ) ≥ hw ( xi ) ) in both numerator and denominator in ( 1 ) makes the optimization non-tractable . To tackle this , we use a loss function ` ( w ; xs , xi ) as a surrogate function of I ( hw ( xs ) ≥ hw ( xi ) ) . One can consider different surrogate losses , e.g. , hinge loss , squared hinge loss , and smoothed hinge loss , and exponential loss . In this paper , we will consider a smooth surrogate loss function to facilitate the development of an optimization algorithm , e.g. , a squared hinge loss ` ( w ; xs ; xi ) = ( max { m− ( hw ( xi ) −hw ( xs ) ) , 0 } ) 2 , where m is a margin parameter . Note that we do not require ` to be a convex function , hence one can also consider non-convex surrogate loss such as ramp loss . As a result , our problem becomes min w P ( w ) = 1 n+ ∑ xi∈D+ − n∑ s=1 I ( ys = 1 ) ` ( w ; xs ; xi ) n∑ s=1 ` ( w ; xs ; xi ) . ( 2 )

Summary: This paper studies classification in an imbalanced data setting, where the goal is to obtain large area under the precision-recall curve (AUPRC). This paper proposes a stochastic optimization method procedure for directly optimizing the AUPRC. Their procedure is based on optimizing a continuous approximation to the average precision (AP), which is obtained by substituting non-continuous indicator functions with continuous loss surrogates, and directly optimizing the resulting objective.   The optimization algorithm proposed by this paper is called SOAP, and it is essentially a first order optimization method (e.g. SGD or Adam) with a better estimate for the stochastic gradient of the objective. The key for the improved stochastic estimate is to control the variance by tracking an exponential moving average to compute terms in the denominator of the gradient. SOAP has convergence guarantees that the average gradient norm of the surrogate AP objective will go to zero.   Empirically, the authors apply SOAP to a variety of settings with imbalanced datasets such as imbalanced CIFAR, Melanoma datasets in the image classification settings, molecular prediction datasets with graph NNs, and drug discovery datasets with graph NNs. They show that SOAP can outperform competing baselines, particularly in the setting where the dataset is extremely imbalanced.