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∈[α,β]\in[\alpha,\beta] (Figure 1 middle) and TPAUC puts a restriction on the lower bound of TPR and the upper bound of FPR, i.e., TPR≥α\geq\alpha, FPR≤β\leq\beta (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 FPR∈[0,β]\text{FPR}\in[0,\beta] and for TPAUC we focus on maximizing pAUC in the region where FPR≤β\text{FPR}\leq\beta and TPR≥α\text{TPR}\geq\alpha for some α,β∈(0,1)\alpha,\beta\in(0,1). 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 O(1/B)O(1/\sqrt{B}), where BB 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 (α,β)(\alpha,\beta) 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 β\beta and TPR larger than α\alpha. 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 (α,β)(\alpha,\beta) where α>0\alpha>0; 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 (0,β)(0,\beta); (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 F(w)F(\mathbf{w}) is weakly convex if there exists C>0C>0 such that F(w)+C2∥w∥2F(\mathbf{w})+\frac{C}{2}\|\mathbf{w}\|^{2} is a convex function. A function F(w)F(\mathbf{w}) is LL-smooth if its gradient is Lipchitz continuous, i.e., ∥∇F(w)−∇F(w′)∥≤L∥w−w′∥\|\nabla F(\mathbf{w})-\nabla F(\mathbf{w}^{\prime})\|\leq L\|\mathbf{w}-\mathbf{w}^{\prime}\|.

where h(x−)∈[FPR−1(α1),FPR−1(α0)]h(\mathbf{x}_{-})\in[\text{FPR}^{-1}(\alpha_{1}),\text{FPR}^{-1}(\alpha_{0})] 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 k1=⌈n−α0⌉,k2=⌊n−α1⌋k_{1}=\lceil n_{-}\alpha_{0}\rceil,k_{2}=\lfloor n_{-}\alpha_{1}\rfloor. In this work, we will focus on optimizing OPAUC^(h,0,β)\widehat{\text{OPAUC}}(h,0,\beta) for some β∈(0,1)\beta\in(0,1).

Similarly, a non-parametric estimator of TPAUC with FPR≤β,TPR≥α\text{FPR}\leq\beta,\text{TPR}\geq\alpha is given by

where k1=⌊n+α⌋,k2=⌊n−β⌋k_{1}=\lfloor n_{+}\alpha\rfloor,k_{2}=\lfloor n_{-}\beta\rfloor.

By using the CVaR divergence ϕc(t)\phi_{c}(t) for some γ\gamma such that nγn\gamma 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 S−\mathcal{S}_{-}, i.e., S−↓[1,k]\mathcal{S}_{-}^{\downarrow}[1,k] for some fixed kk. 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 ϕ(⋅)=ϕc(⋅)\phi(\cdot)=\phi_{c}(\cdot), we refer to the above estimator (i.e., the objective function) as CVaR-based OPAUC estimator; and when ϕ(⋅)=ϕkl(⋅)\phi(\cdot)=\phi_{kl}(\cdot), 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 sis_{i} can be considered as the threshold variable to select the top-ranked negative examples for each positive data.

By choosing ϕ(⋅)=ϕkl(⋅)\phi(\cdot)=\phi_{kl}(\cdot), then the problem (8) becomes

Remark: Theorem 2 indicates that KLDRO-based OPAUC estimator is a soft estimator, which interpolates between OPAUC(hw,0,1/n−h_{\mathbf{w}},0,1/n_{-}) and OPAUC(hw,0,1h_{\mathbf{w}},0,1) by varying λ\lambda.

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 (w,s)(\mathbf{w},\mathbf{s}). However, this problem is still challenging because the objective function F(w,s)F(\mathbf{w},\mathbf{s}) is non-smooth non-convex. In order to develop a stochastic algorithm with convergence guarantee, we prove that F(w,s)F(\mathbf{w},\mathbf{s}) is weakly convex in terms of (w,s)(\mathbf{w},\mathbf{s}), 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 F(w,s)F(\mathbf{w},\mathbf{s}).

If L(⋅;xi,xj)L(\cdot;\mathbf{x}_{i},\mathbf{x}_{j}) is a LsL_{s}-smooth function for any xi,xj\mathbf{x}_{i},\mathbf{x}_{j}, then F(w,s)F(\mathbf{w},\mathbf{s}) is ρ\rho-weakly convex with ρ=Ls/β\rho=L_{s}/\beta.

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 pijp_{ij} in SOPA-s (step 5) are soft weights between 0 and 1, in contrast to the hard weights pij∈{0,1}p_{ij}\in\{0,1\} in SOPA. Second, the update for wt+1\mathbf{w}_{t+1} is a momentum-based update where γ1∈(0,1)\gamma_{1}\in(0,1). 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 h(⋅;x)h(\cdot;\mathbf{x}) 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 h^(w;x)\hat{h}(\mathbf{w};\mathbf{x}) denote the plain output of the neural network, then the score function h(w;x)=1/(1+exp⁡(−h^(w;x))h(\mathbf{w};\mathbf{x})=1/(1+\exp(-\hat{h}(\mathbf{w};\mathbf{x})) is bounded and smooth. The Lipschitz continuity of h(w;x)h(\mathbf{w};\mathbf{x}) can be guaranteed if w\mathbf{w} is bounded.

We first consider the analysis of SOPA. Since F(w,s)F(\mathbf{w},\mathbf{s}) is non-smooth, for presenting the convergence result, we need to introduce a convergence measure based on the Moreau envelope of F(w,s)F(\mathbf{w},\mathbf{s}) given below for some ρ^>ρ\hat{\rho}>\rho:

It is guaranteed that Fρ^(w,s)F_{\hat{\rho}}(\mathbf{w},\mathbf{s}) is a smooth function (Drusvyatskiy & Paquette 2019). A point (w,s)(\mathbf{w},\mathbf{s}) is called an ϵ\epsilon-nearly stationary solution to F(w,s)F(\mathbf{w},\mathbf{s}) if ∥∇Fρ^(w,s)∥≤ϵ\|\nabla F_{\hat{\rho}}(\mathbf{w},\mathbf{s})\|\leq\epsilon for some ρ^>ρ\hat{\rho}>\rho, where ρ\rho is the weak convexity parameter of FF. 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 F(w)F(\mathbf{w}) in (11) is smooth. Hence, we use the standard convergence measure in terms of gradient norm of F(w)F(\mathbf{w}).

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 L^ϕ(xi,w),xi∈S+\hat{L}_{\phi}(\mathbf{x}_{i},\mathbf{w}),\mathbf{x}_{i}\in\mathcal{S}_{+} and define the following estimator of TPAUC:

Next, we focus on optimizing the soft estimator of TPAUC defined by using ϕ=ϕ′=ϕkl\phi=\phi^{\prime}=\phi_{kl}. First, we have the following form for the estimator.

When ϕ=ϕ′=ϕkl\phi=\phi^{\prime}=\phi_{kl}, we have

For minimizing this function, we formulate the problem as a novel three-level compositional stochastic optimization:

Under Assumption 2, Algorithm 3 with γ0=O(B−ϵ2)\gamma_{0}=O(B_{-}\epsilon^{2}), γ1=O(B+ϵ2)\gamma_{1}=O(B_{+}\epsilon^{2}), γ2=O(min⁡{B−,B+}ϵ2)\gamma_{2}=O(\min\{B_{-},B_{+}\}\epsilon^{2}), η=O(min⁡{γ0B1/n+,γ1,γ2})\eta=O(\min\{\gamma_{0}B_{1}/n_{+},\gamma_{1},\gamma_{2}\}) ensures that after T=O(1min⁡(B+,B−)ϵ4+n+B+B−ϵ4)T=O(\frac{1}{\min(B_{+},B_{-})\epsilon^{4}}+\frac{n_{+}}{B_{+}B_{-}\epsilon^{4}}) iterations we can find an ϵ\epsilon nearly stationary solution of F(w)F(\mathbf{w}), where B+=∣B+∣B_{+}=|\mathcal{B}_{+}| and B−=∣B−∣B_{-}=|\mathcal{B}_{-}|.

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 TPR≥α\text{TPR}\geq\alpha and FPR≤β\text{FPR}\leq\beta in view of the estimator TPAUC^\widehat{\text{TPAUC}} given in (3).

Similar to Theorem 1, we can show that F(w;ϕc,ϕc′)F(\mathbf{w};\phi_{c},\phi_{c}^{\prime}) is equivalent to:

Like F(w,s)F(\mathbf{w},\mathbf{s}) in (9), we can prove the inner function is weakly convex in terms of (w,s,s′)(\mathbf{w},\mathbf{s},s^{\prime}). However, computing an unbiased stochastic gradient in terms of w\mathbf{w} and sis_{i} is also impossible due to that ψi(w;si)\psi_{i}(\mathbf{w};s_{i}) is inside a hinge function. To solve this problem, we can use the conjugate form of the hinge function to convert the minimization of F(w;ϕc,ϕc′)F(\mathbf{w};\phi_{c},\phi^{\prime}_{c}) into a weakly-convex concave min-max problem (Rafique et al. 2020) and we can develop a stochastic algorithm but only with O(1/ϵ6)O(1/\epsilon^{6}) 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 ≤0.3\leq 0.3 and FPR≤0.5\leq 0.5 separately. For TPAUC maximization, we evaluate TPAUC with two settings, i.e, FPR≤0.4\leq 0.4 and TPR≥0.6\geq 0.6, and FPR≤0.5\leq 0.5 and TPR≥0.5\geq 0.5.

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 γ\gamma 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 {10%,30%,50%}\{10\%,30\%,50\%\}, and for MB that optimizes TPAUC we tune the top proportion of negative samples in {30%,40%,50%}\{30\%,40\%,50\%\}, and tune the bottom proportion of positive samples in the range {30%,40%,50%}\{30\%,40\%,50\%\}. For AW-poly, we follow (Yang et al. 2021) and tune its parameter γ\gamma in {101, 34, 11}. For SOPA, we tune the truncated FPR i.e. β\beta in {0.1, 0.3, 0.5}. For SOPA-s, we fix γ0=0.9\gamma_{0}=0.9 and tune the KL-regularization parameter λ\lambda in {0.1, 1.0, 10}, and for SOTA-s, we fix γ0=γ1=0.9\gamma_{0}=\gamma_{1}=0.9, and tune both λ\lambda and λ′\lambda^{\prime} in {0.1, 1.0, 10}. The momentum parameter for updating vt\mathbf{v}_{t} in SOPA-s (i.e., 1−γ11-\gamma_{1}) and SOTA-s (i.e., 1−γ21-\gamma_{2}) 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 λ\lambda for 100 independently randomly generated model parameters w\mathbf{w}, 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 λ\lambda 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 γ0\gamma_{0} in SOPA-s and γ0,γ1\gamma_{0},\gamma_{1} 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≤0.3\leq 0.3 and FPR≤0.5\leq 0.5. 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 γ0\gamma_{0} in SOPA-s and γ0,γ1\gamma_{0},\gamma_{1} 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 γ0\gamma_{0} at {1.0, 0.9, 0.7, 0.5, 0.3, 0.1}; for SOTA-s, we investigate both γ0\gamma_{0} and γ1\gamma_{1} 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 γ0\gamma_{0} 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 ϕ(⋅)=ϕkl(⋅)\phi(\cdot)=\phi_{kl}(\cdot), the L^ϕ(w;xi)\hat{L}_{\phi}(\mathbf{w};\mathbf{x}_{i}) in problem (8) becomes

With Karush-Kuhn-Tucker(KKT) conditions, it is not difficult to have pj⋆=exp⁡(L(w;xi,xj)/λ∑jexp⁡(L(w;xi,xj)/λp_{j}^{\star}=\frac{\exp(L(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j})/\lambda}{\sum_{j}\exp(L(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j})/\lambda}. By plugging in this back we obtain the claimed objective (10). When λ=0\lambda=0, the above becomes the maximal one among {L(w;xi,xj),xj∈S+}\{L(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j}),\mathbf{x}_{j}\in\mathcal{S}_{+}\} for each xi\mathbf{x}_{i}. Then the object is 1n+max⁡xj∈S−L(w;xi,xj)\frac{1}{n_{+}}\max_{\mathbf{x}_{j}\in\mathcal{S}_{-}}L(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j}), which is the surrogate of pAUC with FPR≤1/n−\leq 1/n_{-}. When λ=∞\lambda=\infty, the above becomes the average of {L(w;xi,xj),xj∈S+}\{L(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j}),\mathbf{x}_{j}\in\mathcal{S}_{+}\}, which gives the standard surrogate of AUC. ∎

B.5 Proof of Lemma 2

First note that F(w,s)=1n+∑xi∈S+(si+1βgi(w,si))F(\mathbf{w},\mathbf{s})=\frac{1}{n_{+}}\sum_{\mathbf{x}_{i}\in\mathcal{S}_{+}}\left(s_{i}+\frac{1}{\beta}g_{i}(\mathbf{w},s_{i})\right), and gi(w,si)=1n−∑xj∈S−(L(w;xi,xj)−si)+g_{i}(\mathbf{w},s_{i})=\frac{1}{n_{-}}\sum_{\mathbf{x}_{j}\in\mathcal{S}_{-}}(L(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j})-s_{i})_{+}. We prove that (L(w;xi,xj)−si)+(L(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j})-s_{i})_{+} is weakly convex in terms of (w,si)(\mathbf{w},\mathbf{s}_{i}), i.e. there exists ρ>0\rho>0 such that (L(w;xi,xj)−si)++ρ2∥w∥2+ρ2si2(L(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j})-s_{i})_{+}+\frac{\rho}{2}\|\mathbf{w}\|^{2}+\frac{\rho}{2}s_{i}^{2} is jointly convex in terms of w,s\mathbf{w},\mathbf{s}. To this end, let ψ(⋅)=[⋅]+\psi(\cdot)=[\cdot]_{+} which is convex and Lipchitz continuous, and q(w,si)=L(w;xi,xj)−siq(\mathbf{w},\mathbf{s}_{i})=L(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j})-s_{i}, which is LsL_{s}-smooth function with respect to (w,si)(\mathbf{w},s_{i}) due to that L(w;xi,xj)L(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j}) is LsL_{s}-smooth function. Then for any ω∈ϕ′(ψ(w′,si′))\omega\in\phi^{\prime}(\psi(\mathbf{w}^{\prime},s^{\prime}_{i})) we have

