Non-convex Distributionally Robust Optimization: Non-asymptotic Analysis

Abstract:
Distributionally robust optimization (DRO) is a widely-used approach to learn models that are robust against distribution shift. Compared with the standard optimization setting, the objective function in DRO is more difficult to optimize, and most of the existing theoretical results make strong assumptions on the loss function. In this work we bridge the gap by studying DRO algorithms for general smooth non-convex losses. By carefully exploiting the specific form of the DRO objective, we are able to provide non-asymptotic convergence guarantees even though the objective function is possibly non-convex, non-smooth and has unbounded gradient noise. In particular, we prove that a special algorithm called the mini-batch normalized gradient descent with momentum, can find an -first-order stationary point within O( −4) gradient complexity. We also discuss the conditional valueat-risk (CVaR) setting, where we propose a penalized DRO objective based on a smoothed version of the CVaR that allows us to obtain a similar convergence guarantee. We finally verify our theoretical results in a number of tasks and find that the proposed algorithm can consistently achieve prominent acceleration.

heading: 1 Introduction
text: 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.

heading: 1.1 Related work
text: 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.

heading: 2 Preliminaries

heading: 2.1 Notations and Assumptions
text: Throughout this paper we use ‖ · ‖ to denote the `2-norm in an Euclidean space Rd and use 〈·, ·〉 to denote the standard inner product. For a real number t, denote (t)+ as max(t, 0). For a set C, denote IC(·) as the indicator function such that IC(x) = 0 if x ∈ C and IC(x) = +∞ otherwise. We first list some basic definitions in optimization literature, which will be frequently used in this paper. Definition 2.1. (Lipschitz continuity) A mapping f : X → Rm is G-Lipschitz continuous if for any x, y ∈ X , ‖f(x)− f(y)‖ ≤ G ‖x− y‖. Definition 2.2. (Smoothness) A function f : X → R is L-smooth if it is differentiable on X and the gradient ∇f is L-Lipschitz continuous, i.e. ‖∇f(x)−∇f(y)‖ ≤ L ‖x− y‖ for all x, y ∈ X . We say f is non-smooth if such L does not exist. Definition 2.3. (Conjugate function) For a function ψ : R→ R, the conjugate function ψ∗ is defined as ψ∗(t) := sups∈R (st− ψ(s)). Assumption 2.4. We make the following assumptions throughout the paper:
• Given ξ, the loss function `(x, ξ) is G-Lipschitz continuous and L-smooth with respect to x;
• ψ is a valid divergence function, i.e. a non-negative convex function satisfying ψ(1) = 0 and ψ(t) = +∞ for all t < 0. Furthermore the conjugate ψ∗ is M -smooth.
We finally define the notion of -stationary points for differentiable non-convex functions. Definition 2.5. ( -stationary point) For a differentiable function f : X → R, a point x ∈ X is said to be first-order -stationary if ‖∇f(x)‖ ≤ .

