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

Abstract:
Areas under ROC (AUROC) and precision-recall curves (AUPRC) are common metrics for evaluating classification performance for imbalanced problems. Compared with AUROC, AUPRC is a more appropriate metric for highly imbalanced datasets. While stochastic optimization of AUROC has been studied extensively, principled stochastic optimization of AUPRC has been rarely explored. In this work, we propose a principled technical method to optimize AUPRC for deep learning. Our approach is based on maximizing the averaged precision (AP), which is an unbiased point estimator of AUPRC. We cast the objective into a sum of coupled compositional functions with inner functions dependent on random variables of the outer level. We propose efficient adaptive and non-adaptive stochastic algorithms named SOAP with provable convergence guarantee under mild conditions by leveraging recent advances in stochastic compositional optimization. Extensive experimental results on image and graph datasets demonstrate that our proposed method outperforms prior methods on imbalanced problems in terms of AUPRC. To the best of our knowledge, our work represents the first attempt to optimize AUPRC with provable convergence. The SOAP has been implemented in the libAUC library at https://libauc.org/.

heading: 1 Introduction
text: 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.

heading: 2 Related Work
text: 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.

heading: 3 The Proposed Method
text: 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.

heading: 3.1 Background on AUPRC and its estimator AP
text: 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)

heading: 3.2 Stochastic Optimization of AP (SOAP)
text: We cast the problem into a finite-sum of compositional functions. To this end, let us define a few notations:
g(w;xj ,xi) = [g1(w;xj ,xi), g2(w;xj ,xi)] > = [`(w;xj ,xi)I(yj = 1), `(w;xj ,xi)] > gxi(w) = Exj∼D[g(w;xj ,xi)], (3) where gxi(w) : Rd → R2. Let f(s) = − s1s2 : R 2 → R. Then, we can write the objective function for maximizing AP as a sum of compositional functions:
P (w) = 1
n+ ∑ xi∈D+ f(gxi(w)) = Exi∼D+ [f(gxi(w))]. (4)
We refer to the above problem as an instance of two-level stochastic coupled compositional functions. It is similar to the two-level stochastic compositional functions considered in literature [52, 53] but with a subtle difference. The difference is that in our formulation the inner function gxi(w) = Exj∼D[g(w;xj ,xi)] depends on the random variable xi of the outer level. This difference makes the proposed algorithm slightly complicated by estimating gxi(w) separately for each positive example. It also complicates the analysis of the proposed algorithms. Nevertheless, we can still employ the techniques developed for optimizing stochastic compositional functions to design the algorithms and develop the analysis for optimizing the objective (4).
In order to motivate the proposed method, let us consider how to compute the gradient of P (w). Let the gradient of gxi(w) be denoted by∇wgxi(w)> = (∇w[gxi(w)]1,∇w[gxi(w)]2). Then we have
∇wP (w) = 1
n+ ∑ xi∈D+ ∇wgxi(w)>∇f(gxi(w))
= 1
n+ ∑ xi∈D+ ∇wgxi(w)> ( −1 [gxi(w)]2 , [gxi(w)]1 ([gxi(w)]2) 2 )> .
(5)
The major cost for computing∇wP (w) lies at evaluating gxi(w) and its gradient∇wgxi(w), which involves passing through all examples in D. To this end, we will approximate these quantities by stochastic samples. The gradient∇wgxi(w) can be simply approximated by the stochastic gradient, i.e.,
∇̂wgxi(w) =
( 1 B ∑ xj∈B I(yj = 1)∇`(w;xj ,xi) 1 B ∑ xj∈B∇`(w;xj ,xi) ) , (6)
where B denote a set of B random samples from D. For estimating gxi(w) = Exj∼Dg(w;xj ,xi), however, we need to ensure its approximation error is controllable due to the compositional structure such that the convergence can be guaranteed. We borrow a technique from the literature of stochastic compositional optimization [52] by using moving average estimator for estimating gxi(w) for all positive examples. To this end, we will maintain a matrix u = [u1,u2] with each column indexable by any positive example, i.e., u1xi ,u 2 xi correspond to the moving average estimator of [gxi(w)]1 and [gxi(w)]2, respectively. The matrix u is updated by the subroutine UG in Algorithm 2, where γ ∈ (0, 1) is a parameter. It is notable that in Step 3 of Algorithm 2, we clip the moving average update of u2xi by a lower bound u0, which is a given parameter. This step can ensure the division in computing the stochastic gradient estimator in (7) always valid and is also important for convergence
Algorithm 1: SOAP 1: Input: γ, α, u0, and other parameters for SGD-stype update or Adam-stype update. 2: Initialize w1 ∈ Rd, u ∈ R|n+|×2 3: for t = 1, . . . , T do 4: Draw a batch of B+ positive samples denoted by B+. 5: Draw a batch of B samples denoted by B. 6: u = UG(B,B+,u,wt, γ, u0) 7: Compute (biased) Stochastic Gradient Estimator
G(wt) = 1
B+ ∑ xi∈B+ ∑ xj∈B (u1xi − u 2 xiI(yj = 1))∇`(w;xj ,xi) B(u2xi) 2
(7)
8: Update wt+1 by a SGD-style method or by a Adam-style method
wt+1 = UW(wt, G(wt))
9: end for 10: Return: last solution.
analysis. With these stochastic estimators, we can compute an estimate of∇P (w) by equation (7), where B+ includes a batch of sampled positive data. With this stochastic gradient estimator, we can employ SGD-style method and Adam-style shown in Algorithm 3 to update the model parameter w. The final algorithm named as SOAP is presented in Algorithm 1.
Algorithm 2: UG(B,B+,u,wt, γ, u0) 1: for each positive xi ∈ B+ do 2: Compute
[g̃xi(wt)]1 = 1 |B| ∑ xj∈B yj=1 `(wt;xj ,xi)
[g̃xi(wt)]2 = 1 |B| ∑ xj∈B `(wt;xj ,xi)
3: Compute u1xi = (1− γ)u 1 xi + γ[g̃xi(wt)]1
u2xi = max((1−γ)u 2 xi +γ[g̃xi(wt)]2, u0)
4: end for 5: Return u
Algorithm 3: UW(wt, G(wt)) 1: Option 1: SGD-style update (paras: α)
wt+1 = wt − αG(wt) 2: Option 2: Adam-style update (paras: α, , η1, η2)
ht+1 = η1ht + (1− η1)G(wt) vt+1 = η2v̂t + (1− η2)(G(wt))2
wt+1 = wt − α ht+1√ + v̂t+1
where v̂t = vt (Adam) or v̂t = max(v̂t−1, vt) (AMSGrad)
3: Return: wt+1

heading: 3.3 Convergence Analysis
text: In this subsection, we present the convergence results of SOAP and also highlight its convergence analysis. To this end, we first present the following assumption. Assumption 1. Assume that (a) there exists ∆1 such that P (w1) −minw P (w) ≤ ∆1; (b) there exist C,M > 0 such that `(w;xi,xi) ≥ C for any xi ∈ D+, `(w;xj ,xi) ≤M , and `(w;xj ,xi) is Lipscthiz continuous and smooth with respect to w for any xi ∈ D+,xj ∈ D; (c) there exists V > 0 such that Exj∼D[‖g(w;xj ,xi)− gxi(w)‖2] ≤ V , and Exj∼D[‖∇g(w;xj ,xi)−∇gxi(w)‖2] ≤ V for any xi. With a bounded score function hw(x) the above assumption can be easily satisfied. Based on the above assumption, we can prove that the objective function P (w) is smooth. Lemma 1. Suppose Assumption 1 holds, then there exists L > 0 such that P (·) is L-smooth. In addition, there exists u0 ≥ C/n such that gxi(w) ∈ Ω = {u ∈ R2, 0 ≤ [u]1 ≤ M,u0 ≤ [u]2 ≤ M}, ∀xi ∈ D+.
Next, we highlight the convergence analysis of SOAP employing the SGD-stype update and include that for employing Adam-style update in the supplement. Without loss of generality, we assume |B+| = 1 and the positive sample in B+ is randomly selected from D+ with replacement. When the context is clear, we abuse the notations gi(w) and ui to denote gxi(w) and uxi below, respectively. We first establish the following lemma following the analysis of non-convex optimization.
Lemma 2. With α ≤ 1/2, running T iterations of SOAP (SGD-style) updates, we have α
2 E[ T∑ t=1 ‖∇P (wt)‖2] ≤ E[ ∑ t (P (wt)− P (wt+1))] + αC1 2 E[ T∑ t=1 ‖git(wt)− uit‖2] + α2TC2,
where it denotes the index of the sampled positive data at iteration t, C1 and C2 are proper constants.
Our key contribution is the following lemma that bounds the second term in the above upper bound. Lemma 3. Suppose Assumption 1 holds, with u initialized by (6) for every xi ∈ D+ we have
E[ T∑ t=1 ‖git(wt)− uit‖2] ≤ n+V γ + γV T + 2 n2+α 2TC3 γ2 , (8)
where C3 is a proper constant. Remark: The innovation of proving the above lemma is by grouping uit , t = 1, . . . , T into n+ groups corresponding to the n+ positive examples, and then establishing the recursion of the error ‖git(wt)− uit‖2 within each group, and then summing up these recursions together. Based on the two lemmas above, we establish the following convergence of SOAP with a SGD-style update.
Theorem 1. Suppose Assumption 1 holds, let the parameters be α = 1 n 2/5 + T 3/5 ,γ =
n 2/5 + T 2/5 , ∀ t ∈
1, · · · , T , and T > n+. Then after running T iterations, SOAP with a SGD-style update satisfies E [
1 T T∑ t=1 ‖∇P (wt)‖2
] ≤ O( n 2/5 +
T 2/5 ), where O suppresses constant numbers.
Remark: To the best of our knowledge, this is the first time a stochastic algorithm was proved to converge for AP maximization.
Similarly, we can establish the following convergence of SOAP by employing an Adam-style update, specifically the AMSGrad update. Theorem 2. Suppose Assumption 1 holds, let the parameters η1 ≤ √ η2 ≤ 1, α = 1
n 2/5 + T
3/5 ,γ =
n 2/5 + T 2/5 , ∀ t ∈ 1, · · · , T , and T > n+. Then after running T iterations, SOAP with an AMSGRAD update satisfies E [
1 T T∑ t=1 ‖∇P (wt)‖2
] ≤ O( n 2/5 +
T 2/5 ), where O suppresses constant numbers.