where we use 0≤ω≤10\leq\omega\leq 1. The above inequality implies that [L(w;xi,xj)−si]+[L(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j})-s_{i}]_{+} is LsL_{s}-weakly convex in terms of (w,si)(\mathbf{w},s_{i}) (Davis & Drusvyatskiy 2018), i.e., 1n−∑xj∈S−{(L(w;xi,xj)−si)++Ls2(∥w∥2+∣si∣2)}\frac{1}{n_{-}}\sum_{\mathbf{x}_{j}\in\mathcal{S}_{-}}\left\{(L(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j})-s_{i})_{+}+\frac{L_{s}}{2}(\|\mathbf{w}\|^{2}+|s_{i}|^{2})\right\} is convex. As a result 1n−∑xj∈S−(L(w;xi,xj)−si)++Ls2(∥w∥2+∥si∥2)\frac{1}{n_{-}}\sum_{\mathbf{x}_{j}\in\mathcal{S}_{-}}(L(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j})-s_{i})_{+}+\frac{L_{s}}{2}(\|\mathbf{w}\|^{2}+\|s_{i}\|^{2}) is jointly convex in (w,si)(\mathbf{w},s_{i}). Then F(w,s)+Ls2β(∥w∥2+∑i∣si∣2)F(\mathbf{w},\mathbf{s})+\frac{L_{s}}{2\beta}(\|\mathbf{w}\|^{2}+\sum_{i}|s_{i}|^{2}) is jointly convex in terms of (w,s)(\mathbf{w},\mathbf{s}).

