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 studies the family of distributionally robust optimization (DRO) problems with non-convex objectives and KL-divergence. With the motivation of improving the computational efficiency on existing methods and applying to online settings, the authors proposed a new duality-free online stochastic methods for solving the considered DRO problem. Theoretical analyses on complexities are provided for the proposed algorithm with and without a PL condition. Comparisons with state-of-the-art DRO methods are presented over ERM on imbalanced datasets.