Non-convex Distributionally Robust Optimization: Non-asymptotic Analysis

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):

where $d$ measures the distance between two probability distributions, and the positive number $\epsilon$ 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):

where $\lambda>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 $\psi$-divergence, which is a popular choice in DRO literature (Namkoong and Duchi, 2016; Shapiro, 2017). Specifically, for a non-negative convex function $\psi$ such that $\psi(1)=0$ and two probability distributions $P,Q$ such that $Q$ is absolutely continuous w.r.t. $P$, the $\psi$-divergence between $Q$ and $P$ is defined as

which satisfies $d_{\psi}(Q,P)\geq 0$ and $d_{\psi}(Q,P)=0$ if $Q=P$ a.s.

The above analysis applies to a broad class of divergence functions $\psi$. We further discuss special cases when $\psi$ 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.

Assuming that $\psi^{*}$ is further Lipschitz, in Section 3.4 we prove that vinilla SGD suffices to achieve the $\mathcal{O}(\epsilon^{-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 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 $\psi$-divergence. $\psi$-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, $\chi^{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 $\mathcal{O}(\sqrt{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 $\mathcal{O}(\epsilon^{-6})$ complexity. We contribute to this line of work by proposing a different divergence with similar behavior as CVaR and an $\mathcal{O}(\epsilon^{-4})$ complexity.

Preliminaries

We make the following assumptions throughout the paper:

$\psi$ is a valid divergence function, i.e. a non-negative convex function satisfying $\psi(1)=0$ and $\psi(t)=+\infty$ for all $t<0$. Furthermore the conjugate $\psi^{*}$ is $M$-smooth.

We finally define the notion of $\epsilon$-stationary points for differentiable non-convex functions.

2 Equivalent formulation of the DRO objective

The aim of this paper is to find an $\epsilon$-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))

Under the Assumption 2.4, $\Psi(x)$ is differentiable, and $\nabla\Psi(x)=\nabla_{x}\mathcal{L}(x,\eta)$ for any $\eta\in\arg\min_{\eta^{\prime}}\mathcal{L}(x,{\eta^{\prime}})$.

Note that the $\eta$ in Lemma 2.6 may not be unique but the values of $\nabla_{x}\mathcal{L}(x,\eta)$ are all equal. Since $\Psi(x)$ is differentiable, the $\epsilon$-stationary points are well-defined. We now prove that the problem of finding an $\epsilon$-stationary point of $\Psi(x)$ is equivalent to finding an $\epsilon$-stationary point of a rescaled version of $\mathcal{L}(x,\eta)$.

Under the Assumption 2.4, if for some $(x,\eta)$ the following holds: $\|\nabla_{x}\mathcal{L}(x,\eta)\|+G|\nabla_{\eta}\mathcal{L}(x,\eta)|\leq\epsilon$, then $x$ is an $\epsilon$-stationary point of $\Psi(x)$. Furthermore, define a rescaled function

then $\|\nabla\widehat{\mathcal{L}}(x,\eta)\|\leq\epsilon/\sqrt{2}$ implies that $x$ is an $\epsilon$-stationary point of $\Psi(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 $\epsilon$-stationary point of $\widehat{\mathcal{L}}(x,\eta)$ such that $\|\nabla\widehat{\mathcal{L}}(x,\eta)\|\leq\epsilon$ (ignoring numerical constant $\sqrt{2}$). As a result, we will mainly work with $\widehat{\mathcal{L}}$ in subsequent analysis. The property of the objective function 5 heavily depends on $\psi^{*}$. We list some popular choices of $\psi$ together with the corresponding $\psi^{*}$ in Table 1. They serve as motivating examples of our subsequent analysis.

Analysis of general non-convex DRO

Now consider the DRO where $\psi$ is chosen as the commonly used $\chi^{2}$-divergence. Fix $\lambda=1$ and $\eta=0$. Based on the expression of $\psi^{*}(t)$ in Table 1, the DRO objective function 5 thus takes the form $\widehat{\mathcal{L}}(x,0;\xi)=\frac{1}{4}\left[x^{2}\left(1+\frac{\xi}{x^{2}+1}\right)^{2}+2\right]^{2}-1$, which is a quartic-like function. It follows that

$\widehat{\mathcal{L}}(x,0;\xi)=\Theta(x^{4})$ for large $x$ and therefore $\widehat{\mathcal{L}}(x,0;\xi)$ is not globally smooth;

$\nabla_{x}\widehat{\mathcal{L}}(x,0;\xi)=x^{3}+2x\xi+2x+\mathcal{O}(1)$ for large $x$ and the stochastic gradient variance $\operatorname{Var}[\nabla_{x}\widehat{\mathcal{L}}(x,0;\xi)]=\Theta(x^{2})$ which is unbounded globally.

2 Main results

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:

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.

Under Assumptions 2.4 and 3.2, the gradient estimators of (5) satisfies the following property:

Under Assumption 2.4, for any pair of parameters $(x,\eta)$ and $(x^{\prime},\eta^{\prime})$, we have the following property for the gradient of $\widehat{\mathcal{L}}$:

Note that 7 reduces to the standard notion of smoothness if the term $\frac{L}{G}\|\nabla\widehat{\mathcal{L}}(x,\eta)\|$ 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.

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.

Suppose that $F$ satisfies the following conditions:

(Generalized smoothness) $\|\nabla F(w_{1})-\nabla F(w_{2})\|\leq(K_{0}+K_{1}\|\nabla F(w_{1})\|)\|w_{1}-w_{2}\|$ holds for any $w_{1},w_{2}$;

Substituting Lemmas 3.4 and 3.3 into Theorem 3.5 immediately yields the final result:

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 $\epsilon$-stationary point of $\Psi(x)$ is

Corollary 3.6 shows that Algorithm 1 finds an $\epsilon$-stationary point with complexity $\mathcal{O}(\epsilon^{-4})$, which is the same as standard smooth non-convex optimization. Also note that the bound in Theorem 3.5 does not depend on $K_{1}$ and $\Gamma$ as long as $\epsilon$ 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 $\Gamma^{2}$, the required number of iterations $T$ is inversely propositional to $\Gamma^{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.

3 Proof sketch of Theorem 3.5

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:

(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:

The above lemma suggests that the algorithm should take a small step size when $\left\|\nabla F(x)\right\|$ 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:

Consider the algorithm that starts at $w_{0}$ and makes updates $w_{t+1}=w_{t}-\gamma\frac{m_{t+1}}{\left\|m_{t+1}\right\|}$ where $\{m_{t}\}$ is an arbitrary sequence of points. Define $\delta_{t}:=m_{t+1}-\nabla F(w_{t})$ be the estimation error. If $\gamma=O(1/K_{1})$, then

which is $\gamma\|\nabla F(w_{t})\|-2\gamma\|\delta_{t}\|-\mathcal{O}(\gamma^{2})$ for small $\gamma$. Therefore the objective function $F(w)$ decreases if $\|\delta_{t}\|<1/2\cdot\|\nabla F(w_{t})\|$, i.e. a small estimation error. However, $\delta_{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 $\{m_{t}\}$ to reduce the noise. Formally, let $\beta$ be the momentum factor and define $\hat{\delta_{t}}=\hat{\nabla}F(w_{t})-\nabla F(w_{t})$, then using the recursive equation of momentum $m_{t}$ in Algorithm 1 we can show that

Finally, for a suitable choice of $\gamma$ we can obtain the minimum gradient complexity bound on $T$.

4 Dealing with the CVaR case

Previous analysis applies to any divergence function $\psi$ as long as $\psi^{*}$ is smooth. This includes some popular choices such as the $\chi^{2}$-divergence, but not the CVaR. In the case of CVaR, $\psi^{*}$ 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 $\alpha$, $\psi_{\alpha}(t)$ takes zero when $t\in[0,1/\alpha)$ and takes infinity otherwise. Instead, we consider the following smoothed version of CVaR:

It is easy to see that $\psi_{\alpha}^{\text{smo}}$ is a valid divergence. The corresponding conjugate function is

The following propositions demonstrate that $\psi_{\alpha}^{\text{smo}}$ is indeed a smoothed approximation of CVaR.

Fix $0<\alpha<1$. When $\lambda\rightarrow 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 $\lambda$.

$\psi^{\text{smo},*}_{\alpha}(t)$ is $\frac{1}{\alpha}$-Lipschitz and $\frac{1}{4\alpha}$-smooth.

Based on Proposition 3.10, we can then use Corollary 3.6 to obtain the gradient complexity (taking $M=1/4\alpha$).

Note that $\psi^{\text{smo},*}_{\alpha}(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 $\widehat{\mathcal{L}}(x,\eta,\xi)$ can be bounded, as shown in the following lemma:

Suppose Assumption 2.4 holds. For smoothed CVaR, the DRO objective 5 satisfies

Moreover, $\widehat{\mathcal{L}}(x,\eta)$ is $K$-smooth with $K=\frac{L}{\alpha}+\frac{G^{2}}{2\lambda\alpha}$.

Equipped with the above lemma, we can obtain the following guarantee for smoothed CVaR, which shows that vanilla SGD suffices for convergence.

Suppose that $\psi=\psi_{\alpha}^{\text{smo}}$ and Assumption 2.4 holds. If we run SGD with properly selected hyper-parameters on the loss $\widehat{\mathcal{L}}(x,\eta)$, then the gradient complexity of finding an $\epsilon$-stationary point of $\Psi(x)$ is $\mathcal{O}\left(\alpha^{-3}\lambda^{-1}G^{2}(G^{2}+\lambda L)\Delta\epsilon^{-4}\right)$, where $\Delta=\mathcal{L}(x_{0},\eta_{0})-\inf_{x}\Psi(x)$.

The above theorem shows a similar convergence rate compared with Corollary 3.6 in terms of $\epsilon$ and $G$, and the dependency on $\lambda$ is even better. Therefore the Lipschitz property of $\psi^{*}$ is very useful, in that it is now possible to use a simpler algorithm while achieving a similar (or even better) bound.

Experiments

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 $\lambda=0.1$ and the CVaR coefficient $\alpha=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.

2 Experimental results

Results are demonstrated in Figure 1. For each figure, we plot the value of the DRO objective $\Psi(x)$ through the training process. Here we calculate $\Psi(x)=\min_{\eta}\mathcal{L}(x,\eta)$ at each epoch based on a convex optimization on $\eta$ until convergence (rather than using ${\mathcal{L}}(x,\eta)$ with the current parameter $\eta$ directly).

Experimental result for $\chi^{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 $\chi^{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. $\lambda$, the number of epochs, and the learning rate schedule).

Discussion

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 $\lambda$. Secondly, while this paper mainly considers smooth $\psi^{*}$, in some cases $\psi^{*}$ 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 $\psi$-divergence. Finally, we are looking forward to seeing more applications of DRO in real-world problems.

Acknowledgement

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 PKU-Baidu Fund. Jikai Jin is partially supported by the elite undergraduate training program of School of Mathematical Sciences in Peking University.

References

Appendix A Equivalent formulation of the DRO objective

We first assume $\Psi(x)$ is non-smooth and non-convex. To measure the convergence of non-smooth non-convex optimization, we define the notion called the generalized gradient [Clarke, 1990, Chapter 2].

and the generalized gradient at $x$ is the set

Interested readers may refer to the book [Clarke, 1990] for an in-depth exploration of this concept. Importantly, $\partial f(x)$ is a non-empty closed convex set; $\partial f(x)$ degenerates to a single point $\{\nabla f(x)\}$ if $f$ is smooth, and $\partial f(x)$ is equivalent to the sub-gradient if $f$ is convex. If $x$ is a local minima (or maxima) for $f(x)$, then $0\in\partial f(x)$. The following proposition gives the relationship between generalized gradient and (conventional) gradient.

([Clarke, 1990, Section 2.2]) If function $f$ is differentiable at $x$, then $f$ is local Lipschitz near $x$ and $\partial f(x)=\{\nabla f(x)\}$. Conversely, if $f$ is local Lipschitz near $x$ and $\partial f(x)$ reduces to a singleton point $\{g\}$, then $f$ is differentiable at $x$ and $\nabla f(x)=g$.

A.2 Proof of Lemma 2.6

We first present a basic lemma which provides a rule to calculate generalized gradients of the pointwise maxima of a function family.

where $T(x)=\left\{t\in T:F(x)=f(x,t)\right\}$, $\partial_{x}f(x,t)$ is the partial generalized gradient and $\operatorname{Conv}$ denotes a convex hull of a point set.

Assume Assumption 2.4 holds. Fix a point $x_{0}\in\mathcal{X}$. Denote $\eta_{0}\in\operatorname{argmin}_{\eta}\mathcal{L}(x_{0},\eta)$ be an arbitrary minima. Then for any point $x\in\mathcal{B}_{r}(x_{0})$ near $x_{0},$ there exists $\eta_{x}\in\operatorname{argmin}_{\eta}\mathcal{L}(x,\eta)$, such that $|\eta_{0}-\eta_{x}|\leq Gr$.

Using the condition that $\psi^{*}$ is convex and differentiable, we have $\eta_{x}\in\operatorname{argmin}_{\eta}\mathcal{L}(x,\eta)$ if and only if $\nabla_{\eta}\mathcal{L}(x,\eta)=0$. Namely,

Since $\nabla_{\eta}\mathcal{L}(x,\eta)$ is continuous in $\eta$, there must exists an $\eta_{x}\in[\eta_{0}-Gr,\eta_{0}+Gr]$, such that

Therefore $\eta_{x}\in\operatorname{argmin}_{\eta}\mathcal{L}(x,\eta)$. $\square$

For any point $x_{0}$, we can use Lemma A.4 by substituting $\mathcal{X}=\mathcal{B}_{r}(x_{0})$ and $\mathcal{T}=[\eta_{0}-Gr,\eta_{0}+Gr]$. The procedure is as follows:

We finally prove below (Lemma A.6) that $\{\nabla_{x}\mathcal{L}(x,\eta):\eta\in\operatorname{argmin}_{\eta}\mathcal{L}(x,\eta)\}$ is a singleton set. Then Proposition A.3 indicates that $\Psi(x)$ is differentiable, and the generalized gradient reduces to gradient such that $\nabla\Psi(x)=\nabla_{x}\mathcal{L}(x,\eta)$ for any $\eta\in\operatorname{argmin}_{\eta}\mathcal{L}(x,\eta)$. Thus we complete the proof of Lemma 2.6.

Assume Assumption 2.4 holds. For any $\eta_{1},\eta_{2}\in\operatorname{argmin}_{\eta}\mathcal{L}(x,\eta)$, we have $\nabla_{x}\mathcal{L}(x,\eta_{1})=\nabla_{x}\mathcal{L}(x,\eta_{2})$.

Denote $X(x,\eta)$, $Y(x)$ be two random functions defined by

which depend on the random variable $\xi$. Rewrite the gradient of $\mathcal{L}(x,\eta)$ as follows:

Note that $(\psi^{*})^{\prime}$ is monotonically increasing (due to the convexity of $\psi^{*}$), thus $X(x,\eta)$ is monotonically decreasing in $\eta$. It follows that

A.3 Proof of Theorem 2.7

Now, suppose that we have obtained a pair $(x,\eta)$ s.t. $\left\|\nabla_{x}\mathcal{L}(x,\eta)\right\|+G\left|\nabla_{\eta}\mathcal{L}(x,\eta)\right|\leq\epsilon$. Let $x$ be fixed and $\eta^{*}\in\mathop{\arg\min}_{\eta}\mathcal{L}(x,\eta)$. Then we have

where we use the fact that $(\psi^{*})^{\prime}$ is monotone increasing (due to the comvexity of $\psi^{*}$. Hence, using Lemma 2.6 we obtain

Now suppose that $\|\nabla\widehat{\mathcal{L}}(x,\eta)\|\leq\epsilon/\sqrt{2}$. Then

Using $(a+b)^{2}\leq 2(a^{2}+b^{2})$ we obtain

Appendix B The Stochastic Projected Gradient Descent algorithm for DRO with bounded loss

In this section we use a simple projected gradient method to minimize the DRO objective (5) and analyze its convergence rate under the assumption that the loss is bounded. Since this section is not so related to the main result in our paper, we mainly provide the gradient complexity bound in terms of $\epsilon$ for finding an $\epsilon$-stationary point without delving into problem-dependent parameters.

It turns out that we can restrict the feasible region to $\mathcal{X}\times\left[{U},{V}\right]$ where $\left[{U},{V}\right]$ is a finite interval.

Under the Assumptions 2.4 and B.1, the DRO problem is equivalent to

where $U=-\frac{\lambda C_{\psi}}{G}$ and $V=\frac{B-\lambda C_{\psi}}{G}$ are real numbers and $C_{\psi}$ is a constant depending only on $\psi$.

Note that $(\psi^{*})^{\prime}$ is a function satisfying the following properties:

$(\psi^{*})^{\prime}$ is monotonically increasing;

$0\leq\lim\limits_{s\to-\infty}(\psi^{*})^{\prime}(s)\leq 1$. This is because $\lim\limits_{s\to-\infty}\frac{\psi^{*}(s)}{s}=\lim\limits_{s\to-\infty}\inf_{t\geq 0}t-\frac{\psi(t)}{s}=\min\{t:\psi(t)<+\infty\}\in$ since $\psi(1)=0$;

$\lim\limits_{s\to+\infty}(\psi^{*})^{\prime}(s)\geq 1$ (possibly be $+\infty$). This is because $\frac{\psi^{*}(s)}{s}=\sup_{t\geq 0}t-\frac{\psi(t)}{s}\geq 1$ for $s>0$ since $\psi(1)=0$.

Therefore there exists a constant $C_{\psi}$ depending only on $\psi$ such that $(\psi^{*})^{\prime}(C_{\psi})=1$.

For any $x\in\mathcal{X}$, the optimal $\eta^{*}$ satisfies the following equation:

We now show there exists an optimal $\eta^{*}$ such that $G\eta^{*}\in[-\lambda C_{\psi},B-\lambda C_{\psi}]$. In fact, we have

Suppose Assumption 2.4 holds. Under Proposition B.2, $\mathcal{L}$ is $K$ smooth on $\mathcal{X}\times[U,V]$, where $K$ only depends on $\psi,\lambda,M,B,G$ and $L$.

Therefore $\mathcal{L}$ is smooth. $\square$

Suppose Assumption 2.4 holds. Under Proposition B.2, the stochastic gradients are unbiased estimates of the true gradients $\nabla_{x}\mathcal{L}$ and $\nabla_{\eta}\mathcal{L}$ and are uniformly bounded over $\mathcal{X}\times{[U,V]}$, by a constant $\Lambda$ which only depends on $\psi,\lambda,M,B,G$ and $L$.

Following [Ghadimi et al., 2016, Reddi et al., 2016], in constrained optimization we typically consider a generalized gradient defined as

[Ghadimi et al., 2016, Corollary 3], combined with Proposition B.4 implies that if $\gamma=1/2L$,

For any $\epsilon>0$, we choose $T=2K\Delta\epsilon^{-2}$ and $S=12\Lambda^{2}\epsilon^{-2}$, then 24 implies that

Thus the sample complexity of Algorithm 1 for finding $\epsilon$-stationary point is upper bounded by $24K\Lambda^{2}\Delta\epsilon^{-4}$. In this case, with pobability $\geq 0.5$ the gradient norm is upper bounded by $2\epsilon$, the conclusion follows. $\square$

as the interval constraint for $\eta$. Using parameters specified in Theorem B.5, Algorithm 2 arrives at $(x,\eta)$ with $\left\|\nabla\Psi(x)\right\|\leq\epsilon$ with probability $\geq 0.5$.

However, $\eta^{+}\leq\eta$, therefore it can only be that $\eta^{+}=\eta=\eta_{0}$. Therefore we still have

Appendix C Proofs in Section 3.2

In this section we present the proof of main results in Section 3.2. For convenience we restate the results before proving them.

Under Assumptions 2.4 and 3.2, the gradient estimators of (5) satisfies the following property:

For a random vector $X$, define the sum of its element-wise variance as

Then it is easy to check that, for i.i.d. random vectors $X_{1},X_{2}$ we have

We first bound the variance of the stochastic gradient $\nabla_{x}\widehat{\mathcal{L}}(x,\eta;\xi)$. Indeed we have

Here in 33 we use that fact that $(a+b)^{2}\leq 2(a^{2}+b^{2})$ for any $a,b$; in 34 we use Assumption 2.4. Now we deal with the first term. Using $2(a-1)^{2}+2\geq a^{2}$ for any $a$, we have

Under Assumption 2.4, for any pair of parameters $(x,\eta)$ and $(x^{\prime},\eta^{\prime})$, we have the following property for the gradient of $\widehat{\mathcal{L}}$:

First write $\nabla\widehat{\mathcal{L}}(x,\eta)$ as

We then split $\nabla\mathcal{L}(x,\eta)-\nabla\mathcal{L}(x^{\prime},\eta^{\prime})$ into two terms $A+B$, where

where we use $(\psi^{*})^{\prime}(s)\geq 0$ for all $s$. $B$ can be bounded as follows:

Combining the above inequalities, we obtain

C.2 Proof of Theorem 3.5

We formalize the generalized smoothness property into a definition.

We now present a descent inequality for $(K_{0},K_{1})$-smooth functions which will be used in subsequent analysis.

(Descent Inequality) Let $F$ be $(K_{0},K_{1})$-smooth, then for any point $x$ and direction $z$ the following holds:

C.2.2 Properties of the normalized update

Let $\mu\geq 0$ be a real constant. For any vectors $u$ and $v$,

Now we can characterize the behavior of normalization-based algorithms in terms of function value descent.

Consider the algorithm that starts at $w_{0}$ and makes updates $w_{t+1}=w_{t}-\gamma\frac{m_{t+1}}{\left\|m_{t+1}\right\|}$ where $\{m_{t}\}$ is an arbitrary sequence of points. Define $\delta_{t}:=m_{t+1}-\nabla F(w_{t})$ be the estimation error. Then

Since $\|w_{t+1}-w_{t}\|=\gamma$, by Lemma C.4 we have

where in the second inequality we use Lemma C.5. $\square$

C.2.3 A general convergence result

Instead of directly focusing on the specific problem of DRO, we first provide convergence guarantee for Algorithm 2 under general smoothness and noise assumptions.

Suppose that $F$ is $(K_{0},K_{1})$-smooth and the stochastic gradient estimator $\nabla F(w,\xi)$ is unbiased and satisfies

Define the estimation errors $\delta_{t}:=m_{t+1}-\nabla F(w_{t})$. Denote $H(a,b):=\nabla F(a)-\nabla F(b)$. We can upper bound $H(a,b)$ using the definition of $(K_{0},K_{1})$-smoothness:

Using the definition of momentum $m_{t}$ and $H(a,b)$, we can get a recursive formula on $\delta_{t}$:

Using triangle inequality and plugging in the estimate 46, we have

Taking a telescope summation of 48 we obtain

Now we take expectation of $\left\|\sum_{\tau=0}^{t}\beta^{\tau}\hat{\delta}_{t-\tau}\right\|$ over all the randomness. We will prove a core lemma ( Lemma C.9) later which shows

If we choose $\gamma=\frac{1}{8}(\min(K_{1}^{-1},K_{0}^{-1}\epsilon)(1-\beta)$, and $S=64\Gamma^{2}$, then

Set $1-\beta=\min(4\Lambda^{-2}\Gamma^{2}\epsilon^{2},1)$ and $m_{0}=\|\nabla F(w_{0})\|$, then

If $\epsilon\leq\mathcal{\min}\left(\frac{K_{0}}{K_{1}},\frac{\Lambda}{2\Gamma}\right)$, then the gradient complexity is $512\Lambda^{2}\Delta K_{0}\epsilon^{-4}$. $\square$

Suppose the DRO problem 3 satisfies Assumptions 2.4 and 3.2. Using Algorithm 1 with a constant batch size 4096, the gradient complexity for finding an $\epsilon$-stationary point of $\Psi(x)$ is

Lemmas 3.3 and 3.4 imply that the conditions in Theorem 3.5 for $\widehat{\mathcal{L}}(x,\eta)$ are satisfied with $K_{0}=L+2G^{2}\lambda^{-1}M,\Gamma^{2}=64,\Lambda^{2}=11G^{2}M^{2}\lambda^{-2}\sigma^{2}+8G^{2}$. The main result immediately follows from Theorems 3.5 and 2.7. $\square$

We now return to prove the core lemma that is used in 50.

Let $\hat{\delta}_{t}=\hat{\nabla}F(w_{t})-\nabla F(w_{t})$ be the stochastic noise. Then

We prove the following result: for each $i\in\{0,1,\cdots,t+1\}$, the following inequality holds:

It is easy to see that Lemma C.9 follows by setting $i=t+1$ in 53.

We prove 53 by induction. When $i=0$, 53 holds obviously. Now suppose 53 holds for $i$, and we want to prove that 53 holds for $i+1$.

Appendix D Proofs in Section 3.4

In this section we prove the main result of Section 3.4 for smoothed CVaR. Recall the expressions

The following proposition shows that $\psi^{\text{smo},*}_{\alpha}$ is Lipschitz-continuous and smooth.

$\psi^{\text{smo},*}_{\alpha}(t)$ is $\frac{1}{\alpha}$-Lipschitz and $\frac{1}{4\alpha}$-smooth.

where we use $\alpha(1-\alpha)\leq\frac{1}{4}$. Hence the conclusion follows. $\square$

Fix $0<\alpha<1$. When $\lambda\rightarrow 0$, the solution of the DRO problem 5 for smoothed CVaR tends to the solution for the standard CVaR.

For the standard CVaR, the DRO problem can be written as

which is irrelevant to $\lambda$. For smoothed CVaR, the DRO problem can be written as

Suppose Assumption 2.4 holds. For smoothed CVaR, the DRO objective 5 satisfies

Moreover, $\widehat{\mathcal{L}}(x,\eta)$ is $K$-smooth with $K=\frac{L}{\alpha}+\frac{G^{2}}{2\lambda\alpha}$.

since $\psi^{*}$ is non-decreasing and $\frac{1}{\alpha}$-Lipschitz continuous.

We also have $\left\|\nabla_{\eta}\widehat{\mathcal{L}}(x,\eta;\xi)\right\|\leq\alpha^{-1}G$. Therefore $\left\|\nabla\widehat{\mathcal{L}}(x,\eta)\right\|^{2}\leq 2\alpha^{-2}G^{2}$.

Now we turn to the smoothness of $\mathcal{L}$. For any $(x,\eta)$ and $(x^{\prime},\eta^{\prime})$ we decouple $\nabla\widehat{\mathcal{L}}(x,\eta)-\nabla\widehat{\mathcal{L}}(x^{\prime},\eta^{\prime})$ into $A+B$ using the same approach as in 40. Now different from 41, $A$ can be bounded by

using the Lipschitz property of $\psi^{*}$. The bound for $B$ is the same as 42:

Hence $\mathcal{L}$ is $K$-smooth as desired. $\square$

Suppose that $\psi=\psi_{\alpha}^{\text{smo}}$ and Assumption 2.4 holds. If we run SGD with properly selected hyper-parameters on the loss $\widehat{\mathcal{L}}(x,\eta)$, then the gradient complexity of finding an $\epsilon$-stationary point of $\Psi(x)$ is $\mathcal{O}\left(\alpha^{-3}\lambda^{-1}G^{2}(G^{2}+\lambda L)\Delta\epsilon^{-4}\right)$, where $\Delta=\mathcal{L}(x_{0},\eta_{0})-\inf_{x}\Psi(x)$.

It is well-known [Ghadimi and Lan, 2013] that the complexity of SGD for finding an $\epsilon$-stationary point is $\mathcal{O}(\Delta K\sigma^{2}\epsilon^{-4})$ if the objective function is $K$-smooth and $\sigma^{2}$ is an upper bound of the variance of stochastic gradients. Now the proof can be completed by using Lemma D.3. $\square$

Appendix E Experiment

Imbalanced CIFAR-10. To demonstrate the effectiveness of our method in DRO-classification setting, we construction an imbalanced classification dataset. The original version of CIFAR-10 contains 50,000 training images and 10,000 validation images of size 32$\times$32 with 10. To create their imbalanced version, we reduce the number of training examples per class and keep the validation set unchanged. We consider the type of random imbalance and use $\rho_{i}$ to denote the sample ratio of $i$th class between the imbalanced and original dataset. $\rho=\{0.804,0.543,0.997,0.593,0.390,0.285,0.959,0.806,0.967,0.660\}$

E.2 Implementation details

For every training task we jointly tune the parameters learning rate for baseline and our method by grid search and pick the one that achieves the fastest optimization. By default we set momentum = 0.9 for all experiments and $\epsilon=0.1$ for normalized SGD. We use batch size n = 128 throughout.

Hyper-parameter for $\chi^{2}$ penalized DRO. In regression setting, we use SGD with lr=0.0002 as our baseline algorithm and set lr=0.005 for normalized SGD. In classification setting, we set lr=0.005 and 0.01 for baseline and our method, respectively.

Hyper-parameter for smoothed CVaR. In smooth CVaR, we also divide experiment into two part, regression and classification task. We train CVaR with lr = (0.00005, 0.00005) and smoothed CVaR with lr = (0.001, 0.0001) in regression and classification setting.