B.6 Proof of Theorem 3

Let F(w,s)F(\mathbf{w},\mathbf{s}) denote the objective function, and let v=(w,s)\mathbf{v}=(\mathbf{w},\mathbf{s}). Define F1/ρ^(v)=min⁡uF(u)+ρ^2∥u−v∥2F_{1/\hat{\rho}}(\mathbf{v})=\min_{\mathbf{u}}F(\mathbf{u})+\frac{\hat{\rho}}{2}\|\mathbf{u}-\mathbf{v}\|^{2} for some ρ^>ρ\hat{\rho}>\rho and the minimizer is denoted by proxF/ρ^(v)\text{prox}_{F/\hat{\rho}}(\mathbf{v}). Let v^t=proxF/ρ^(vt)\widehat{\mathbf{v}}_{t}=\text{prox}_{F/\hat{\rho}}(\mathbf{v}_{t}). Define ∥v−v′∥2=∥w−w′∥2+∥s−s′∥2\|\mathbf{v}-\mathbf{v}^{\prime}\|^{2}=\|\mathbf{w}-\mathbf{w}^{\prime}\|^{2}+\|\mathbf{s}-\mathbf{s}^{\prime}\|^{2}.

Summing the above inequality over i∈B+i\in\mathcal{B}_{+}, we have