heading: 4 Experiments
text: In this section, we evaluate the proposed method through comprehensive experiments on imbalanced datasets. We show that the proposed method can outperform prior state-of-the-art methods for imbalanced classification problems. In addition, we conduct experiments on (i) the effects of imbalance ratio; (ii) the insensitivity to batch size and (iii) the convergence speed on testing data; and observe that our method (i) is more advantageous when data is more imbalanced, (ii) is not sensitive to batch size, and (iii) converges faster than baseline methods.
Our proposed optimization algorithm is independent of specific datasets and tasks. Therefore, we perform experiments on both graph and image prediction tasks. In particular, the graph prediction tasks in the contexts of molecular property prediction and drug discovery suffer from very severe imbalance problems as positive labels are very rare while negative samples are abundantly available. Thus, we choose to use graph data intensively in our experiments. Additionally, the graph data we use allow us to vary the imbalance ratio to observe the performance change of different methods.
In all experiments, we compare our method with the following baseline methods. CB-CE refers to a method using a class-balanced weighed cross entropy loss function, in which the weights for positive and negative samples are adjusted with the strategy proposed by Cui et al. [13]. Focal is to up-weight the penalty on hard examples using focal loss [32]. LDAM refers to training with labeldistribution-aware margin loss [7]. AUC-M is an AUROC maximization method using a surrogate loss [60]. In addition, we compare with three methods for optimizing AUPRC or AP, namely, the MinMax method [15] - a method for optimizing a discrete approximation of AUPRC, SmoothAP [4] - a method that optimizes a smoothed approximation of AP, and FastAP - a method that uses soft histogram binning to approximate the gradient of AP [6]. For all of these methods, we use the
SGD-style with momentum optimization for image prediction tasks and the Adam-style optimization algorithms for graph prediction tasks and unless specified otherwise. We refer to imbalance ratio as the number of positive samples over the total number of examples of a considered set. The hyper-parameters of all methods are fine tuned using cross-validation with training/validation splits mentioned below. For AP maximization methods, we use a sigmoid function to produce the prediction score. For simplicity, we set u0 = 0 for SOAP and encounter no numerical problems in experiments. As SOAP requires positive samples for updating u to approximate the gradient of surrogate objective, we use a data sampler which samples a few positive examples (e.g., 2) and some negative examples per iteration. The same sampler applies to all methods for fair comparison. The code for reproducing the results is released here [44].