heading: 2.2 Equivalent formulation of the DRO objective
text: The aim of this paper is to find an -stationary point of problem (3). However, the original formulation (3) involves a max operation over distributions which makes optimization challenging. By duality arguments we can show that the DRO objective (3) can be equivalently written as (see detailed derivations in [Levy et al., 2020, Section A.1.2])
Ψ(x) = min η∈R
λEξ∼Pψ∗ ( `(x; ξ)− η
λ
) + η. (4)
Thus, to minimize Ψ(x) in (4), one can jointly minimize L(x, η) := Eξ∼P [ λψ∗ ( `(x;ξ)−η
λ
) + η ]
over (x, η) ∈ X × R ⊂ Rn+1. This can be seen as a standard stochastic optimization problem. The remaining thing is to show one can find an -stationary point of Ψ(x) by optimizing L(x, η) instead. We first present a lemma that gives connection of the gradient of Ψ(x) to the gradient of L(x, η). Lemma 2.6. Under the Assumption 2.4, Ψ(x) is differentiable, and ∇Ψ(x) = ∇xL(x, η) for any η ∈ arg minη′ L(x, η′).
Note that the η in Lemma 2.6 may not be unique but the values of ∇xL(x, η) are all equal. Since Ψ(x) is differentiable, the -stationary points are well-defined. We now prove that the problem of finding an -stationary point of Ψ(x) is equivalent to finding an -stationary point of a rescaled version of L(x, η). Theorem 2.7. Under the Assumption 2.4, if for some (x, η) the following holds: ‖∇xL(x, η)‖ + G|∇ηL(x, η)| ≤ , then x is an -stationary point of Ψ(x). Furthermore, define a rescaled function
L̂(x, η) = L(x,Gη) := Eξ∼P [ λψ∗ ( `(x; ξ)−Gη
λ
) +Gη ] , (5)
then ‖∇L̂(x, η)‖ ≤ / √ 2 implies that x is an -stationary point of Ψ(x).
The proof of Lemma 2.6 and Theorem 2.7 can be found in Appendix A. From the above theorem it suffices to find an -stationary point of L̂(x, η) such that ‖∇L̂(x, η)‖ ≤ (ignoring numerical constant √ 2). As a result, we will mainly work with L̂ in subsequent analysis. The property of the objective function (5) heavily depends on ψ∗. We list some popular choices of ψ together with the corresponding ψ∗ in Table 1. They serve as motivating examples of our subsequent analysis.

heading: 3 Analysis of general non-convex DRO
text: In this section we will analyze the DRO problem with general smooth non-convex loss functions `. We first discuss the challenges appearing in our analysis, then show how to leverage the specific structure of the objective function in order to overcome these challenges. Specifically, we show that our proposed algorithm can achieve a non-asymptotic complexity of O( −4).

heading: 3.1 Challenges in non-convex DRO
text: A standard result in optimization literature states that if the objective function is smooth and the stochastic gradient is unbiased and has bounded variance3, then standard stochastic gradient descent (SGD) algorithms can provably find an -first-order stationary point underO( −4) gradient complexity [Ghadimi and Lan, 2013]. Here the smoothness and bounded variance property are crucial for the convergence of SGD [Zhang et al., 2019]. However, we find that both assumptions are violated in non-convex DRO, even if the inner loss `(x, ξ) is smooth and the stochastic noise is bounded for both `(x, ·) and ∇x`(x, ·). We present a counter example to illustrate this point, in which we can gain some insight about the structure of the DRO objective.
Example 3.1. Consider the loss `(x; ξ) = x2 (
1 + ξx2+1
)2 which is a quadratic-like function with
noise ξ, where ξ is a Rademacher variable drawn from {−1,+1} with equal probabilities. Then a straightforward calculation shows that the loss ` has the following properties:
• (Smoothness) For any ξ ∈ {−1,+1}, `(x, ξ) is L-smooth with respect to x for L = 8; • (Bounded variance) For any x ∈ R, Eξ [( `(x, ξ)− x2 )2] = 4x 4 (x2+1)2 + x4
(x2+1)4 ≤ 4. It then follows that Varξ[`(x, ξ)] ≤ 4;
• (Bounded variance for gradient) Similarly we can check that the gradient of ` also has bounded variance. Moreover, the variance tends to zero when x→∞.
Now consider the DRO where ψ is chosen as the commonly used χ2-divergence. Fix λ = 1 and η = 0. Based on the expression of ψ∗(t) in Table 1, the DRO objective function (5) thus takes the
form L̂(x, 0; ξ) = 14
[ x2 (
1 + ξx2+1
)2 + 2 ]2 − 1, which is a quartic-like function. It follows that
• L̂(x, 0; ξ) = Θ(x4) for large x and therefore L̂(x, 0; ξ) is not globally smooth;
• ∇xL̂(x, 0; ξ) = x3 + 2xξ + 2x + O(1) for large x and the stochastic gradient variance Var[∇xL̂(x, 0; ξ)] = Θ(x2) which is unbounded globally.
As we can see from the above example, both the local smoothness and the gradient variance of L̂ strongly rely on the scale of x. Indeed, in general non-convex DRO both the two quantities have a positive correlation with the magnitude of `. As shown in Appendix B, if we make the additional assumption that ` is bounded by a small constant, then the smoothness and gradient noise can be controlled in a straightforward way, and we show that a projected stochastic gradient method can be applied in this setting. However, such bounded loss assumption is quite restrictive and not satisfactory.
3Eξ∼P ‖∇x`(x, ξ)−∇x`(x)‖2 ≤ σ2 for some σ and all x ∈ X where `(x) = Eξ∼P `(x, ξ).

heading: 3.2 Main results
text: In this section, we present the main theoretical result of this paper. All proofs can be founded in Appendix C. We make the following assumption on the noise of the stochastic loss: Assumption 3.2. We assume that for all x ∈ X , the stochastic loss has bounded variance, i.e. Eξ∼P (`(x, ξ)− `(x))2 ≤ σ2 where `(x) = Eξ∼P `(x, ξ).
We now provide formal statements of the key properties mentioned above, which show that both the gradient variance and the local smoothness can be controlled in terms of the gradient norm. Lemma 3.3. Under Assumptions 2.4 and 3.2, the gradient estimators of (5) satisfies the following property:
Eξ‖∇L̂(x, η, ξ)−∇L̂(x, η)‖2 ≤ 11G2M2λ−2σ2 + 8(G2 + ‖∇L̂(x, η)‖2) (6)
Lemma 3.4. Under Assumption 2.4, for any pair of parameters (x, η) and (x′, η′), we have the following property for the gradient of L̂:
‖∇L̂(x, η)−∇L̂(x′, η′)‖ ≤ ( K + LG‖∇L̂(x, η)‖ ) ‖(x− x′, η − η′)‖ (7)
where K = L+ 2G2λ−1M .
Note that (7) reduces to the standard notion of smoothness if the term LG‖∇L̂(x, η)‖ is absent. Thus the inequality (7) can be seen as a generalized smoothness condition. Zhang et al. [2020b] for the first time proposed such generalized smoothness for twice-differentiable functions in a different form, and Zhang et al. [2020a] further gave a comprehensive analysis of algorithms for optimizing generalized smooth functions. However, all these works make strong assumptions on the gradient noise and can not be applied in our setting.
Instead, we propose to use the mini-batch normalized SGD with momentum algorithm for non-convex DRO, shown in Algorithm 1. The algorithm has been theoretically analysed in [Cutkosky and Mehta, 2020] for optimizing standard smooth non-convex functions. Compared with Cutkosky and Mehta [2020], we use mini-batches in each iteration in order to ensure convergence in our setting.
Algorithm 1: Mini-batch Normalized SGD with Momentum Input :The objective function F (w), distribution P , initial point w0, initial momentum m0,
learning rate γ, momentum factor β, batch size S and total number of iterations T 1 for t← 1 to T do 2 ∇̂F (wt−1)← 1S ∑S i=1∇F (wt−1, ξ (i) t−1) where {ξ (i) t−1}Si=1 are i.i.d. samples drawn from P 3 mt ← βmt−1 + (1− β)∇̂F (wt−1) 4 wt ← wt−1 − γ
mt ‖mt‖
The following main theorem establishes convergence guarantee of Algorithm 1. We further provide a sketch of proof in Section 3.3, where we can gain insights on how normalization and momentum techniques help tackle the difficulties shown in Lemmas 3.3 and 3.4. Theorem 3.5. Suppose that F satisfies the following conditions:
• (Generalized smoothness) ‖∇F (w1)−∇F (w2)‖ ≤ (K0 +K1‖∇F (w1)‖)‖w1−w2‖ holds for any w1, w2;
• (Gradient variance) The stochastic gradient∇F (w, ξ) is unbiased (∇F (w) = Eξ∇F (w, ξ)) and satisfies Eξ ‖∇F (w, ξ)−∇F (w)‖2 ≤ Γ2 ‖∇F (w)‖2 + Λ2 for some Γ and Λ.
Let {wt} be the sequence produced by Algorithm 1. Then with a mini-batch size S = Θ(Γ2) and a suitable choice of parameters γ and β, for any small = O(min(K0/K1,Λ/Γ)), we need at most O ( ∆K0Λ 2 −4 )
gradient complexity to guarantee that we find an -stationary point in expectation, i.e. 1T ∑T−1 t=0 E‖∇F (wt)‖ ≤ where ∆ = F (w0)− infw∈Rd F (w).
Substituting Lemmas 3.3 and 3.4 into Theorem 3.5 immediately yields the final result:
Corollary 3.6. Suppose the DRO problem (3) satisfies Assumptions 2.4 and 3.2. Using Algorithm 1 with a constant batch size, the gradient complexity for finding an -stationary point of Ψ(x) is
O ( G2 ( M2σ2λ−2 + 1 ) ( λ−1MG2 + L ) ∆ −4 ) .
Corollary 3.6 shows that Algorithm 1 finds an -stationary point with complexity O( −4), which is the same as standard smooth non-convex optimization. Also note that the bound in Theorem 3.5 does not depend on K1 and Γ as long as is sufficiently small. In other words, Algorithm 1 is well-adapted to the non-smoothness and unbounded noise in our setting. We also point out that although the batch size is chosen propositional to Γ2, the required number of iterations T is inversely propositional to Γ2, therefore the total number of stochastic gradient computations remains the same.
Finally, note that Theorem 3.5 is stated in a general form and is not limited to DRO setting. It greatly extends the results in Zhang et al. [2020a,b] by relaxing their noise assumptions, and demonstrates the effectiveness of combining adaptive gradients with momentum for optimizing ill-conditioned objective functions. More importantly, our algorithm is to some extent similar to currently widely used optimizers in practice, e.g. Adam. We believe our result can shed light on why these optimizers often show superior performance in real applications.

heading: 3.3 Proof sketch of Theorem 3.5
text: Below we present our proof sketch, in which the motivation of using Algorithm 1 will be clear. Similar to standard analysis in non-convex optimization, we first derive a descent inequality for functions satisfying the generalized smoothness: Lemma 3.7. (Descent inequality) Let F (x) be a function satisfying the generalized smoothness condition in Theorem 3.5. Then for any point x and direction z the following holds:
F (x− z) ≤ F (x)− 〈∇F (x) , z〉+ 1 2 (K0 +K1 ‖∇F (x)‖) ‖z‖2 . (8)
The above lemma suggests that the algorithm should take a small step size when ‖∇F (x)‖ is large in order to decrease F . This is the main motivation of considering a normalized update. Indeed, after some careful calculation we can prove the following result: Lemma 3.8. Consider the algorithm that starts at w0 and makes updates wt+1 = wt − γ mt+1‖mt+1‖ where {mt} is an arbitrary sequence of points. Define δt := mt+1 − ∇F (wt) be the estimation error. If γ = O(1/K1), then
F (wt)− F (wt+1) ≥ ( γ − 1
2 K1γ
2 ) ‖∇F (wt)‖ − 1
2 K0γ
2 − 2γ‖δt‖ (9)
which is γ‖∇F (wt)‖−2γ‖δt‖−O(γ2) for small γ. Therefore the objective function F (w) decreases if ‖δt‖ < 1/2 · ‖∇F (wt)‖, i.e. a small estimation error. However, δt is related to the stochastic gradient noise which can be very large due to Lemma 3.3. This motivates us to the use the momentum technique for the choice of {mt} to reduce the noise. Formally, let β be the momentum factor and define δ̂t = ∇̂F (wt)−∇F (wt), then using the recursive equation of momentum mt in Algorithm 1 we can show that
δt = β t−1∑ τ=0 βτ (∇F (wt−τ−1)−∇F (wt−τ )) + (1− β) t−1∑ τ=0 βτ δ̂t−τ + (1− β)βtδ̂0. (10)
The first term of the right hand side in (10) can be bounded using the generalized smoothness condition, and the core procedure is to bound the second term using a careful analysis of conditional expectation and the independence of noises {δ̂t} (see Lemma C.9 in Appendix). Finally, the use of mini-batches of size Θ(Γ2), a carefully chosen β and a small enough γ ensure that ∑T−1 t=0 ‖δt‖ <
c ∑T−1 t=0 (E‖∇F (wt)‖ + O( )) where c < 1/2. This guarantees that the right hand side of (9) is overall positive, and by taking summation over t in (9) we have that
F (w0)− F (wT ) ≥ (1− 2c)γ T−1∑ t=0 ‖∇F (wt)‖ − O(γ2T − γT ).
namely, 1
T T−1∑ t=0 ‖∇F (wt)‖ ≤ O ( ∆ γT + γ + ) .
Finally, for a suitable choice of γ we can obtain the minimum gradient complexity bound on T .

heading: 3.4 Dealing with the CVaR case
text: Previous analysis applies to any divergence function ψ as long as ψ∗ is smooth. This includes some popular choices such as the χ2-divergence, but not the CVaR. In the case of CVaR, ψ∗ is not differentiable as shown in Table 1, which is undesirable from an optimization viewpoint. In this section we introduce a smoothed version of CVaR. The conjugate function of the smoothed CVaR is also smooth, so that the results in Section 3.2 can be directly applied in this setting.
For standard CVaR at level α, ψα(t) takes zero when t ∈ [0, 1/α) and takes infinity otherwise. Instead, we consider the following smoothed version of CVaR:
ψsmoα (t) =
{ t log t+ 1−αtα log 1−αt 1−α t ∈ [0, 1/α)
+∞ otherwise (11)
It is easy to see that ψsmoα is a valid divergence. The corresponding conjugate function is
ψsmo,∗α (t) = 1
α log(1− α+ α exp(t)). (12)
The following propositions demonstrate that ψsmoα is indeed a smoothed approximation of CVaR.
Proposition 3.9. Fix 0 < α < 1. When λ→ 0+, the solution of the DRO problem (5) for smoothed CVaR tends to the solution for the standard CVaR. Note that the solution of the standard CVaR does not depend on λ.
Proposition 3.10. ψsmo,∗α (t) is 1α -Lipschitz and 1 4α -smooth.
Based on Proposition 3.10, we can then use Corollary 3.6 to obtain the gradient complexity (taking M = 1/4α).
Note that ψsmo,∗α (t) is not only smooth but also Lipschitz. In this setting, we can in fact obtain a stronger result than the general one provided in Corollary 3.6. Specifically, the gradient noise and smoothness of the objective function L̂(x, η, ξ) can be bounded, as shown in the following lemma: Lemma 3.11. Suppose Assumption 2.4 holds. For smoothed CVaR, the DRO objective (5) satisfies
E‖∇L̂(x, η, ξ)‖2 ≤ 2α−2G2. (13)
Moreover, L̂(x, η) is K-smooth with K = Lα + G2 2λα .
Equipped with the above lemma, we can obtain the following guarantee for smoothed CVaR, which shows that vanilla SGD suffices for convergence.
Theorem 3.12. Suppose that ψ = ψsmoα and Assumption 2.4 holds. If we run SGD with properly selected hyper-parameters on the loss L̂(x, η), then the gradient complexity of finding an -stationary point of Ψ(x) is O ( α−3λ−1G2(G2 + λL)∆ −4 ) , where ∆ = L(x0, η0)− infx Ψ(x).
The above theorem shows a similar convergence rate compared with Corollary 3.6 in terms of and G, and the dependency on λ is even better. Therefore the Lipschitz property of ψ∗ is very useful, in that it is now possible to use a simpler algorithm while achieving a similar (or even better) bound.

heading: 4 Experiments
text: We perform two sets of experiments to verify our theoretical results. In the first set of experiments, we consider the setting in Section 3.2, where the loss `(x; ξ) is highly non-convex and unbounded, and ψ is chosen to be the commonly used χ2-divergence such that its conjugate is smooth. We will show that (i) the vanilla SGD algorithm cannot optimize this loss efficiently due to the non-smoothness of the DRO objective; (ii) by simply adopting the normalized momentum algorithm, the optimization process can be greatly accelerated. In the second set of experiments, we deal with the CVaR setting in Section 3.4. We will show that by employing the smooth approximation of CVaR defined in (11) and (12), the optimization can be greatly accelerated.

heading: 4.1 Experimental settings
text: Tasks. We consider two tasks: the classification task and the regression task. While classification is more common in machine learning, here we may also highlight the regression task, since recent studies show that DRO may be more suitable for non-classification problems in which the metric of interest is continuous as opposed to the 0-1 loss [Hu et al., 2018, Levy et al., 2020].
Datasets. We choose the AFAD-Full dataset for regression and CIFAR-10 dataset for classification. AFAD-Full [Niu et al., 2016] is a regression task to predict the age of human from the facial information, which contains more than 160K facial images and the corresponding age labels ranging from 15 to 75. Note that AFAD-Full is an imbalanced dataset where the ages of two thirds of the whole dataset are between 18 and 30. Following the experimental setting in [Chang et al., 2011, Chen et al., 2013, Niu et al., 2016], we split the whole dataset into a training set comprised of 80% data and a test set comprised of the remaining 20% data. CIFAR-10 dataset is a classification task consisting of 10 classes with 5000 images for each class. To demonstrate the effectiveness of our method in DRO setting, we adopt the setting in Chou et al. [2020] to construct an imbalanced CIFAR-10 by randomly sampling each category at different ratio. See Appendix for more details.
Model. For all experiments in this paper, we use the standard ResNet-18 model in [He et al., 2016]. The output has 10 logits for CIFAR-10 classification task, and has a single logit for regression.
Training details. We choose the penalty coefficient λ = 0.1 and the CVaR coefficient α = 0.02 in all experiments. For each algorithm, we tune the learning rate hyper-parameter from a grid search and pick the one that achieves the fastest optimization speed. The momentum factor is taken to 0.9 in all experiments, and the mini-batch size is chosen to be 128. We train the model for 100 epochs on CIFAR-10 dataset and 200 epochs on AFAD-Full dataset. Other training details can be found in Appendix E.

heading: 4.2 Experimental results
text: Results are demonstrated in Figure 1. For each figure, we plot the value of the DRO objective Ψ(x) through the training process. Here we calculate Ψ(x) = minη L(x, η) at each epoch based on a convex optimization on η until convergence (rather than using L(x, η) with the current parameter η directly).
Experimental result for χ2 penalized DRO. Figure 1(a) and Figure 1(b) plot the training curve of the DRO objective using different algorithms. It can be seen that in both regression and classification, vanilla SGD converges slowly, and using normalized momentum algorithm significantly improves the convergence speed. For example, in regression task SGD does not converge after 100 epochs while normalized momentum algorithm converges just after 25 epochs. These results highly consist with our theoretical findings, which shows that due to the non-smoothness of the DRO loss, vanilla SGD may not be able to optimize the loss well; In contrast, normalized momentum utilizes the relationship between local smoothness and gradient magnitude, and achieves better performance.
Experimental result for smoothed CVaR. Figure 1(c) and Figure 1(d) plot the training curves for different training losses: CVaR and smoothed CVaR. Note that the evaluation metrics (y-axis) in these figures are all chosen to be CVaR, even when the training objective is smoothed CVaR. In this way we can make a fair comparison of optimization speed based on these training curves. Firstly, it can be seen that the optimization of CVaR is very hard due to the non-smoothness, and the training curves have lots of spikes. In contrast, the optimization of smoothed CVaR is much easier for both tasks, and the final loss is significantly lower. Such experimental results show the benefit of our proposed smoothed CVaR for optimization.
Test performance. We also measure the test performance of trained models to see whether a better optimizer can also improve test accuracy. Due to space limitation, in the main text we provide results of χ2 penalized DRO problem for classification using unbalanced CIFAR-10 dataset, which is listed in Table 2. Other results can be found in Appendix E. It can be seen that the model trained using normalized SGD with momentum achieves higher test accuracy on all class, and especially, the worst-performing class. Since the experiments in this paper is mainly designed to compare algorithms rather than to achieve best performance, better performance is likely to be reached if adjusting the hyper-parameters (e.g. λ, the number of epochs, and the learning rate schedule).

heading: 5 Discussion
text: Conclusion. In this paper we provide non-asymptotic analysis of first-order algorithms for the DRO problem with unbounded and non-convex loss. Specifically, we write the original DRO problem as a non-smooth non-convex optimization problem, and we propose an efficient normalization-based algorithm to solve it. The general result of Theorem 3.5 might be of independent value and is not limited to DRO setting. We hope that this work can also bring inspiration to the study of other non-smooth non-convex optimization problems.
Limitations. Despite the theoretical grounds and promising experimental justifications, there are some interesting questions that remain unexplored. Firstly, it may be possible to obtain better complexities on problem-dependent parameters, e.g. G and λ. Secondly, while this paper mainly considers smooth ψ∗, in some cases ψ∗ may be non-smooth (e.g. for KL-divergence) or even not continuous. In future we hope to discover approaches that can deal with more general classes of ψ-divergence. Finally, we are looking forward to seeing more applications of DRO in real-world problems.

heading: Acknowledgement
text: This work was supported by Key-Area Research and Development Program of Guangdong Province (No. 2019B121204008), National Key R&D Program of China (2018YFB1402600), BJNSF (L172037) and Beijing Academy of Artificial Intelligence. Project 2020BD006 supported by PKUBaidu Fund. Jikai Jin is partially supported by the elite undergraduate training program of School of Mathematical Sciences in Peking University.

Review summaries:
The authors propose an algorithm that they call mini-batch normalized gradient descent with momentum. This algorithm can find an \epislon-first-order stationary point under O(\epsilon^{-4}) gradient complexity, even though the objective function is non-convex and non-smooth and has unbounded gradient noise. Their work is limited to where all loss functions are G−Lipschitz and L−smooth in x and conjugate function of a divergence function is M−smooth. Their main algorithm goes like: get a mini-batch gradient estimator (using mini-batch to ensure convergence) and update momentum using that estimator and update their variable w using normalized momentum. They presented two numerical results: one classification, the other regression. They compared their algorithm with vanilla SGD regarding convergence speed.

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.

This paper studies distributionally robust optimization with the loss function being nonconvex and the uncertainty set defined by a divergence function. Particularly, this paper gives a very interesting generalized smoothness condition, shows the connection of the normalized gradient descent with this smoothness condition, and gives a non-asymptotic convergence rate. Finally, experimental results verify the theory.