Taking expectation and re-arrange, we have

Let ρ^=3ρ/2\hat{\rho}=3\rho/2. 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 min⁡w1n∑zi∈Dfi(gi(w))\min_{\mathbf{w}}\frac{1}{n}\sum_{\mathbf{z}_{i}\in\mathcal{D}}f_{i}(g_{i}(\mathbf{w})) and the convergence proof just follows the proof of Theorem 1 in Wang & Yang 2022.

B.8 Proof of Theorem 5

If gig_{i} is CgC_{g}-Lipschitz, LgL_{g}-smooth and f1,f2f_{1},f_{2} are CfC_{f}-Lipschitz, LgL_{g}-smooth, FF in (12) is LFL_{F}-smooth and LF=LfCf2Cg2+Cf2Lg+CfCgLfL_{F}=L_{f}C_{f}^{2}C_{g}^{2}+C_{f}^{2}L_{g}+C_{f}C_{g}L_{f}.

We can show that ∥1n∑i∈S(∇f2(gi(w))∇gi(w)−∇f2(gi(w′))∇gi(w′))∥≤(CfLg+CgLf)∥w−w′∥\left\|\frac{1}{n}\sum_{i\in\mathcal{S}}(\nabla f_{2}(g_{i}(\mathbf{w}))\nabla g_{i}(\mathbf{w})-\nabla f_{2}(g_{i}(\mathbf{w}^{\prime}))\nabla g_{i}(\mathbf{w}^{\prime}))\right\|\leq(C_{f}L_{g}+C_{g}L_{f})\left\|\mathbf{w}-\mathbf{w}^{\prime}\right\|. Thus,

