Multi-block Min-max Bilevel Optimization with Applications in Multi-task Deep AUC Maximization
Anonymous Author(s)
AffiliationAddress email
Abstract
In this paper, we study multi-block min-max bilevel optimization problems, wherethe upper level is non-convex strongly-concave minimax objective and the lowerlevel is a strongly convex objective, and there are multiple blocks of dual variablesand lower level problems.
Due to the intertwined multi-block min-max bilevelstructure, the computational cost at each iteration could be prohibitively high,especially with a large number of blocks.
To tackle this challenge, we present asingle-loop randomized stochastic algorithm, which requires updates for only aconstant number of blocks at each iteration.
Under some mild assumptions on theproblem, we establish its sample complexity of O (1 /ϵ 4 ) for finding an ϵ -stationarypoint.
This matches the optimal complexity for solving stochastic nonconvexoptimization under a general unbiased stochastic oracle model.
Moreover, weprovide two applications of the proposed method in multi-task deep AUC (areaunder ROC curve) maximization and multi-task deep partial AUC maximization.
Experimental results validate our theory and demonstrate the effectiveness of ourmethod on problems with hundreds of tasks.
Introduction
We consider multi-block min-max bilevel optimization problem of the following formulationmin
[Equation]
(upper)
[Equation]
for
[Equation] [Equation]
m
[Equation]
(lower) (1) where f i and g
i are smooth functions and
A ⊂ R d α is a convex set.
In particular, in this paperwe assume that for each i ∈ { 1 , . . .
, m } , f i ( x , α i , y i ) is strongly concave in the dual variable α ibut can be nonconvex in the primal variable x , and g i ( x , y i ) is strongly convex in y i .
The upperproblem min x ∈ R dx max α ∈A m F ( x , α ) is a min-max optimization problem where the lower problems{ y i ( x ) =
arg min
y i ∈
R dy g
i ( x , y i ) } mi =1 are involved as variables.
For each block i , the upper-levelobjective f i ( x , α i , y i ) and the lower-level objective g i ( x , y i ) depend only on its corresponding blockof variables α and y , i.e. α i and y i .
This problem has important applications in machine learning,e.g., multi-task deep AUC maximization as presented in section 3.
Tackling problem (1) is challenging as it involves solving a min-max problem with coupled multipleminimization problems simultaneously.
The naive way for solving it is to do multiple gradient ascentsand descents for α and y , respectively, to ensure a good estimation of the gradient for updating x .
However, this approach has two major drawbacks.
As the algorithm involves a double loop structure,it can be computationally expensive and give suboptimal theoretical complexity.
On the other hand,the multi-block structure requires data sampling from distributions
P i , Q i for all blocks, which maylead to an impractical demand for memory.
1.1 Related work
Min-max Bilevel optimization.
To the best of our knowledge, the only existing work that provides astochastic algorithm with provable convergence guarantee on min-max bilevel problems is [8].
Theypropose a single loop bi-time scale stochastic algorithm based on gradient descent ascent, and provethat it converges to an ϵ -stationary point with an oracle complexity of O ( ϵ − 5 ) .
Nevertheless, thisconvergence result is established for a special case where f ( x , · , y ) is a linear function.
Stochastic Nonconvex Strongly Concave Min-max Problems.
The considered problem is alsoclosely related to non-convex strongly concave min-max problems, which have been studied extensively recently.
To the best of our knowledge, [28] establishes the first result on non-smoothnonconvex concave min-max problems.
They prove a convergence to nearly stationary point of theprimal objective function with an oracle complexity in the order of O ( ϵ − 4 ) for non-convex stronglyconcave min-max problems with a certain special structure.
The same order of oracle complexity isachieved in [34] without relying on any special structure.
These two works use two-loop algorithms.
There are some studies focusing on single-loop algorithms.
[23] analyzes a single-loop stochasticgradient descent ascent (SGDA) method for smooth nonconvex strongly concave problem, whichachieves O ( ϵ − 4 ) complexity but with a large mini-batch size.
In [11], the same order of complexityis achieved without large mini-batch by employing the stochastic moving average estimator.
Somerecent works improve the sample complexity to O ( ϵ − 3 ) under the Lipschitz continuous oracle modelfor the stochastic gradient using less practical variance reduction techniques [15, 26, 29].
[39]establishes lower complexity bounds for non-concex strongly concave min-max problem underboth general and finite-sum setting and proposes accelerated algorithms that nearly match the lowerbounds.
Stochastic Nonconvex Bilevel optimization.
The considered problem belongs to a general familyof non-convex bilevel optimization problems.
Non-asymptotic convergence results for nonconvexstochastic bilevel optimization (SBO) with a strongly convex lower problem has been establishedin several recent studies [4, 7, 11, 13, 18].
As the one who gives the first results for this problem,[7] proposes a double-loop algorithm with O ( ϵ − 6 ) oracle complexity for finding an ϵ -stationarypoint of the objective function.
[18] improves the complexity order to O
( ϵ − 4 ) , but suffers froma large mini-batch size.
[13] proposes a single-loop algorithm with two time-scale updates thatachieves an oracle complexity of (cid:101)
O
( ϵ − 5 ) .
Recently, [11] improves the oracle complexity to thestate-of-the-art oracle complexity (cid:101) O
( ϵ − 4 ) by proposing a single-loop algorithm based on a movingaverage estimator.
[5] presents a new analysis for (double-loop) SGD-type updates showing thatan improved sample complexity O
( ϵ − 4 ) can be achieved.
There are studies that further improvethe complexity to O
( ϵ − 3 ) by leveraging the Lipschitz continuous conditions of stochastic oracles[4, 9, 20].
[22] considers bilevel optimization under distributed setting and proposed algorithmsachieving state-of-the-art complexities.
However, none of these works tackle multi-block min-maxbilevel optimization problems directly.
Multi-block Bilevel Optimization.
There are some recent studies considering bilevel optimizationwith multi-block structure.
[9] extends their single-block bilevel optimization algorithm to multi-blockstructure.
Their algorithm requires two independently sampled block batches and for all sampledblocks, each variable needs update using variance reduction technique STORM [6].
For unsampledblocks, an update involving constant factor multiplication is also required.
Under Lipschitz continuousconditions on stochastic oracles, the complexity is no worse than O ( m/ϵ 3 ) with m blocks.
A morerecent work [27] considers top-K NDCG optimization, which is formulated as a compositional bileveloptimization with multi-block structure.
Their method simplifies the updates by sampling only oneblock batch in each iteration and requires updates only for the sampled blocks.
Their method achievescomplexity of O ( m/ϵ 4 ) .
We use a similar approach as the latter work for estimating the hessianinverse in a block-wise manner.
However, this paper differs from [27] in that we tackle a more generalmulti-block min-max bilevel problems without assuming a particular form of the objective.
1.2 Our Contributions
We present a simple single loop single timescale stochastic method with randomized block-samplingfor solving a general form of multi-block min-max bilevel optimization problem under the nonconvexstrongly concave (upper) strongly convex (lower) setting.
It employs SGD for updating selected y i fortheir corresponding lower-level problems, employs SGA for updating the selected α i , and employs amomentum update for the primal variable x based on the sampled α
i , y i .
Then we show, theoretically,that it converges to ϵ -stationary point with complexity O
( ϵ − 4 ) under a general unbiased stochasticoracle model.
This oracle complexity matches the optimal complexity order for solving stochastic nonconvex optimization under a general unbiased stochastic oracle model [1].
Finally, we presenttwo applications of multi-block min-max bilevel optimization in deep AUC maximization: multi-taskdeep AUC maximization and multi-task deep partial AUC maximization.
Empirical results show theeffectiveness of the proposed method.
Preliminaries
Notations.
Let ∥ · ∥ denote the Euclidean norm of a vector or the spectral norm of a matrix.
For a twice differentiable function f : X × Y → R , ∇ x f ( x, y ) (resp.
∇ y f
( x, y ) ) denotes itspartial gradient taken w.r.t x (resp.
y ), and ∇ 2 xy f ( x, y ) (resp.
∇ 2 yy f ( x, y ) ) denotes the Jacobianof ∇ x f ( x, y ) w.r.t x (resp.
∇ y f
( x, y ) w.r.t y).
We let f ( · ; B ) represent the unbiased stochasticoracle of f ( · ) with a sample batch B as the input.
The unbiased stochastic oracle is said to havebounded variance σ 2 if E [ ∥ f ( · ; B ) − f ( · ) ∥ 2 ] ≤ σ 2 .
A mapping f : X → R is C -Lipschitzcontinuous
if ∥ f ( x ) − f ( x ′ )
∥ ≤ C ∥ x − x ′ ∥
∀ x, x ′ ∈ X .
Function f is L -smooth if its gradient∇ f ( · ) is L -Lipschitz continuous.
A function g : X → R is λ -strongly convex
if ∀ x, x ′ ∈ X ,g ( x ) ≥ g ( x ′ )
+ ∇ g ( x ′ ) T ( x − x ′ )
+ λ
2 ∥
x
− x ′ ∥ 2 .
A function g : X → R is λ -strongly
concave
if− g ( x ) is λ -strongly convex.
Let Π A denote a projection function onto a convex set A .
For notationsimplicity, we use S to denote the set of all block indices, i.e. S = { 1 , . . .
, m } {+.
+}
We state the definition of ϵ -stationary point as following.
Definition 2.1.
Consider a differentiable function F ( x ) , a point x is called ϵ -stationary if ∥∇ F ( x ) ∥ ≤ϵ .
A stochastic algorithm is said to achieve an ϵ -stationary point if E [ ∥∇ F (¯ x t ) ∥ ] ≤ ϵ , where ¯ x t isthe algorithm output at the t -th iteration and the expectation is taken over the randomness of thealgorithm until the iteration t .
Assumptions.
Before presenting our algorithm, we make the following well-behaving assumptions.
Assumption 2.2.
For functions f i and g i , we assume that the following conditions hold for all i ∈ S
• f i ( x , α i , y i ) is µ f -strongly concave in terms of α i .
g i ( x, y i ) is µ g -strongly convex in terms
of y i .
• f
i , g i are C f , C g -Lipschitz continuous and L f , L g -smooth respectively.
• ∥∇ 2 xy g i ( x , y i )
∥ 2 ≤ C 2 gxy , ∇
yy g i ( x , y i ; ζ ) ⪰ µ g I .
• ∇ 2 xy g i ( x , y i ) , ∇ 2 yy g
i
( x , y i ) are L gxy , L gyy -Lipschitz continuous respectively.
Moreover, the gradients of functions f i and g i can only be accessed through unbiased oracles withbounded variance.
Assumption 2.3.
The unbiased stochastic oracles ∇ x f i ( x , α i , y i ; B ) , ∇ α f i ( x , α i , y i ; B ) ,∇ y f i ( x , α i , y i ; B ) , ∇ y g
i ( x , y i ; B ) , ∇ 2 xy g i ( x , y i ; B ) , ∇ 2 yy g
i ( x , y i ; B ) have variances bounded by σ 2 |B|for all i ∈ S , where |B| denotes the size of the sampled batch B .
These assumptions are similar to those made in many existing works for SBO [4, 7, 13, 18].
Moving average gradient estimator.
Algorithms based on moving average estimators have achievedthe state-of-the-art oracle complexity in both min-max and bilevel optimizations [10].
Here wegive a brief introduction to the moving average estimator.
For solving a nonconvex minimizationproblem min x ∈ R d F ( x ) through an unbiased oracle O F ( x ) , i.e. E [ O F ( x )]
= ∇ F ( x ) , the stochasticmomentum method (stochastic heavy-ball method) that employs moving average updates is given by
[Equation]
where β and η are momentum parameter and learning rate, respectively.
As a moving average of thehistorical gradient estimator, the sequence of v t +1 could achieve an effect of variance diminishingacross a long run ([30]).
2.1 The Proposed Algorithm
In this section, we propose a simple single loop stochastic algorithm 1 to solve the multi-blockmin-max bilevel optimization problem.
At the beginning of each iteration, we first sample a set of
Require: α 0 , y 0 , H 0 , z 0 , x1: for t = 0 , 1 , . . .
, T do 2:
Sample tasks I t ∈ S .
Sample data batch B ti for each i ∈ I t .
4: α t +1 i = (cid:26)
[Equation]
f
[Equation]
B
[Equation] [Equation]
α
[Equation]
o.w.
5: y t +
[Equation]
= (cid:26)
[Equation]
y
[Equation]
o.w.
Update estimator H t +1 of [ ∇
yy g
i ( x t , y ti )] − 1 by (2) 7:
Update gradient estimator ∆ t +1 of ∇ x F ( x t ) by (3) 8: z t +1 =
(1 − β 0 )
z t + β 0 ∆ t + 9: x t +1 = x
t − η 0 z t + 10: end for blocks I t and a data batch B ti for each selected block i ∈ I t .
Then we update estimators of α i ( x t )and y i ( x t ) for all selected blocks i ∈
I t using one step of SGA and SGD.
Then, we compute anestimator of the hessian inverse H t +1 of the lower-level objective and compute a gradient estimator∆ t +1 of the upper-level objective.
Finally, we compute the moving average estimator z t +1 of ∇ F ( x t )and update x t +1 .
Note that the design of our algorithm on the min-max bilevel optimization part isinspired by [10], and it is similar to their momentum-based algorithms PDSM (for min-max problem)and SMB (for bilevel problem) in their paper.
In fact, if we set the number of blocks to be one andremove y and the lower level problems, then Algorithm 1 is the same as PDSM.
Similarly, if thenumber of blocks is one and the dual variables in the upper level problem are removed, then theproposed algorithm becomes similar as SMB except for the hessian inverse update with reasonsexplain shortly.
In other words, our proposed method is a generalized form of momentum-basedalgorithm for min-max and bilevel optimization problems.
Additionally, Algorithm 1 only updates
O (1) blocks of dual variables α i and the variables
y i of the lower-level problems.
These make theanalysis of Algorithm 1 much more involved.
To further understand Algorithm 1, we first define the objective function
[Equation]
where y i ( x ) =
arg min y
i g i ( x , y i ) and α i ( x ) :
= arg max α i f i ( x , α i , y i ( x )) , so that the Problem (1)can be rewritten as min x F ( x ) .
The updates for α t +1 i ’s and y t
+1 i ’s are intuitive since the gradientestimations of ∇ α f i ( x t , α ti , y ti ) and ∇ y g
i ( x t , y ti ) are directly available from the unbiased stochasticoracles.
However, since functions y i ( x ) and α i ( x ) are implicit, estimating the gradient ∇ F ( x ) isdifficult.
In fact, one may apply the corollary of Theorem 1 in [3] to get:
[Equation]
A standard approach in bilevel optimization literature [7] for computing ∇ y
i ( x ) is to derive ∇ y
i ( x ) =−∇
2 xy g i
( x , y i ( x ))
[ ∇ 2 yy g
i
( x , y i ( x ))] − 1 from the optimality condition of y i ( x ) .
Therefore, the gradientwe are looking for is given by
[Equation]
All components in this gradient can be easily obtained from unbiased stochastic oracles except forthe inverse of hessian [ ∇ 2 yy g
i ( x , y i ( x ))] − 1 for all blocks.
This could be problematic in the sense oftheory and practical implementation.
For practical implementation, we do not want to update thehessian inverse estimators for all blocks, which is prohibitive when the number of blocks is large.
Acommon approach used in the literature of SBO is to use Neumann series [7]:
Choose q ∈ { 1 , . .
k t } randomly,
[Equation]
where k t is the number of samples { ξ ti } k t i =1 for estimating the hessian inverse.
This estimator isa biased one and its error w.r.t to ∇
yy g i ( x t , y ti ) − 1 is controlled by the number of samples k t [7].
However, it is problematic to employ the above estimator for only the sampled blocks i ∈
I t becausethe error for those not sampled cannot be controlled.
To address this issue, we use a different approach for estimating the hessian inverse by only updatingthe estimators for those sampled blocks [27].
The idea is to maintain a momentum term s t
+1 i for eachblock that stores historical information on the hessian estimator ∇
yy g i ( x t , y ti ; B ti ) .
And the hessianinverse is approximated by directly computing the inverse of s t
+1 i , i.e.,
[Equation]
o
[Equation]
w
[Equation]
In terms of theoretical analysis, we are not bounding the individual error H t
+1 i − ∇
yy g i ( x
[Equation]
, y ti ) − 1 forall blocks, but the cumulative error for all blocks across all iterations.
This is exhibited in Lemma 2.6.
As the conclusion of the above discussion, the gradient estimator of ∇ F ( w t ) is given by
[Equation]
We maintain a moving average estimator z t +1 for ∆ t +1 and finally update x t +1 using z t +1 .
2.2 Convergence Analysis
In this section, we present a brief convergence analysis of Algorithm 1.
The detailed theorems andproofs are deferred to the appendix.
The key point of this analysis is the gap between the true gradient∇ F ( x t ) and its estimator z t +1 .
To this end, we define
[Equation]
One may notice that the estimator ∆ t +1 is in fact approximating ∇ x F ( x t , α t , y t ) instead of ∇ x F ( x t ) .
We exploit the moving average formulation of z t +1 and decompose the gap into two parts, ∥ ∆ t +1 −∇ x F ( x t , α t , y t ) ∥ 2 and ∥∇ x F ( x t , α t , y t ) − ∇ x F ( x t ) ∥ 2 .
These two gaps are determined by howwell y t , α t and H ti approximate y
( x t ) , α ( x t ) and [ ∇
yy g i ( x t , y ti )] − 1 , respectively.
In other words,we aim to bound the following three errors, ∥ y ( x t ) − y t ∥
=: δ y,t , ∥ α ( x t ) − α t ∥ 2 :=
δ α,t and∥ [ ∇ 2 yy g
i ( x t , y ti )] − 1 − H ti ∥ 2 .
We first bound the variance E [ δ y,t ] by proving the following lemma.
Lemma 2.4. Consider the updates for y t in Algorithm 1, under Assumption 2.2 and 2.3, withη 2 ≤ min { µ g L 2 g , 2 m | I t | µ g } we have
[Equation]
∥
[Equation]
One may also bound the second variance E [ δ α,t ] based on the previous variance E [ δ y,t ] following asimilar strategy.
Lemma 2.5.
Consider the updates for α t in Algorithm 1, under Assumption 2.2, 2.3, with η 1 ≤min (cid:110) µ f L 2 f , 1 µ f , 4 m µ f | I t | (cid:111) , we have
T
[Equation]
Due to the lower bound assumption of ∇
yy g i ( x , y i ; ζ ) in Assumption 2.2, the error in Hessianapproximation can be bounded by bounding δ gyy,t := (cid:80) i ∈S ∥ s ti − ∇
yy g i ( x t , y i ( x t ))
∥ 2 .
We provethe following lemma.
Lemma 2.6.
Under Assumption 2.2,2.3, considering the momentum method (2) for the update ofhessian, with β 1 ≤ 1 we have
T (cid:88)
[Equation]
It then follows the convergence theorem for Algorithm 1.
Theorem 2.7.
Under Assumption 2.2, 2.3 and with a proper settings of parameters η 1 , η 2 , β 1 =
O (cid:0) |B
ti | ϵ 2 (cid:1) , β
= O (cid:0) min {| I t | , |B ti |} ϵ 2 (cid:1) and η 0 = O (cid:16) min (cid:110) min {| I t | , |B ti |} ϵ 2 , |B
ti
|| I t | ϵ 2 m (cid:111)(cid:17) , Algorithm 1 ensures
that after T = O (cid:16) max (cid:110) m | I t ||B ti | ϵ 4 , [Table] min {| [Table]
| , |B ti | [Table] (cid:111)(cid:17) iterations we can find anϵ -stationary solution of F ( x ) , i.e., E [ ∥∇ F ( x τ ) ∥ 2 ] ≤ ϵ 2 for a randomly selected τ ∈ { 0 , . . .
, T } .
Remark.
In Theorem 2.7, there is no condition on the sizes of data batch B ti nor block batch
I tfor the algorithm to converge.
Hence, their sizes can be as small as one.
The order of complexityis O (1 /ϵ 4 ) , which matches the optimal complexity for nonconvex optimization under a generalunbiased stochastic oracle model.
In addition, there is parallel speed up by increasing batch sizes fordata samples and task samples due to the scaling in terms of | I t | and |B ti | in the iteration complexity.
Applications in Multi-task Deep (Partial) AUC Maximization
In this section, we present two applications of multi-block min-max bilevel optimization: multi-taskdeep AUC maximization and deep partial AUC maximization.
(partial) AUC is a performance measureof classifiers for imbalanced data.
Recent studies have shown great success of deep AUC maximizationin various domains (e.g., medical image classification and molecular property prediction) [24, 31, 38].
However, efficient algorithms for multi-task deep (partial) AUC maximization have not been welldeveloped.
For multi-task deep AUC maximization, we solve an existing formulation by our algorithm.
For multi-task deep partial AUC maximization, we propose a new bilevel formulation and solve it byour algorithm.
3.1 Muti-task Deep AUC maximization
Following the previous work [24, 38], deep AUC maximization problem can be formulated as anon-convex strongly concave min-max optimization problem min w ,a,b max α ∈A L AUC ( w , a, b, α ) .
However, training a deep neural network from scratch by optimizing AUC loss does not necessarilylead to a good performance[37].
To address this issue, [37] proposed a compositional training strategyfor deep AUC maximization:
[Equation]
where L CE denotes the cross-entropy loss.
The outer objective remains to be the AUC loss, whilethe inner objective is a gradient descent step of minimizing the traditional cross-entropy loss.
Thismethod has shown superior performance on various datasets [37].
We extend this formulation tomulti-task problems and reformulate it into a multi-block min-max bilevel optimization:
[Equation]
where w l denotes the weight for the encoder network that is shared for all tasks, and w h =( w 1 h , . .
. , w mh ) denote the task-owned classification heads.
The upper objective is strongly concave in terms of dual variables α i and the lower level objective is strongly convex in terms of u i .
The hessian of the lower-level objective is the identity matrix.
Hence, there is no need to track andestimate the hessian matrix.
3.2 Multi-task Deep Partial AUC Maximization
Some real-world applications (e.g., medical diagnosis [2]) cannot tolerate a model with a high False
Positive Rate (FPR) even though it has significant performance in AUC.
Hence, a measure of interestis one-way partial AUC (pAUC), which puts a restriction on the range of FPR (i.e., FPR ∈ [ ρ l , ρ ] ,where 0 ≤ ρ l ≤ ρ ≤ 1 ).
Below, we focus on the case ρ l = 0 .
However, our method can be easilyextended for handling ρ l > 0 .
Let D
+ , D − denote the set of positive and negative data for a binaryclassification task, respectively.
Let D −
[ K ] denote the top-K negative examples according to theirprediction scores.
Let n + , n − denote the number of positive and negative samples respectively.
Then
Require: α 0 , λ 0 , H 0 , z 0 , w 0 , a 0 , b1: for t = 0 , 1 , . . .
, T do 2:
Draw task batch I t ⊂ S .
Draw data sample batch B tk for each k ∈
I t 4:
For sampled tasks k ∈ I t , update 5: α t +1 k ← α tk + η 1 G α ( w t , α tk , λ tk ; B tk ) ⋄ G α ( · ) denotes a stochastic gradient w.r.t α k 6: λ t +1 k ← λ tk − η 2 ∇ λ L k
( λ tk , w t ; B tk ) 7:
H t +1 k ← (1 − β 1 ) H tk + β 1 ∇ 2 λλ L k ( λ tk , w t ; B tk ) 8:
Compute loss G t +1 according to (25) ⋄ G t +1 denotes an appropriate loss 9:
Update gradient estimator ∆ t +1 ← autograd ( G t +1 ) 10: z t +1 ← (1 − β 0 )
z t + β 0 ∆ t + 11: ( w t +1 , a t +1 , b t +1 )
← ( w t , a t , b t ) − η
z t + 12: end for we have partial AUC optimization with a pairwise square loss formulated as following [35]:min
[Equation]
x
[Equation] [Equation]
[Equation]
where K = n − ρ , c is a constant and h w ( · ) denotes the prediction score on a data.
A key challenge forsolving the above problem is to deal with the non-differentiable top-K selector x j ∈ D − [ K ] , whichdepends on the model parameters w .
This challenge has been recently tackled in [36, 40].
We focuson the comparison with the first work as it is optimization oriented similar to ours and also has thestate-of-the-art performance.
They formulate the problem into either a weakly convex minimizationor approximate it by a smooth objective in a compositional form.
A caveat of their algorithms (named
SOPA, SOPA-s) is that they need to maintain and update n + auxiliary variables with one for eachpositive data.
If we apply their algorithms for multi-task problems, one needs to maintain and update 249 (cid:80) mi =1 n i + auxiliary variables, which could dramatically slow down the convergence.
To address this problem, we first transform it into a min-max optimization problem.
Let a ( w ) =
1 n + (cid:80) x i ∈D + h w ( x i ) and b ( w ) =
[Equation]
∈D
[Equation]
[ K ] h w ( x j ) .
Then we can write the problem as (cf.
Appendix B for a derivation)
[Equation]
To tackle non-continuous non-differentiable indicator function I ( · ) we can replace it by a sigmodfunction.
To tackle non-differentiability of the top-K selector x j ∈ D − [ K ] , we follow [27] andformulate it as lower-level optimization problem, i.e., x j ∈ D − [ K ] is equivalent to h w ( x j ) > λ ( w ) ,where λ ( w ) represents the K + 1 -th largest scores among all negative examples, which can beapproximated by a solution from a smooth strongly convex minimization problem as following:
[Equation]
where ε, τ 1 , τ 2 are small constants.
Based on above, the multi-task deep partial AUC minimizationproblem can be formulated as a multi-block min-max bilevel optimization problem give by:min w , a , b
[Equation]
[Table]
where D k + / − denote the positive/negative data set of the k -th task, h w ( x ; k ) denote the predictionscore for the k -th classifier, and ϕ ( s )
= 1 1+exp( − s ) is the sigmoid function.
The upper-level objectiveis strongly concave in terms of α k and the lower-level objective is strongly convex in terms of λ k .
We develop a tailored algorithm based on Algorithm 1 for solving the above formulation of multi-taskpAUC maximization as shown in Algorithm 2.
For the hessian update, the momentum update (2)is efficient due to that the each lower-level problem is only one-dimensional.
For simplicity ofimplementation, we define a loss G t +1 (25) in the Appendix B, on which auto-differentiation can bedirectly applied for computing a gradient estimator.
We refer readers to Appendix B for a detailedexplanation and derivation of G t +1 .
Experiments
4.1 Multi-task Deep AUC Maximization with Compositional Training
Data.
We use four datasets, namely CIFAR100, CheXpert, CelebA and ogbg-molpcba.
CIFAR-100[21] is an image dataset consisting of 60 ,
32 × 32 color images in 100 classes.
Hence, there are 100tasks for CIFAR100.
We follow 45 , 000 / 5 , 000 / 10 , 000 split to construct training/validation/testingdatasets.
CelebA [25] is a large-scale face attributes dataset with more than 200K celebrity images,each with 40 attribute annotations (i.e., 40 tasks).
We use the recommended training/validation/testingsplit as 162 , 770 / 19 , 866 / 19 , 961 .
CheXpert [17] is a dataset that contains 224,316 chest radiographswith 14 observations.
Since the official testing dataset is not open to public, we take the officialvalidation set as the testing data, and take the last 1000 images in the training dataset for validation.
Due to the absence of positive samples for the observation Fracture in the testing dataset, we ignorethis label and only consider the rest 13 observations (i.e., 13 tasks).
The last dataset ogbg-molpcba isa molecular property prediction graph dataset [14].
It consists of 437,929 graphs with 128 labels (i.e.,128 tasks).
We follow scaffold splitting procedure as recommended in [32].
Models.
We use ResNet18 [12] for CIFAR-100 and CelebA, and ImageNet pretrained DenseNet121[16] for CheXpert.
For ogbg-molpcba, we use Graph Isomorphism Network (GIN)
Setup.
We compare our method for optimizing the multi-task AUC maximization with compositionaltraining denoted by mAUC-CT (ours) with a baseline that directly optimizes multi-task min-max
AUC loss denoted by mAUC (baseline).
We do not compare with other straightforward baselines(e.g.
, optimizing the CE loss and the focal loss) since they have been shown to be inferior than AUCmaximization methods for imbalanced data in many previous works [38, 40].
For both methods, thelearning rates η 2 , η 1 , η 0 are set to be the same and tuned in { 0 . 01 , 0 .
The learningrates decay by a factor of 10 at the 4 th and 30 th epoch for CheXpert and CelebA, respectively.
Nolearning rate decay is applied for CIFAR-100 and ogbg-molpcba.
The moving average parameterβ 0 and ˜ η in the lower level problem of mAUC-CT (ours) are tuned in { 0 . 1 , 0 . 5 , 0 . 9 } .
Regarding thetask sampling, for datasets CIFAR-100 and ogbg-molpcba, 10 tasks are sampled to be updated ineach iteration, and for each sampled task, we independently sample a data batch of size 128 .
For theother two datasets with fewer tasks, CheXpert and CelebA, we sample one task at each iteration.
Thebatch size for data samples is 32 for CheXpert, and 128 for CelebA.
We run both methods the samenumber of epochs which varies on different data, 2000 epochs for CIFAR100, 6 epochs for CheXpert,40 epochs for CelebA and 100 epochs for obgb-molpcba.
Results.
In Table 1, we report the testing AUC score with the model selected according to thebest performance on validation datasets.
Comparing with optimizing AUC loss directly, mAUC-CT(ours) achieves better performance on all tested datasets.
We show the results of an ablation study in
Figure 1, which verifies convergence has a parallel speed-up effect on both the batch sizes of datasamples and task samples.
The algorithm converges faster as either the data or task sample batch sizeincreases.
In Figure 2 (left two), we compare our method with the baseline in terms of convergencespeed on the training data of two datasets, which demonstrate that our method converges faster.
Moreresults on training convergence are included in the appendix.
[Figure]
[Table]
4.2 Multi-task Deep Partial AUC Maximization
Setup.
For this task, we use the same datasets and the same networks for the image datasets andgraph datasets as in the previous subsection.
For baselines, we compare with a naive mini-batch basedmethod (MB) for pAUC maximization [19], and a state-of-the-art pAUC maximization method SOPAs [40].
Following previous works [38, 40], for pAUC maximization methods we use a pretrainedencoder network by optimizing the CE loss as the initial encoder and learn the whole network bymaximizing pAUC.
We also report the performance of optimizing the CE loss for a refernece.
Forall methods, the learning rate is tuned in { 0 . 0001 , 0 . 0005 , 0 . 001 , 0 .
The hyperparametersselection of MMB-pAUC are: η 1 and η 2 ∈ { 0 . 5 , 0 .
01 } , β 1 ∈ { 0 . 99 , 0 . 9 , 0 . 5 , 0 . 1 , 0 .
01 } andβ 0 ∈ { 0 . 9 , 0 . 99 } .
The momentum parameters in SOPA-s are tuned in the same range and their λparameter in { 0 . 1 , 1 , 10 } as in [40].
The margin parameter in the surrogate loss (e.g., c ) is set to be 1 .
Regarding the task sampling, we sample one task at each iteration for ogbg-molpcba and CheXpert,sample 10 tasks for CIFAR100, and sample 4 tasks for CelebA.
The data sample batch size is 32 for
CheXpert, and 64 for others.
For smaller datasets (CIFAR100 and ogbg-molpcba), we run 100 epochsfor each, and we decay the learning rate by a factor of 10 at the 50-th epoch.
For larger datasets(CelebA, CheXpert), we run 50 and 5 epochs respectively.
Results.
The partial AUC scores with FPR ≤ 0 .
1 on the testing data of different methods are shownin Table 2.
From the results, we can see that our methods perform better than baseline methods witha significant margin.
In Figure 2 (right two), we compare our method with the baselines in terms ofconvergence speed on the training data of two datasets, which demonstrate that our method convergesfaster.
More results on training convergence are included in the appendix.
Conclusion and Future Work
We have developed a simple single loop randomized stochastic algorithm for solving multi-blockmin-max bilevel optimization problems.
This algorithm requires updates for only constant number ofblocks in each iteration.
We showed that it achieves an oracle complexity of O ( ϵ − 4 ) , which matchesthe optimal complexity order for solving stochastic nonconvex optimization under a general unbiasedstochastic oracle model.
One limitation of the proposed algorithm is that the computation of hessianinverse estimator could be high for high dimensional problems.
We plan to address this issue in thefuture work. At the same time, we hope our work inspires others to find more novel applications ofour idea.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A](d)
Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]2.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes]
(c) Did you report error bars (e
g., with respect to the random seed after running experiments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, or cloud provider)?
[N/A]
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes]
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]