heading: 4.1 Image Classification
text: Data. We first conduct experiments on three image datasets: CIFAR10, CIFAR100 and Melanoma dataset [49]. We construct imbalanced version of CIFAR10 and CIFAR100 for binary classification. In particular, for each dataset we manually take the last half of classes as positive class and first half of classes as negative class. To construct highly imbalanced data, we remove 98% of the positive images from the training data and keep the test data unchanged (i.e., the testing data is still balanced). And we split the training dataset into train/validation set at 80%/20% ratio. The Melanoma dataset is from a medical image Kaggle competition, which serves as a natural real imbalanced image dataset. It contains 33,126 labeled medical images, among which 584 images are related to malignant melanoma and labelled as positive samples. Since the test set used by Kaggle organization is not available, we manually split the training data into train/validation/test set at 80%/10%/10% ratio and report the achieved AUPRC on the test set by our method and baselines. The images of Melanoma dataset are always resized to have a resolution of 384× 384 in our experiments. Setup. We use two ResNet [23] models, i.e., ResNet18 and ResNet34, as the backbone networks for image classification. For all methods except for CE, the ResNet models are initialized with a model pre-trained by CE with a SGD optimizer. We tune the learning rate in a range {1e-5, 1e-4, 1e-3, 1e-2} and the weight decay parameter in a range {1e-6, 1e-5, 1e-4}. Then the last fully connected layer is randomly re-initialized and the network is trained by different methods with the same weight decay parameter but other hyper-parameters individually tuned for fair comparison, e.g., we tune γ of SOAP in a range {0.9, 0.99,0.999}, and tune m in {0.5, 1, 2, 5, 10}. We refer to this scheme as two-stage training, which is widely used for imbalanced data [60]. We consistently observe that this strategy can bring the model to a good initialization state and improve the final performance of our method and baselines.
Results. Table 1 shows the AUPRC on testing sets of CIFAR-10 and CIFAR-100. We report the results on Melanoma in Table 3. We can observe that the proposed method SOAP outperforms all baselines. It is also striking to see that on Melanoma dataset, our proposed SOAP can outperform all baselines by a large margin, and all other methods have very poor performance. The reason is that the testing set of Melanoma is also imbalanced (imbalanced ratio=1.72%), while the testing sets of CIFAR-10 and CIFAR-100 are balanced. We also observe that the AUROC maximization (AUC-M) does not necessarily optimize AUPRC. We also plot the final PR curves in Figure 3 in the supplement.