Consider a sequence wt+1=wt−ηmt\mathbf{w}_{t+1}=\mathbf{w}_{t}-\eta\mathbf{m}_{t} and the LFL_{F}-smooth function FF and the step size ηLF≤1/2\eta L_{F}\leq 1/2.

where Δt≔mt−∇F(wt)\Delta_{t}\coloneqq\mathbf{m}_{t}-\nabla F(\mathbf{w}_{t})

For the gradient estimator mt\mathbf{m}_{t} in SOTA-s and Δt=mt−∇F(wt)\Delta_{t}=\mathbf{m}_{t}-\nabla F(\mathbf{w}_{t}),

where we denote Δt≔mt+1−∇F(wt)\Delta_{t}\coloneqq\mathbf{m}_{t+1}-\nabla F(\mathbf{w}_{t}), Ξt≔ut+1−g(wt)\Xi_{t}\coloneqq\mathbf{u}_{t+1}-\mathbf{g}(\mathbf{w}_{t}) and Ψt≔vt+1−1n∑i∈Sf2(gi(wt))\Psi_{t}\coloneqq v_{t+1}-\frac{1}{n}\sum_{i\in\mathcal{S}}f_{2}(g_{i}(\mathbf{w}_{t})).

Based on the update rule mt+1=(1−γ2)mt+γ2G(wt+1)\mathbf{m}_{t+1}=(1-\gamma_{2})\mathbf{m}_{t}+\gamma_{2}G(\mathbf{w}_{t+1}), we have

Based on the Young’s inequality for products, we have 2⟨a,b⟩≤∥a∥2c2+2∥b∥2c2\left\langle\mathbf{a},\mathbf{b}\right\rangle\leq\frac{\left\|\mathbf{a}\right\|^{2}c}{2}+\frac{2\left\|\mathbf{b}\right\|^{2}}{c} for c>0c>0.

where we denote Ξt≔ut−g(wt)\Xi_{t}\coloneqq\mathbf{u}_{t}-\mathbf{g}(\mathbf{w}_{t}) and Ψt≔vt−1n∑i∈Sf2(gi(wt))\Psi_{t}\coloneqq v_{t}-\frac{1}{n}\sum_{i\in\mathcal{S}}f_{2}(g_{i}(\mathbf{w}_{t})). ∎

For the function value estimator vt+1v_{t+1} in SOTA-s and Ψt≔vt−1n∑i∈Sf2(gi(wt))\Psi_{t}\coloneqq v_{t}-\frac{1}{n}\sum_{i\in\mathcal{S}}f_{2}(g_{i}(\mathbf{w}_{t})),

According to the update of vv in SOTA-s, we have

