An Online Method for A Class of Distributionally Robust Optimization with Non-convex Objectives

1 Introduction . Distributionally robust optimization ( DRO ) has received tremendous attention in machine learning due to its capability to handle noisy data , adversarial data and imbalanced classification data [ 42 , 33 , 4 ] . Given a set of observed data { z1 , . . . , zn } , where zi = ( xi , yi ) , a DRO formulation can be written as : min w∈Rd max p∈∆n Fp ( w ) = n∑ i=1 pi ` ( w ; zi ) − h ( p,1/n ) + r ( w ) , ( 1 ) where w denotes the model parameter , ∆n = { p ∈ Rn : ∑ i pi = 1 , pi ≥ 0 } denotes a ndimensional simplex , ` ( w ; z ) denotes a loss function on data z , h ( p,1/n ) is a divergence measure between p and uniform probabilities 1/n , and r ( w ) is convex regularizer of w. When ` ( w ; z ) is a convex function ( e.g. , for learning a linear model ) , many stochastic primal-dual methods can be employed for solving the above min-max problem [ 35 , 20 , 49 , 48 , 32 ] . When ` ( w ; z ) is a non-convex function ( e.g. , for learning a deep neural network ) , some recent studies also proposed stochastic primal-dual methods [ 41 , 48 ] . ∗The first two authors make equal contributions . Correspondence to tianbao-yang @ uiowa.edu . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) , virtual . However , stochastic primal-dual methods for solving DRO problems with a non-convex ` ( w ; z ) loss function ( e.g. , the predictive model is a deep neural network ) suffer from several drawbacks . First , primal-dual methods need to maintain and update a high-dimensional dual variable p ∈ Rn for large-scale data , whose memory cost is as high as O ( n ) per-iteration . Second , existing primal-dual methods usually need to sample data according to probabilities p in order to update w , which brings additional costs than random sampling . Although random sampling can be used for computing the stochastic gradient in terms of w , the resulting stochastic gradient could have n-times larger variance than using non-uniform sampling according to p ( please refer to the supplement for a simple illustration ) . Third , due to the constraint on p ∈ ∆n , the min-max formulation ( 1 ) is not friendly to online learning in which the data is received sequentially and n is rarely known in prior . Can we design an efficient online algorithm to address the DRO formulation ( 1 ) without dealingwith p ∈ Rn for a non-convex objective that is applicable to deep learning ? To address this question , we restrict our attention to a family of DRO problems , in which the KL divergence h ( p,1/n ) = λ ∑ i pi log ( npi ) is used for regularizing the dual variables p , where λ > 0 is a regularization parameter . We note that this consideration does not impose strong restriction to the modeling capability . It has been shown that for a family of divergence functions h ( p,1/n ) , different DRO formulations are statistically equivalent to a certain degree [ 12 ] . The proposed method is based on an equivalent minimization formulation for h ( p,1/n ) = λ ∑ i pi log ( npi ) . In particular , by maximizing over p exactly , ( 1 ) is equivalent to min w∈Rd { Fdro ( w ) = λ log ( 1/n n∑ i=1 exp ( ` ( w ; zi ) /λ ) ) ) + r ( w ) } . ( 2 ) In an online learning setting , we can consider a more general formulation : min w∈Rd { Fdro ( w ) = λ log ( Ez exp ( ` ( w ; z ) /λ ) ) + r ( w ) } . ( 3 ) The above problem is an instance of stochastic compositional problems of the following form : min w∈Rd F ( w ) : = f ( Ez [ gz ( w ) ] ) + r ( w ) , ( 4 ) by setting f ( s ) = λ log ( s ) , s ≥ 1 and gz ( w ) = exp ( ` ( w ; z ) /λ ) . Stochastic algorithms have been developed for solving the above compositional problems . [ 44 ] proposed the first stochastic algorithms for solving ( 4 ) , which are easy to implement . However , their sample complexities are sub-optimal for solving ( 4 ) . Recently , a series of works have tried to improve the convergence rate by using advanced variance reduction techniques ( e.g. , SVRG [ 19 ] , SPIDER [ 13 ] , SARAH [ 36 ] ) . However , most of them require using a mega mini-batch size in the order of O ( n ) or O ( 1/ ) 2 at every iteration or many iterations for updating w , which hinders their applications on large-scale problems . In addition , these algorithms usually use a constant step size , which may harm the generalization performance . This paper aims to develop more practical stochastic algorithms for solving ( 3 ) without suffering from the above issues in order to enable practitioners to explore the capability of DRO for deep learning with irregular data ( e.g. , imbalanced data , noisy data ) . To this end , we proposed an online stochastic method ( COVER ) and its restarted variant ( RECOVER ) . We establish a state-of-the-art complexity of COVER for finding an -stationary solution and a state-of-the-art complexity of RECOVER under a Polyak-Łojasiewicz ( PL ) condition of the problem . PL condition has been widely explored for developing practical optimization algorithms for deep learning [ 52 ] . Compared with other stochastic algorithms , the practical advantages of RECOVER are : 1 . RECOVER is an online duality-free algorithm for addressing large-scale KL regularized DRO problem that is independent of the high dimensional dual variable p ∈ Rn , which makes it suitable for deep learning applications . 2 . RECOVER also enjoys the benefits of stagewise training similar to existing stochastic methods for deep learning [ 52 ] , i.e. , the step size is decreased geometrically in a stagewise manner . 2 is either the objective gap accuracy F ( w ) − minF ( w ) ≤ or the gradient norm square bound ‖∇F ( w ) ‖2 ≤ In addition , this paper also makes several important theoretical contributions for stochastic non-convex optimization , including 1 . We establish a nearly optimal complexity for finding -stationary point , i.e. , ‖∇F ( w ) ‖2 ≤ , for a class of two-level compositional problems in the order of Õ ( 1/ 3/2 ) without a large mini-batch size , which is better than existing results [ 44 , 45 , 14 , 5 ] . 2 . We etablish an optimal complexity for finding -optimal solution under an µ−PL condition for a class of two-level compositional problems in the order of O ( 1/ ( µ ) ) without a large mini-batch size , which is better than existing results [ 53 ] . A theoretical comparison between our results and existing results is shown in Table 1 . Empirical studies vividly demonstrate the effectiveness of RECOVER for deep learning on imbalanced data . 2 Related Work . DRO has been extensively studied in machine learning [ 31 , 11 , 40 ] , statistics , and operations research [ 42 ] . In [ 33 ] , the authors proved that minimizing the DRO formulation with a quadratic regularization in a constraint form is equivalent to minimizing the sum of the empirical loss and a variance regularization defined on itself . Variance regularization can enjoy better generalization error compared with the empirical loss minimization [ 33 ] , and was also observed to be effective for imbalanced data [ 33 , 56 ] . Recently , [ 12 ] also establishes this equivalence for a broader family of regularization function h ( p,1/n ) including the KL divergence . Several recent studies have developed stochastic primal-dual methods for solving DRO with a nonconvex loss function ` ( w ; z ) assuming it is smooth or weakly convex [ 41 , 26 , 30 , 48 ] . [ 41 ] proposed the first primal-dual methods for solving weakly convex concave problems . For online problems , their algorithms for finding an -stationary solution whose gradient norm square is less than have a sample complexity of O ( 1/ 2 ) or O ( 1/ 3 ) with or without leveraging the strong concavity of h ( p,1/n ) for finding an -stationary point . Recently Liu et al . [ 28 ] proposed to leverage the PL condition of the objective function to improve the convergence for a non-convex min-max formulation of AUC maximization . Then , a PES-SGDA algorithm is proposed to solve a more general class of non-convex min-max problems by leveraging the PL condition [ 15 ] . Both [ 28 ] and [ 15 ] have used geometrically decreasing step sizes in a stagewise manner . However , their algorithms ’ complexity is in the order of O ( 1/µ2 ) , which is worse than O ( 1/µ ) achieved in this paper . Similarly , [ 50 ] also leveraged PL conditions to solve non-convex min-max problems and has a sample complexity of O ( 1/µ2 ) . Nevertheless , the step size of their algorithm is decreased polynomially in the order of O ( 1/t ) , which usually yields poor performance for deep learning . All the methods reviewed above require maintaining and updating both the primal variable w and a high dimensional dual variable p ∈ Rn . Recently , Levy et al . [ 23 ] considered different formulations of DRO , which includes our considered KL-regularized DRO formulation as a special case . Their assumed that the loss function is convex and proposed a stochastic method with a sample complexity O ( 1/ 3 ) for sovling the KL-regularized DRO formulation . In contrast , we provide a better sample complexity in the order of O ( 1/ ) under a PL condition without convexity assumption . Additionally , their method requires a large batch size in the order of O ( 1/ ) , while our method only requires a constant batch size which is more practical . We also notice that a recent work [ 24 ] and its extended version [ 25 ] have considered a formulation similar to ( 2 ) and proposed a stochastic algorithm . However , their algorithm has a slower convergence rate with an O ( 1/ 2 ) complexity for finding an - stationary point and anO ( 1/ ( µ2 ) ) complexity for finding an -optimal solution under a PL condition . Our work is a concurrent work appearing online earlier than [ 24 ] . To the best of our knowledge , this is the first work trying to solve the non-convex DRO problem with a duality-free stochastic method by formulating the min-max formulation into an equivalent stochastic compositional problem . There are extensive studies for solving stochastic compositional problems . [ 44 ] considered a more general family of stochastic compositional problems and proposed two algorithms . When the objective function is non-convex , their algorithm ’ s complexity isO ( 1/ 7/2 ) for finding an -stationary solution . This complexity was improved in their later works [ 14 ] , reducing to O ( 1/ 2 ) . When the objective is smooth , several papers proposed to use variance reduction techniques ( e.g. , SPIDER , SARAH ) to improve the complexity for finding a stationary point [ 53 , 18 , 55 , 27 ] . The best sample complexity achieved for online problems is O ( 1/ 3/2 ) [ 53 , 55 ] . [ 53 ] also considered the PL condition for developing a faster algorithm called restarted CIVR , whose sample complexity is O ( 1/µ ) . However , these variance reduction-based methods require using a very large mini-batch size at many iterations , which has detrimental influence on training deep neural networks [ 43 ] . To address this issue , [ 9 ] proposed a new technique called STORM that integrates momentum and the recursive variance reduction technique for solving stochastic smooth non-convex optimization . Their algorithm does not require a large mini-batch size at every iterations and enjoys a sample complexity of O ( log2/3 ( 1/ ) / 3/2 ) for finding an -stationary point . However , their algorithm uses a polynomially decreasing step size , which is not practical for deep learning , and is not directly applicable to stochastic composite problems .

Summary: This paper propose an online algorithm to solve a class of Distributionally Robust Optimization problem with non-convex objective functions. The standard primal-dual method typically suffer from optimization over high dimensional variable, which is impractical in many real world applications. The sota complexity of proposed method is provided with and without the PL condition. Empirical experiments demonstrate the superiority of the proposed algorithm.