heading: 4.2 Graph Classification for Molecular Property Prediction
text: Data. To further demonstrate the advantages of our method, we conduct experiments on two graph classification datasets. We use the datasets HIV and MUV from the MoleculeNet [55], which is a benchmark for molecular property prediction. The HIV dataset has 41,913 molecules from the Drug Therapeutics Program (DTP), and the positive samples are molecules tested to have inhibition ability to HIV. The MUV dataset has 93,127 molecules from the PubChem library, and molecules are labelled by whether a bioassay property exists or not. Note that the MUV dataset provides labels of 17 properties in total and we only conduct experiments to predict the third property as this property is more imbalanced. The percentage of positive samples in HIV and MUV datasets are 3.51% and 0.20%, respectively. We use the split of train/validation/test set provided by MoleculeNet. Molecules are treated as 2D graphs in our experiments, and we use the feature extraction procedure of MoleculeKit [54] to obtain node features of graphs. The same data preprocessing is used for all of our experiments on graph data.
Setup. Many recent studies have shown that graph neural networks (GNNs) are powerful models for graph data analysis [29, 17, 16]. Hence, we use three different GNNs as the backbone network for graph classification, including the message passing neural network (MPNN) [19], an invariant of graph isomorphism network [56] named by GINE [26], and the multi-level message passing neural network (ML-MPNN) proposed by Wang et al. [54]. We use the same two-stage training scheme with a similar hyper-parameter tuning. We pre-train the networks by Adam with 100 epochs and a tuned initial learning rate 0.0005, which is decayed by half after 50 epochs.
Results. The achieved AUPRC on the test set by all methods are presented in Table 2. Results show that our method can outperform all baselines by a large margin in terms of AUPRC, regardless of which model structure is used. These results clearly demonstrate that our method is effective for classification problems in which the sample distribution is highly imbalanced between classes.