Denoting Ψt≔∥vt−1n∑i∈Sf2(gi(wt))∥2\Psi_{t}\coloneqq\left\|v_{t}-\frac{1}{n}\sum_{i\in\mathcal{S}}f_{2}(g_{i}(\mathbf{w}_{t}))\right\|^{2}, we have

Re-arranging the terms and telescoping (19) from t=1t=1 to TT 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 1/(γ0B+)≥max⁡(10Cf2/n,1/n)1/(\gamma_{0}B_{+})\geq\max(10C_{f}^{2}/n,1/n), 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 η≤min⁡{γ24LF,γ0B+40nCfCgC1Lf2+5Cf2,γ115Cf2CgC1Lf}\eta\leq\min\left\{\frac{\gamma_{2}}{4L_{F}},\frac{\gamma_{0}B_{+}}{40nC_{f}C_{g}C_{1}\sqrt{L_{f}^{2}+5C_{f}^{2}}},\frac{\gamma_{1}}{15C_{f}^{2}C_{g}C_{1}L_{f}}\right\}, we have

Besides, Lemma 2 in Wang & Yang 2022 and Lemma 11 imply that

Set γ0=B−ϵ2400Cf2C12Lf2σ2(1+5Cf2)\gamma_{0}=\frac{B_{-}\epsilon^{2}}{400C_{f}^{2}C_{1}^{2}L_{f}^{2}\sigma^{2}(1+5C_{f}^{2})}, γ1=B+ϵ2200Cf4C12Lf2\gamma_{1}=\frac{B_{+}\epsilon^{2}}{200C_{f}^{4}C_{1}^{2}L_{f}^{2}}, γ2=min⁡{B−,B+}ϵ220Cf4(ζ2+Cg2)\gamma_{2}=\frac{\min\left\{B_{-},B_{+}\right\}\epsilon^{2}}{20C_{f}^{4}(\zeta^{2}+C_{g}^{2})}, 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 (min⁡xf(x)−s)+=min⁡x(f(x)−s)+(\min_{x}f(x)-s)_{+}=\min_{x}(f(x)-s)_{+}. Using the conjugate of [⋅]+[\cdot]_{+}, we have

Define Fi(w,si,π,ui)=π+1αui(si+1βψi(w;si)−π)F_{i}(\mathbf{w},s_{i},\pi,u_{i})=\pi+\frac{1}{\alpha}u_{i}(s_{i}+\frac{1}{\beta}\psi_{i}(\mathbf{w};s_{i})-\pi) such that F(w,s,π,u)=1n+∑xi∈S+Fi(w,si,π,ui)F(\mathbf{w},\mathbf{s},\pi,\mathbf{u})=\frac{1}{n_{+}}\sum_{\mathbf{x}_{i}\in\mathcal{S}_{+}}F_{i}(\mathbf{w},s_{i},\pi,u_{i}). Based on a minibatch B−⊆S−\mathcal{B}_{-}\subseteq\mathcal{S}_{-}, we can estimate Fi(w,si,π,ui)F_{i}(\mathbf{w},s_{i},\pi,u_{i}) by Fi(w,si,π,ui;B−)≔π+1αui(si+1βψi(w;si;B−)−π)F_{i}(\mathbf{w},s_{i},\pi,u_{i};\mathcal{B}_{-})\coloneqq\pi+\frac{1}{\alpha}u_{i}(s_{i}+\frac{1}{\beta}\psi_{i}(\mathbf{w};s_{i};\mathcal{B}_{-})-\pi) and ψi(w;si;B−)≔1B−∑xj∈B−(L(w;xi,xj)−si)+\psi_{i}(\mathbf{w};s_{i};\mathcal{B}_{-})\coloneqq\frac{1}{B_{-}}\sum_{\mathbf{x}_{j}\in\mathcal{B}_{-}}(L(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j})-s_{i})_{+}. We consider the function F(w,s,π,u)F(\mathbf{w},\mathbf{s},\pi,\mathbf{u}), which can be proved to be weakly convex w.r.t. (w,s,π)(\mathbf{w},\mathbf{s},\pi) and concave w.r.t. u\mathbf{u}. Hence, we can use the stagewise proximal point method to solve the problem (Rafique et al. 2020). At the kk-th stage, we solve the following problem approximately:

