Non-convex Distributionally Robust Optimization: Non-asymptotic Analysis

1 Introduction . For a classical machine learning problem , the goal is typically to train a model over a training set that achieves good performance on a test set , where both the training set and the test set are drawn from the same distribution P . While such an assumption is reasonable and simple for theoretical analysis , it is often not the case in real applications . For example , this setting may be improper when there is a gap between training and test distribution ( e.g . in domain adaptation tasks ) [ Zhang et al. , 2021 ] , when there is severe class imbalance in the training set [ Sagawa et al. , 2020 ] , when fairness in minority groups is an important consideration [ Hashimoto et al. , 2018 ] , or when the deployed model is exposed to adversarial attacks [ Sinha et al. , 2018 ] . Distributionally robust optimization ( DRO ) , as a popular approach to deal with the above situations , has attracted great interest for the machine learning research communities in recent years . In contrast to classic machine learning problems , for DRO it is desired that the trained model still has good performance under distribution shift . Specifically , DRO proposes to minimize the worst-case loss over a set of probability distributions Q around P . This can be formulated as the following constrained optimization problem [ Rahimian and Mehrotra , 2019 , Shapiro , 2017 ] : minimizex∈X Ψ ( x ) : = sup Q∈U ( P ) Eξ∼Q [ ` ( x ; ξ ) ] ( 1 ) ∗Equal Contribution , alphabetical order . †Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . where x ∈ X is the parameter to be optimized , ξ is a sample randomly drawn from distribution Q , and ` ( x ; ξ ) is the loss function so that Eξ∼Q [ ` ( x ; ξ ) ] is the expected loss over distribution Q . The DRO objective Ψ ( x ) is therefore the worst-case loss when the distribution P is shifted to Q . The set U ( P ) is called the uncertainty set and typically defined as U ( P ) : = { Q : d ( Q , P ) ≤ } ( 2 ) where d measures the distance between two probability distributions , and the positive number corresponds to the magnitude of the uncertainty set . Instead of imposing a hard constrained uncertainty set , sometimes it is more preferred to use a soft penalty term , resulting in the penalized DRO problem [ Sinha et al. , 2018 ] : minimizex∈X Ψ ( x ) : = sup Q { Eξ∼Q [ ` ( x ; ξ ) ] − λd ( Q , P ) } ( 3 ) where λ > 0 is the regularization coefficient . There are many possible choices of d. A detailed discussion of different distance measures and their properties can be found in Rahimian and Mehrotra [ 2019 ] . In this paper we consider a general class of distances d called the ψ-divergence , which is a popular choice in DRO literature [ Namkoong and Duchi , 2016 , Shapiro , 2017 ] . Specifically , for a non-negative convex function ψ such that ψ ( 1 ) = 0 and two probability distributions P , Q such that Q is absolutely continuous w.r.t . P , the ψ-divergence between Q and P is defined as dψ ( Q , P ) : = ∫ ψ ( dQ dP ) dP . which satisfies dψ ( Q , P ) ≥ 0 and dψ ( Q , P ) = 0 if Q = P a.s . The main focus of this paper is to study efficient first-order optimization algorithms for DRO problem ( 3 ) for non-convex losses ` ( x , ξ ) . While non-convex models ( especially deep neural networks ) have been extensively used in DRO setting ( e.g . Sagawa et al . [ 2020 ] ) , theoretical analysis about the convergence speed is still lacking . Most previous works ( e.g . Levy et al . [ 2020 ] ) assume the loss ` ( · , ξ ) is convex , and in this case ( 3 ) is equivalent to a convex optimization problem ( see Section 2 for details ) . Recently some works provide convergence rates of algorithms for non-convex losses in certain special cases , e.g . the divergence measure ψ is chosen as the conditional-value-at-risk ( CVaR ) and the loss function has some nice structural properties [ Soma and Yoshida , 2020 , Kalogerias , 2020 ] . Gürbüzbalaban et al . [ 2020 ] considered a more general setting but only proved an asymptotic convergence result for non-convex DRO . Compared with these works , we provide the first non-asymptotic analysis of optimization algorithms for DRO with general smooth non-convex losses ` ( x , ξ ) and general ψ-divergence . In this setting , there are two major difficulties we must encounter : ( i ) the DRO objective Ψ ( x ) is non-convex and can become arbitrarily non-smooth , causing standard techniques in smooth non-convex optimization fail to provide a good convergence guarantee ; ( ii ) the noise of the stochastic gradient of Ψ ( x ) can be arbitrarily large and unbounded even if we assume the gradient of the inner loss ` ( x , ξ ) has bounded variance . To tackle these challenges , we propose to optimize the DRO objective using mini-batch normalized SGD with momentum , and we are able to prove an O ( −4 ) complexity of this algorithm . The core technique here is to exploit the specific structure of Ψ ( x ) , which shows that ( i ) the DRO objective satisfies a generalized smoothness condition [ Zhang et al. , 2020a , b ] and ( ii ) the variance of the stochastic gradient can be bounded by the true gradient . This motivates us to adopt the special algorithm that combines gradient normalization and momentum techniques into SGD , by which both non-smoothness and unbounded noise can be tackled , finally resulting in an O ( −4 ) complexity similar to standard smooth non-convex optimization . The above analysis applies to a broad class of divergence functions ψ . We further discuss special cases when ψ has additional properties . In particular , to handle the CVaR case ( a non-differentiable loss ) , we propose a divergence function which is a smoothed variant of CVaR and is further Lipschitz . In this case we show that a convergence guarantee can be established using vanilla SGD , and an similar complexity bound holds . We highlight that the algorithm and analysis in this paper are not limited to DRO setting , and are described in the context of a general class of optimization problem . Our analysis clearly demonstrates the effectiveness of gradient normalization and momentum techniques in optimizing ill-conditioned objective functions . We believe our result can shed light on why some popular optimizers , in particular Adam [ Kingma and Ba , 2015 ] , often exhibit superior performance in real applications . Contributions . We summarize our main results and contributions below . Let ψ∗ be the conjugate function of ψ ( see Definition 2.3 ) . For non-convex optimization problems , since obtaining the global minima is NP-hard in general , this paper adopts the commonly used ( relaxed ) criteria : to find an -approximate first-order stationary point of the function Ψ ( see Definition 2.5 ) . We measure the complexity of optimization algorithms by the number of computations of the stochastic gradient ∇ ` ( x , ξ ) to reach an -stationary point . • Assuming that ψ∗ is smooth and the loss ` is Lipschitz and smooth ( possibly non-convex or unbounded ) , we show in Section 3.2 that the mini-batch normalized momentum algorithm ( cf . Algorithm 1 ) has a complexity of O ( −4 ) . • Assuming that ψ∗ is further Lipschitz , in Section 3.4 we prove that vinilla SGD suffices to achieve the O ( −4 ) complexity . As a special case , we propose a new divergence which is a smoothed approximation of CVaR . • We conduct experiments to verify our theoretical results . We observe that our proposed methods significantly accelerate the optimization process , and also demonstrates superior test performance . 1.1 Related work . Constrained DRO and Penalized DRO . There are two existing formulations of the DRO problem : the constrained DRO and the penalized DRO . The constrained DRO formulation ( 1 ) has been studied in a number of works [ Namkoong and Duchi , 2016 , Shapiro , 2017 , Duchi and Namkoong , 2018 ] , while other works consider the penalty-based formulation ( 3 ) [ Sinha et al. , 2018 , Levy et al. , 2020 ] . From a Lagrangian perspective , the two formulations are equivalent ; however , the dual objective of the constrained formulation is sometimes hard to solve as pointed out in [ Namkoong and Duchi , 2016 , Duchi and Namkoong , 2018 ] . In this paper we focus on the penalty-based version and provide the first non-asymptotic analysis in the non-convex setting . Moreover , we do not make the assumption that the loss is bounded , as assumed in Levy et al . [ 2020 ] in the convex setting . DRO with ψ-divergence . ψ-divergence is one of the most common choices in DRO literature to measure the distance between probability distributions . It encompasses a variety of popular functions such as KL-divergence , χ2-divergence , and the conditional-value-at-risk ( CVaR ) , etc . Table 1 gives detailed descriptions for these functions . For CVaR , Namkoong and Duchi [ 2016 ] proposed a mirror-descent method which achieves O ( √ T ) regret . Levy et al . [ 2020 ] proposed a stochastic gradient-based method with optimal convergence rate in the convex setting . They also discussed an alternative approach based on the dual formulation which they call Dual SGM . In the non-convex setting , Soma and Yoshida [ 2020 ] proposed a smoothed approximation of CVaR and obtain an O ( −6 ) complexity . We contribute to this line of work by proposing a different divergence with similar behavior as CVaR and an O ( −4 ) complexity . For χ2 divergence , Hashimoto et al . [ 2018 ] considered a constrained formulation of DRO but did not provide theoretical guarantees . Levy et al . [ 2020 ] proposed algorithms based on an multi-level MonteCarlo stochastic gradient estimator , and provide convergence guarantees in the convex setting . In contrast , we consider general smooth non-convex loss function ` and provide convergence guarantee for χ2 divergence as a special case of Corollary 3.6 . Non-smooth non-convex optimization . Conventional non-convex optimization typically focuses on smooth objective functions . For general smooth non-convex stochastic optimization , it is already known that the best possible gradient complexity for finding an -approximate stationary point is O ( −4 ) [ Arjevani et al. , 2019 ] , which is achieved by SGD based algorithms [ Ghadimi and Lan , 2013 ] . However , the optimization can be much harder for non-smooth non-convex objective functions , and there are limited results in this setting . Ruszczynski [ 2020 ] proposed a stochastic gradient-based method which converges to a stationary point with probability one , under the assumption that the feasible region is bounded . For unconstrained optimization , Zhang et al . [ 2020c ] showed that it is intractable to find an -stationary point for some Lipschitz and Hadamard semi-differentiable function . When the function is weakly convex , Davis and Drusvyatskiy [ 2019 ] showed that the projected SGD converges to the stationary point of a Moreau envelope , and a recent work [ Mai and Johansson , 2020 ] extended this result to SGD with momentum . In this paper , we show that for smooth non-convex loss ` , DRO can be formulated as a non-smooth non-convex optimization problem , but the special property of the DRO objective makes it possible to find an -stationary point within O ( −4 ) complexity .

Summary: This paper bridges the gap by studying DRO algorithms for general smooth non-convex losses. By carefully exploiting the specific form of the DRO objective with $\psi$ divergence, this paper provides non-asymptotic convergence guarantees even though the objective function is possibly non-convex, non-smooth, and has unbounded gradient noise. In particular, the authors prove that a special algorithm called mini-batch normalized gradient descent with momentum can find an $\epsilon$-first-order stationary point within $O(\epsilon^{-4})$ gradient complexity. Furthermore, for the conditional value-at-risk (CVaR) setting, this paper proposes a penalized DRO objective based on a smoothed version of the CVaR and obtains better complexity. The theoretical results are verified in several tasks, and the proposed algorithm can consistently achieve prominent acceleration.