heading: 4.3 Graph Classification for Drug Discovery
text: Data. In addition to molecular property prediction, we explore applying our method to drug discovery. Recent studies have shown that GNNs are effective in drug discovery through predicting the antibacterial property of chemical compounds [51]. Such application scenarios involves training a GNN model on labeled datasets and making predictions on a large library of chemical compounds so as to discover new antibiotic. However, because the positive samples in the training data, i.e., compounds known to have antibacterial property, are very rare, there exists very severe class imbalance.
We show that our method can serve as a useful solution to the above problem. We conduct experiments on the MIT AICURES dataset from an open challenge (https://www.aicures.mit.edu/tasks)
in drug discovery. The dataset consists of 2097 molecules. There are 48 positive samples that have antibacterial activity to Pseudomonas aeruginosa, which is the pathogen leading to secondary lungs infections of COVID-19 patients. We conduct experiments on three random train/validation/test splits at 80%/10%/10% ratio, and report the average AUPRC on the test set over three splits.
Setup. Following the setup in Sec. 4.2, we use three GNNs: MPNN, GINE and ML-MPNN. We use the same two-stage training scheme with a similar hyper-parameter tuning. We pre-train GNNs by the Adam method for 100 epochs with a batch size of 64 and a tuned learning rate of 0.0005, which is decayed by half at the 50th epoch. Due to the limit of space, Table 3 only reports GINE and MPNN results. Please refer to Table 6 in the supplement for the full results of all three GNNs.
Results. The average test AUPRC from three independent runs over three splits are summarized in Table 3, Table 6. We can see that our SOAP can consistently outperform all baselines on all three GNN models. Our proposed optimization method can significantly improve the achieved AUPRC of GNN models, indicating that models tend to assign higher confidence scores to molecules with antibacterial activity. This can help identify a larger number of candidate drugs.
We have employed the proposed AUPRC maximization method for improving the testing performance on MIT AICures Challenge and achieved the 1st place. For details, please refer to [54].

heading: 4.4 Ablation Studies
text: Effects of Imbalance Ratio. We now study the effects of imbalance ratio on the performance improvements of our method. We use two datasets Tox21 and ToxCast from the MoleculeNet [55]. The Tox21 and ToxCast contain 8014 and 8589 molecules, respectively. There are 12 property prediction tasks in Tox21, and we conduct experiments on Task 0 and Task 2. Similarly, we select Task 12 and Task 8 of ToxCast for experiments. We use the split of train/validation/test set provided by MoleculeNet. The imbalanced ratios on the training sets are 4.14% for Task 0 of Tox21, 12.00% for Task 2 of Tox21, 2.97% for Task 12 of ToxCast, 8.67% for Task 8 of ToxCast.
Following Sec. 4.2, we test three neural network models MPNN, GINE and ML-MPNN. The hyperparameters for training models are also the same as those in Sec. 4.2. We present the results of Tox21
and ToxCast in Table 5 in the supplement. Our SOAP can consistently achieve improved performance when the data is extremely imbalanced. However, it sometimes fails to do so if the imbalance ratio is not too low. Clearly, the improvements from our method are higher when the imbalance ratio of labels is lower. In other words, our method is more advantageous for data with extreme class imbalance.
Insensitivity to Batch Size. We conduct experiments on CIFAR-10 and CIFAR-100 data by varying the mini-batch size for the SOAP algorithm and report results in Figure 2 (Left most). We can see that SOAP is not sensitive to the mini-batch size. This is consistent with our theory. In contrast, many previous methods for AP maximization are sensitive to the mini-batch size [47, 48, 6].
Convergence Speed. We report the convergence curves of different methods for maximizing AUPRC or AP in Figure 1 on different datasets. We can see that the proposed SOAP algorithms converge much faster than other baseline methods.
More Surrogate Losses. To verify the generality of SOAP, we evaluate the performance of SOAP with two more different surrogate loss functions `(w;xs,xi) as a surrogate function of the indicator I(hw(xs) ≥ hw(xi)), namely, the logistic loss, `(w;xs,xi) = − log 11+exp(−c(`(hw(xi)−hw(xs))) , and the sigmoid loss, `(w;xs,xi) = 11+exp(c(`(hw(xi)−hw(xs))) where c is a hyperparameter. We tune c ∈ {1, 2} in our experiments. We conduct experiments on CIFAR10, CIFAR100 following the experimental setting in Section 4.1 for the image data. For the graph data, we conduct experiments on HIV, MUV data following the experimental setting in Section 4.2. We report the results in Table 4. We can observe that SOAP has similar results with different surrogate loss functions.
Consistency. Finally, we show the consistency between the Surrogate Objective -P (w) and AP by plotting the convergence curves on different datasets in Figure 2 (Right two). It is obvious two see the consistency between our surrogate objective and the true AP.

heading: 5 Conclusions and Outlook
text: In this work, we have proposed a stochastic method to optimize AUPRC that can be used in deep learning for tackling highly imbalanced data. Our approach is based on maximizing the averaged precision, and we cast the objective into a sum of coupled compositional functions. We proposed efficient adaptive and non-adaptive stochastic algorithms with provable convergence guarantee to compute the solutions. Extensive experimental results on graph and image datasets demonstrate that our proposed method can achieve promising results, especially when the class distribution is highly imbalanced. One limitation of SOAP is its convergence rate is still slow. In the future, we will consider to improve the convergence rate to address the limitation of the present work.

heading: Acknowledgments
text: We thank Bokun Wang for discussing the proofs, and thank anonymous reviewers for constructive comments. Q.Q contributed to the algorithm design, analysis, and experiments under supervision of T.Y. Y.L and Z.X contributed to the experiments under supervision of S.J. Q.Q and T.Y were partially supported by NSF Career Award #1844403, NSF Award #2110545 and NSF Award #1933212. Y.L, Z.X and S.J were partially supported by NSF IIS-1955189.

Review summaries:
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. 

This paper investigates the imbalanced classification problems under the areas under precision-recall curves (AUPRC). To handle the nondifferentiability of the indicator functions in  the averaged precision (AP) estimators, they propose a differentiable surrogate functions such as squared hinge loss to replace the indicator function, leading to a compositional optimization problem. To solve this compositional problem, the authors use the moving average forward estimators in Wang et al., 2017, to estimate g_{xi}(w) in the gradient estimator of the objective function. For the outer updates, the authors introduce both the SGD-type and Adam-type gradient updates. Theoretically, this paper is the first one to provide the finite-time convergence analysis as well as a computational complexity for stochastic optimization of AUPRC. I am happy to see that the proposed method perform quite well in various imbalanced classification problems. The experiments are clear, comprehensive and involve a variety of datasets like CIFAR10, 100, HIV, MUV, and data for drug discovery, and test different types of neural networks such as ResNet. Under such a setup, they show their method outperforms existing state-of-the-arts via a notable margin (particularly for those highly imbalanced datasets). 

The author presented a technique for optimizing the Area under the Precision-Recall curve (AUCPR). The proposed technique works by replacing the indicator function in the empirical AUCPR formulation with a smooth loss function and then decompose the formula as a sum of compositional functions. The author then proposed a learning algorithm based on stochastic compositional optimization. Finally, the author shows the effectiveness of the proposed model on various classification tasks.