Assume max⁡(∣πk,t∣,∣sik,t∣,L(wk,t;xi,xj),∥∇L(wk,t;xi,xj)∥)≤C\max(|\pi^{k,t}|,|s_{i}^{k,t}|,L(\mathbf{w}^{k,t};\mathbf{x}_{i},\mathbf{x}_{j}),\|\nabla L(\mathbf{w}^{k,t};\mathbf{x}_{i},\mathbf{x}_{j})\|)\leq C. Then, we have

Let v≔(w,s,π)\mathbf{v}\coloneqq(\mathbf{w},\mathbf{s},\pi). We first show that F(v,u)F(\mathbf{v},\mathbf{u}) is weakly convex in terms of v\mathbf{v} for any u\mathbf{u}.

Under Assumption 2, then F(v,u)F(\mathbf{v},\mathbf{u}) is ρ/(αβ)\rho/(\alpha\beta)-weakly convex in terms of v\mathbf{v} for any u\mathbf{u}, where ρ\rho is the smoothness constant of L(w;xi,xj)L(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j}) w.r.t. w\mathbf{w}.

Following similar analysis of Lemma 2, we can show that F(v,u)+ρ2αβ∥v∥2=F(v,u)+1n+α∑iui(ρ2β∥w∥2+ρ2β∣si∣2+ρ2β∣π∣2)+ρ2n+αβ∑i(1−ui)(∥w∥2+∣si∣2+∣π∣2)+(n+−1)2n+αβ∥s∥2F(\mathbf{v},\mathbf{u})+\frac{\rho}{2\alpha\beta}\|\mathbf{v}\|^{2}=F(\mathbf{v},\mathbf{u})+\frac{1}{n_{+}\alpha}\sum_{i}u_{i}(\frac{\rho}{2\beta}\|\mathbf{w}\|^{2}+\frac{\rho}{2\beta}|s_{i}|^{2}+\frac{\rho}{2\beta}|\pi|^{2})+\frac{\rho}{2n_{+}\alpha\beta}\sum_{i}(1-u_{i})(\|\mathbf{w}\|^{2}+|s_{i}|^{2}+|\pi|^{2})+\frac{(n_{+}-1)}{2n_{+}\alpha\beta}\|\mathbf{s}\|^{2} is jointly convex in terms w,s,π\mathbf{w},\mathbf{s},\pi for any u∈\mathbf{u}\in. Then F(v,u)F(\mathbf{v},\mathbf{u}) is ρ′=ραβ\rho^{\prime}=\frac{\rho}{\alpha\beta}-weakly convex in terms of v\mathbf{v} for any u\mathbf{u}. ∎

Due to the update of xt+1\mathbf{x}_{t+1}, we have

Assume there exists C>0C>0 such that max⁡(∣st′∣,∣st,i∣,L(wt;xi,xj),∥∇L(wt;xi,xj)∥)≤C\max(|s^{\prime}_{t}|,|s_{t,i}|,L(\mathbf{w}_{t};\mathbf{x}_{i},\mathbf{x}_{j}),\|\nabla L(\mathbf{w}_{t};\mathbf{x}_{i},\mathbf{x}_{j})\|)\leq C at every stage. Let 1/γ≥ρ,η1k=η2k=η3k=η4k=∝1/k,Tk∝n+k21/\gamma\geq\rho,\eta^{k}_{1}=\eta^{k}_{2}=\eta^{k}_{3}=\eta^{k}_{4}=\propto 1/k,T_{k}\propto n_{+}k^{2}. SOTA ensures that after T=O(n+/ϵ6)T=O(n_{+}/\epsilon^{6}) iterations we can find an ϵ\epsilon-nearly stationary solution for min⁡w,s,s′F(w,s,s′)\min_{\mathbf{w},\mathbf{s},s^{\prime}}F(\mathbf{w},\mathbf{s},s^{\prime}).

Let R1(w)≔12γ∥w−wk,0∥2R_{1}(\mathbf{w})\coloneqq\frac{1}{2\gamma}\|\mathbf{w}-\mathbf{w}^{k,0}\|^{2}, R2(s)≔12γ∥s−sk,0∥2R_{2}(\mathbf{s})\coloneqq\frac{1}{2\gamma}\|\mathbf{s}-\mathbf{s}^{k,0}\|^{2}, and R3(π)=12γ∣π−πk,0∣2R_{3}(\pi)=\frac{1}{2\gamma}|\pi-\pi^{k,0}|^{2}. Apply Lemma 13 to wk,t+1\mathbf{w}^{k,t+1}.

Take expectation on both sides conditioned on the randomness that occurred before the tt-th iteration in the kk-th stage.

Similarly, apply Lemma 13 to sik,t+1\mathbf{s}_{i}^{k,t+1} and πk,t+1\pi^{k,t+1} and take the conditional expectations.

Note that Fk(vk,t,uk,t)−Fk(v,uk,t)=F(vk,t,uk,t)−F(v,uk,t)+12γ∥vk,t−vk,0∥22−12γ∥v−vk,0∥22F_{k}(\mathbf{v}^{k,t},\mathbf{u}^{k,t})-F_{k}(\mathbf{v},\mathbf{u}^{k,t})=F(\mathbf{v}^{k,t},\mathbf{u}^{k,t})-F(\mathbf{v},\mathbf{u}^{k,t})+\frac{1}{2\gamma}\|\mathbf{v}^{k,t}-\mathbf{v}^{k,0}\|_{2}^{2}-\frac{1}{2\gamma}\|\mathbf{v}-\mathbf{v}^{k,0}\|_{2}^{2}. If 1γ≥ρ′=ραβ\frac{1}{\gamma}\geq\rho^{\prime}=\frac{\rho}{\alpha\beta}, we have Fk(v,u)F_{k}(\mathbf{v},\mathbf{u}) is convex w.r.t. v\mathbf{v} such that Fk(vk,t,uk,t)−Fk(v,uk,t)≤(vk,t−v)⊤∇vF(vk,t,uk,t)+(vk,t−v)⊤∇vR(vk,t,uk,t)F_{k}(\mathbf{v}^{k,t},\mathbf{u}^{k,t})-F_{k}(\mathbf{v},\mathbf{u}^{k,t})\leq(\mathbf{v}^{k,t}-\mathbf{v})^{\top}\nabla_{\mathbf{v}}F(\mathbf{v}^{k,t},\mathbf{u}^{k,t})+(\mathbf{v}^{k,t}-\mathbf{v})^{\top}\nabla_{\mathbf{v}}R(\mathbf{v}^{k,t},\mathbf{u}^{k,t}). where R=(R1,R2,R3)⊤R=(R_{1},R_{2},R_{3})^{\top}. Note that (vk,t−v)⊤∇vR(vk,t,uk,t)=1γ(vk,t−v)⊤(vk,t−vk,0)≥R(vk,t)−R(v)(\mathbf{v}^{k,t}-\mathbf{v})^{\top}\nabla_{\mathbf{v}}R(\mathbf{v}^{k,t},\mathbf{u}^{k,t})=\frac{1}{\gamma}(\mathbf{v}^{k,t}-\mathbf{v})^{\top}(\mathbf{v}^{k,t}-\mathbf{v}^{k,0})\geq R(\mathbf{v}^{k,t})-R(\mathbf{v}). Adding the above inequalities for w,s,π\mathbf{w},\mathbf{s},\pi together we have

where C12≔Cα2β2C_{1}^{2}\coloneqq\frac{C}{\alpha^{2}\beta^{2}}, C22=1α2(1+1/β)2C_{2}^{2}=\frac{1}{\alpha^{2}}(1+1/\beta)^{2}, C32≔(1+1/α)2C_{3}^{2}\coloneqq(1+1/\alpha)^{2}. Applying the same analysis to the update of uk,t\mathbf{u}^{k,t}, we have

We do not assume u\mathbf{u} is independent of the randomness in the updates of our algorithm. As a result,

According to the initialization in Algorithm 4, for any v=(w,s,π)\mathbf{v}=(\mathbf{w},\mathbf{s},\pi) and u\mathbf{u} we have

Note that ∥s∥2≤n+C2\|\mathbf{s}\|^{2}\leq n_{+}C^{2} and ∥u−1∥2≤n+\|\mathbf{u}-1\|^{2}\leq n_{+}. It remains to apply the analysis in (Rafique et al. 2020, Theorem 4.1) to derive the convergence for the Moreau envelope of F(w,s,π)F(\mathbf{w},\mathbf{s},\pi) with a complexity in the order of O(n+/ϵ6)O(n_{+}/\epsilon^{6}) by setting ηk∝1/k\eta_{k}\propto 1/k and Tk∝n+k2T_{k}\propto n_{+}k^{2} and the total number of stages K=O(1/ϵ2)K=O(1/\epsilon^{2}). ∎