DISTRIBUTIONALLY ROBUST OPTIMIZATION LEADS
Year: 2018

Abstract:
In this paper, we adopt distributionally robust optimization (DRO) (Ben-Tal et al., 2013) in hope to achieve a better generalization in deep learning tasks. We establish the generalization guarantees and analyze the localized Rademacher complexity for DRO, and conduct experiments to show that DRO obtains a better performance. We reveal the profound connection between SGD and DRO, i.e., selecting a batch can be viewed as choosing a distribution over the training set. From this perspective, we prove that SGD is prone to escape from bad stationary points and small batch SGD outperforms large batch SGD. We give an upper bound for the robust loss when SGD converges and keeps stable. We propose a novel Weighted SGD (WSGD) algorithm framework, which assigns high-variance weights to the data of the current batch. We devise a practical implement of WSGD that can directly optimize the robust loss. We test our algorithm on CIFAR-10 and CIFAR100, and WSGD achieves significant improvements over the conventional SGD.

heading: 1 INTRODUCTION
text: In recent years, supervised learning based on deep neural networks (DNNs) has achieved state-ofthe-art performances in various tasks of different domains, such as computer vision and natural language processing (Goodfellow et al., 2016). However, theoretical explanations for the generalization ability of deep learning remain challenging. Roughly speaking, supervised learning aims to learn a model from the training data. Let Θ denote the parameter space of the model, and Z denote the sample space. Let `(z, θ) : Z ×Θ→ R+ denote the loss function related to the parameter θ and the instance z. In the general setting of supervised learning, there is a fixed but unknown distribution over Z (denoted P ). We attempt to find a θ ∈ Θ, which gives a small generalization error. That is, we address the following optimization problem:
min θ∈Θ
R(θ) , EP `(z, θ). (1)
However, the real distribution P is unknown, so we consider an empirical distribution. Given the training samples S = (z1, z2, · · · , zm), the empirical distribution is defined as P̂m(z) = 1 m ∑m i=1 I[z=zi], where I is the indicator function. The loss related to this distribution is called the empirical loss, which is
R̂(θ) , EP̂m`(z, θ) = 1
m m∑ i=1 `(zi, θ), (2)
and we minimize the empirical loss instead. This procedure is referred to as Empirical Risk Minimization (ERM).
The overwhelming capacity and the paucity of data make deep learning architectures susceptible to over-fitting. This urges us to propose a better approach to the minimization problem (1). Suppose there is a distribution familyM = {Pλ : λ ∈ Λ} parametrized by λ, and the underlying distribution P = Pλ0 ∈ M, where λ0 is fixed but unknown. It is a common case in deep learning tasks that the structure of the data is complicated while the amount of data is small. Hence, when we use an estimator λ̂ to estimate the parameter λ0, the variance of λ̂ would be quite large. Thus, rather than directly minimizing the true expected loss R(θ) = EPλ0 `(z, θ), we minimize a distributionally
robust loss alternatively, which can ensure the true expected loss to be small with high probability. The distributionally robust loss is defined as:
R(θ,K) , sup λ∈K EPλ`(z, θ),
where K ⊂ Λ is a neighbourhood of the estimator λ̂. This procedure is called Distributionally Robust Optimization (DRO). Because we only have access to the given training dataset, we construct an unbiased1 estimation of the expected loss EPλ`(z, θ):
1
m m∑ i=1 Pλ(zi) Pλ0(zi) `(zi, θ).
This can be viewed as the expected loss with respect to a perturbed empirical generalized distribution2 P̂λ = 1m ∑m i=1 Pλ(z) Pλ0 (z)
I[z=zi], so the DRO problem becomes minimizing the empirical distributionally robust loss:
R̂(θ,K) , sup λ∈K EP̂λ`(z, θ) = sup λ∈K m∑ i=1 Pλ(zi) mPλ0(zi) `(zi, θ).
Keeping the motivations above in mind, we focus our attention on the robust loss, defined as the maximal possible weighted loss over a pre-specified weight set P for simplicity:
R̂S,P(θ) , sup p∈P m∑ i=1 pi`(zi, θ).
In deep learning, the high capacity of the model and the non-convexity of the objective function lead to an extremely complicated landscape and induce a lot of minima, and different minima have different generalization performances. In this case, an optimization algorithm is both an optimizer and a minimum selector, and the minimum preference of the optimization algorithm is crucial in deep learning. Stochastic Gradient Descent (SGD) and its variants are widely used to minimize the empirical loss (2). The simplest mini-batch SGD consists of iterative sampling of a batch of data and then updating the parameter:
θt+1 = θt − ηt ( 1 |Bt| ∑ zi∈Bt ∇θ`(zi, θt) ) ,
where Bt ⊂ S is the sampled mini-batch, ηt is the learning rate, and |Bt| denotes the batch size. We interpret SGD from another perspective: at each iteration, SGD randomly chooses an empirical distribution over the data and performs a gradient descent update w.r.t. the corresponding loss. This interpretation shows a profound connection between SGD and DRO. Motivated by the connection, we develop a brand new version of SGD that can achieve a better generalization performance. The main contribution of this paper is summarized as follows:
• We clarify the connection between mini-batch SGD and DRO. We show that SGD helps escape from bad stationary points and the robust loss of local minima where SGD converges is not large.
• We analyze the generalization property of robust loss and show the trade off: the robust loss generalizes well when P is relatively small, but a larger P is beneficial to reducing the localized Rademacher complexity.
• We propose a framework of improved SGD algorithm, which enables the algorithm to explore more empirical distributions over the data. We give a practical improved SGD algorithm that can directly optimize the robust loss.
1To make the fraction reasonable, we assume Pλ0(z) > 0 whenever Pλ(z) > 0. 2Strictly speaking, this is not a distribution because P̂λ(Z) may not equals to 1, but its expectation is 1 (the
expectation is taken over the sample S). Nonetheless, we can view P̂λ as a measure over Z .

heading: 2 RELATED WORK
text: SGD is widely used in large-scale machine learning problems. The convergence theory of SGD in convex cases is solid (Ghadimi & Lan, 2013; Shamir & Zhang, 2013). However, due to the nonconvexity of deep learning models, the behaviour of SGD is more complicated. Dauphin et al. (2014) argued that it is the prevalence of saddle points rather than local minima that causes the difficulty in high dimensional non-convex optimization. Ge et al. (2015) analyzed the property of perturbed GD (i.e., gradient descent with an isotropic noise injected). They showed that perturbed GD has the ability to escape from saddle points in a polynomial number of iterations, and then converges to a second-order stationary point.
It is also well known that adding perturbation and noise at different levels of the training processes helps generalization, such as dropout (Srivastava et al., 2014), and random noise injection to the input data (Chapelle et al., 2001). The good performance of SGD might benefit from the noise. Keskar et al. (2016) discovered that large-batch training will incur a degradation in the generalization performance of the model. They explained this phenomenon with the concept of sharp minima, and carried out numerical experiments to support the idea that SGD with large batch size tends to converge to sharp minima, which leads to a bad generalization ability. However, Dinh et al. (2017) provided theoretical results indicating that most existing definitions of sharpness and flatness of minima can not be applied to some commonly used non-negative homogeneous activation functions such as ReLU.
DRO was originally applied to decision making under uncertainty. Namkoong & Duchi (2017) introduced DRO to machine learning, and viewed it as a variance-based regularization that allows us to automatically trade-off between bias and variance. They proved an O( 1m ) approximation to the generalization errorR(θ), in contrast to the generalO( 1√
m ) approximation of the empirical risk.
Recently, DRO is used to achieve fairness without demographic information in fair machine learning (Hashimoto et al., 2018).
In learning theory, Rademacher complexity is widely used to measure the complexity of a hypothesis set (Bartlett & Mendelson, 2002). Different from the general Rademacher complexity which considers the entire hypothesis set, localized Rademacher complexity only cares about the near optimal hypothesis (Bartlett et al., 2002). In fact, the hypothesis selected by a learning algorithm usually has better performance than the worst hypothesis. Thus, localized Rademacher complexity is usually much smaller than the general Rademacher complexity.

heading: 3 NOTATION AND PRELIMINARIES
text: In our work, we assume the loss function `(z, θ) is bounded between 0 and 1, and is continuously differentiable. Let Σ denote the hyperplane {p : ∑m i=1 pi = 1} and Σ+ denote the simplex
{p : ∑m i=1 pi = 1,p ≥ 0}. Denote p0 = ( 1 m , . . . , 1 m ) >. For p = (p1, p2, . . . , pm)
> ∈ Σ, let R̂S,p(θ) denote ∑m i=1 pi`(zi, θ). Then the ERM loss can be written as R̂S,p0(θ), and the robust loss w.r.t. the weight set P can be viewed as:
R̂S,P(θ) = sup p∈P R̂S,p(θ).
Let conv(P) denote the convex hull of P . Given a v ∈ Σ, let σ(v) be the set of all the vectors whose entries are permutations of v’s. Let P(v) denote the convex hull of σ(v). We define P(k) as the set of all the vectors that have k 1k ’s and (m− k) 0’s:
P(k) = { p = 1
k v : v ∈ {0, 1}m , m∑ i=1 vi = k } .
Clearly, P(k) has a natural connection to SGD with batch size k: for any batch sampled from the dataset, there is a corresponding p ∈ P(k), where pi = 1k if the i-th instance is in the batch. In the remainder of the paper, we assume that P is symmetric, i.e., if p ∈ P , then σ(p) ⊂ P . This follows that p0 ∈ conv(P). And we assume that P is a compact set. We assume the entries of
p ∈ P sums to 1, but we relax the restrict that each entry of p ∈ P is non-negative, i.e., we assume P ⊂ Σ. To describe the properties of the set P , we use the notation RADp(P) to denote the minimum radius of the Lp ball centered at p0 that contains P , and we use radp(P) to denote the maximum radius of the Lp ball centered at p0 whose intersection with Σ is contained within P . We denote ‖P‖p = supp∈P ‖p‖p, and ‖P‖0 as the maximum number of non-zero entries of p ∈ P . Definition 1. A function f : Θ→ R is L-Lipschitz, if:
|f(θ1)− f(θ2)| ≤ L‖θ1 − θ2‖,∀θ1, θ2 ∈ Θ. Definition 2. A function f : Θ→ R is µ-strongly convex if f(θ)− µ2 ‖θ‖ 2 is convex.
Some basic properties of the robust loss R̂S,P(θ) are summarized in Appendix A.

heading: 4 CONNECTION BETWEEN SGD AND DRO
text: SGD has long been regarded as an excellent optimizer. It is generally believed that the good performance of SGD is due to the noise introduced by stochastic gradient. But the existing results commonly view stochastic gradient as an oracle that provides an unbiased estimation of the true gradient, ignoring the fact that the noise is data-dependent. In this section, we directly analyze the good properties of mini-batch SGD from the distributionally robust perspective.
At each iteration, SGD calculates the gradient of R̂S,p(θ), where p is the weight vector corresponding to the current batch, and performs the update θ ← θ − η∇R̂S,p(θ). This can be viewed as changing the distribution over the data. Thus SGD can improve the distributional robustness of the model. The following theorem shows that changing the distribution helps mini-batch SGD escape from some bad stationary points, which is characterized by the heterogeneity of gradients of each individual instance. Theorem 1. Let S be a fixed sample of size m. For a fixed θ ∈ Θ, define a matrix G = (∇θ`(z1, θ), · · · ,∇θ`(zm, θ)), whose i-th column is the gradient of the i-th loss function w.r.t. θ. For SGD with batch size k, denote p as the random weight vector uniformly sampled from P(k). Then the expected squared update amount can be calculated by:
E p ‖∇R̂S,p(θ)‖22 = ‖∇R̂S,p0(θ)‖22 + m− k k(m− 1)
( tr(G>G)
m − ‖∇R̂S,p0(θ)‖22
) .
When tr(G >G) m ‖∇R̂S,p0(θ)‖ 2 2, the gradient of the ERM loss is relatively small, but the average squared length of the gradients g1, · · · , gm is large. This happens only if those gradients lie in some opposite directions and cancel with each other. In this case, if we leave out the instance with the largest gradient norm, then the average gradient of the rest of the sample would be almost as large as the 1m of the largest gradient norm, and the local landscape of the average loss function would shift dramatically. This indicates the model pays too much attention to this specific instance and fails to learn the intrinsic structure of the data, which can be viewed as a signal of over-fitting. We call this a bad stationary point. So the scale of tr(G
>G) m and ‖∇R̂S,p0(θ)‖ 2 2 captures the over-fitting tendency
of a model by measuring the heterogeneity of gradients of the individual instance.
Furthermore, we can know from Theorem 1 that for any fixed θ, E p ‖∇R̂S,p0(θ)‖22 decreases monotonically when k becomes larger, indicating the update amount of small batch SGD is larger than that of large batch SGD in expectation. Therefore, small batch SGD has a stronger ability to capture the non-uniformity of gradients and escape from bad stationary points.
Intuitively, if SGD converges and keeps stable when the distribution changes, the weighted average loss R̂S,p will not be too large, otherwise SGD has the tendency to escape. To theoretically prove it, we need to focus our attention to some simple case. If SGD has escaped from bad stationary points and the deep learning model really captures some intrinsic structure of the data, then it is conceivable that a batch of the instances Bt is able to reflect some common features of the entire sample S, so the local landscape of the average loss function will not shift too much when we only consider a batch of the data. In this case, we have the following theorem.
Theorem 2. Assume the loss function is L-Lipschitz. Suppose the SGD algorithm with constant learning rate η converges to and stays in B, where B = {θ : ‖θ− θ0‖2 ≤ B} is a ball that contains a minimum of ERM loss R̂S,p0 . Suppose that ∀p ∈ P(k), B contains a local minimum of R̂S,p and R̂S,p is µ-strongly convex on B. Further assume that the minimum value of the batch average loss R̂S,p in B is not too large, i.e., min
θ∈B R̂S,p(θ) ≤ min θ∈B R̂S,p0(θ) + δ. Then the robust loss R̂S,P(k)(θ)
can be bounded by:
R̂S,P(k)(θ) ≤ min θ′∈B R̂S,p0(θ ′) +
2B2 µη2 + 2LB + δ, ∀θ ∈ B.

heading: 5 THEORETICAL ANALYSIS

heading: 5.1 GENERALIZATION GUARANTEE OF DRO
text: In this subsection, we introduce the concept of Robust Rademacher Complexity and use it to derive the generalization guarantee for robust loss, which provides justifications for DRO. We bound the expectation of robust loss by empirical robust loss plus a complexity term with high probability. Then we utilize the notion of covering number to show that the generalization of robust loss is not too hard compared to that of ERM loss.
Definition 3 (Empirical Robust Rademacher Complexity). Let S = (z1, · · · , zm) be a fixed sample of size m. The empirical robust Rademacher complexity of the parameter space Θ and the weight set P w.r.t. the sample S is:
R̂S(Θ,P) = E σ [ sup θ∈Θ sup p∈P m∑ i=1 σipi`(zi, θ) ] ,
where σ = (σ1, . . . , σm)> is a random vector with independent entries uniformly from {−1, 1}. Theorem 3. For any δ > 0, with probability at least 1− δ, the following holds for all θ ∈ Θ:
E S′ R̂S′,P(θ) ≤ R̂S,P(θ) + 2R̂S(Θ,P) + 3m‖P‖∞
√ 1
2m log
2 δ .
For a fixed sample S of size m, denote Θ(S) = { (`(z1, θ), · · · , `(zm, θ))> : θ ∈ Θ } and
Θ(S,P) = { (mp1`(z1, θ), · · · ,mpm`(zm, θ))> : θ ∈ Θ,p ∈ P }
. For a norm ‖ · ‖ in Rm, denote the dual norm by ‖ · ‖∗. For > 0, let N (A, ‖ · ‖, ) denote the covering number of the set A ∈ Rm, which is the minimum cardinality of the -net w.r.t. the norm ‖ · ‖. Denote by d∗(m) an upper bound of ‖σ‖∗ = ‖(σ1, · · · , σm)>‖∗. Theorem 4. For any > 0, the empirical Rademacher complexity and the empirical robust Rademacher complexity can be respectively bounded by:
R̂S(Θ) ≤ √
2 logN (Θ(S), ‖ · ‖, ) m + d∗(m) m , (3)
and
R̂S(Θ,P) ≤ (1 + RAD2(P)) √
2 logN (Θ(S,P), ‖ · ‖, ) m + d∗(m) m . (4)
Notice that Θ(S,P) is generated from Θ(S) by having all the points perturbed a little bit, so it is conceivable that these two sets are close, implying that their covering numbers are close. Specifically, if we define
d̃ = sup x∈Θ(S,P) inf y∈Θ(S)
‖x− y‖
as the maximum distance of a point in Θ(S,P) deviating from the set Θ(S), then it follows that an -net of Θ(S) is also an ( + d̃)-net of Θ(S,P). To compare the right-hand sides of the covering
number guarantees (3) and (4), we change to + d̃ in (4), and we have:
(1 + RAD2(P))
√ 2 logN (Θ(S,P), ‖ · ‖, + d̃)
m +
( + d̃)d∗(m)
m ≤ (1 + RAD2(P)) √
2 logN (Θ(S), ‖ · ‖, ) m + d∗(m) m + d̃d∗(m) m .
When we use the Lp norm, one can easily show that d̃ ≤ mRADp(P), and ‖σ‖∗ can also be bounded as ‖σ‖∗ ≤ m 1 q = m p−1 p , so d∗(m) = m p−1 p . Then the extra term d̃d∗(m)m can be bounded by m1− 1 pRADp(P).
Those bounds quantify how the size of P will influence the generalization ability, and give a sufficient condition on the size of P to ensure that the covering number bounds are close. Conceivably, when P is sufficiently small, the generalization of robust loss is not too difficult. It is worth noticing that the inequality N (Θ(S,P), ‖ · ‖, + d̃) ≤ N (Θ(S), ‖ · ‖, ) is quite loose when the set Θ(S) is some sort of “solid,” and this is often the case in deep neural networks, whose capacities are thought to be very high. In one word, it is completely feasible to consider DRO when using DNNs.

heading: 5.2 LOCALIZED RADEMACHER COMPLEXITY
text: Next, we analyze the localized Rademacher complexity based on robust loss. Localized Rademacher complexity measures the complexity of near optimal hypotheses, which is much smaller than the entire hypothesis space. Here we assume the hypothesis set is restricted to a L2 norm ball.3
DRO can be viewed as a variance-based regularization technique (Namkoong & Duchi, 2017). A low robust loss means a low variance of the losses (`(z1, θ), · · · , `(zm, θ)), since otherwise one can choose a p ∈ P that puts more weight on the instances with a large loss, and obtains a high robust loss. A larger size of P means a more radical way to assign weights, indicating that DRO with large P renders a stricter regularization on the variance of the losses. Rademacher complexity measures the ability of a hypothesis set to fit high variance random noises, and naturally a low variance of the losses means a small Rademacher complexity. The following theorem reveals the connection between Rademacher complexity and DRO, and provides a model-free bound on the localized empirical Rademacher complexity. Theorem 5. Assume the loss function is L-Lipschitz w.r.t. θ for all z. Consider the following hypothesis set:
Θc = { θ : ESR̂S,P(θ) ≤ c, ||θ||2 < r } .
For any > 0, let N (Θc, ) denote the covering number of Θc. Then for all 0 < δ < 1, with probability ≥ 1− δ, the localized empirical Rademacher complexity can be bounded as:
R̂S(Θc) ≤ 1√ m
( 1 +
4√ mrad∞(conv(P))
)( c+ ‖P‖0‖P‖∞L+‖P‖∞ √ m
2 log N (Θc, ) δ
) .
Conceivably, whenm and c are fixed, increasing the size of P reduces the set Θc. This bound shows that increasing the size of P also reduces the complexity of Θc under certain conditions. For example, let P be the intersection of Σ and a L∞ ball with radius ρ centered at p0. When ρ is relatively small compared to 1m , increasing ρ does not change ‖P‖∞ too much, but 1 rad∞(conv(P)) is inversely proportional to ρ. Therefore, increasing ρ can reduce the scale of ‖P‖∞rad∞(conv(P)) dramatically.

heading: 6 WEIGHTED SGD
text: Motivated by the DRO and the previous theoretical analyses, we propose a new variant of SGD algorithms in Algorithm 1. We will see that this variant is simple and practical for applications.
3 When we train models with a constant weight decay λwd, generally the L2 norm of final parameters will not excess √ f(0)/λwd, where f is the loss function. This is because when the algorithm minimizing L(θ) = f(θ) + λwd‖θ‖22 outputs θ′, we generally haveL(θ′) ≤ L(0), hence, λwd‖θ′‖22 ≤ f(θ′) + λwd‖θ′‖22 ≤ f(0).
We have emphasized the significance of distributional robustness before. The critical factors that influence distributional robustness are the size of P and the value of the corresponding robust loss R̂S,P(θ). In order to obtain a more distributionally robust model, a direct approach is to increase the size of P . In the conventional SGD, the weight of the data in the current batch is 1k , and the other is 0. Due to the limited computational resources, we still only have access to one mini-batch at each iteration, but we can assign high-variance weights to the instances of the mini-batch. This implies that our algorithm enjoys the idea behind importance sampling. Therefore, our algorithm explores a wider region of empirical distributions, yielding a more distributionally robust solution.
Algorithm 1 Weighted Stochastic Gradient Descent (WSGD) Require: Initial parameter: θ0, Weight Generator: G, Learning Rate: ηt, Total Iteration: T , Train-
ing Data: S = {z1, z2, . . . , zm}. 1: for t = 1 to T do 2: Select a mini-batch: Bt = {zj1 , zj2 , . . . , zj|Bt|} 3: Generate weight: (w1, w2, . . . , w|Bt|)← G 4: Calculate normalization factor: W = ∑|Bt| i=1 wi
5: Calculate stochastic gradient: gt = 1W ∑|Bt| i=1 wi∇θ`(zji , θt−1) 6: Update parameter: θt = θt−1 − ηtgt 7: end for 8: return θT
We turn to the choices of the weight generator. In particular, we recommend two approaches:
• G1(q, r): wi = r if `(zji , θt−1) is among the q|Bt| smallest elements of the losses (`(zj1 , θt−1), · · · , `(zjBt , θt−1)), otherwise wi = 1;
• G2(q, r): Randomly select q|Bt| of the indices {ji} and set their weights to r. Set the others to 1,
where q ∈ [0, 1] is the proportion of the instances whose weight is set to r, and r is a pre-specified constant. Notice that r is not necessarily non-negative.
We have conducted experiments to show that Algorithm 1 with generator G1 achieves a better performance on classification tasks than the conventional SGD (see Section 7). We refer to Algorithm 1 with weight generator G1(q, r) as WSGD(q, r).
Different from the conventional SGD, WSGD(q, r) directly optimizes a robust loss. It is easy to prove that there exists a vector v(q, r) ∈ Σ such that the stochastic gradient in WSGD(q, r) is an unbiased estimation of the gradient of R̂S,P(v(q,r))(θ) at differentiable points (see Appendix C for details). Furthermore, the stochastic gradient in WSGD(q, r) is bounded by the Lipshitz constant of ` when r is non-negative. Taking advantages of previous work, we can give the convergence guarantee for WSGD(q, r) with a non-negative r in the strongly convex case:
Corollary 1 (Shamir & Zhang 2013). Suppose the loss function `(z, θ) is µ-strongly convex and L-Lipshitz. Consider WSGD(q, r) with step size ηt = 1µt and r is non-negative. Let R̂ ∗ S,P(v(q,r)) be
the optimal value of R̂S,P(v(q,r))(θ). Then for all T > 1, it holds that:
E[R̂S,P(v(q,r))(θT )− R̂∗S,P(v(q,r))] ≤ 17L2(1 + log(T ))
µT .
It is worth pointing out that G2 usually yields a lower test loss in our trial. Strictly speaking, this means G2 has a stronger generalization ability. However, the loss we choose in classification tasks is different from the true loss (0-1 loss), thus the algorithm with a low test loss not always has a low prediction error. We still recommendG2 because it may achieve a good performance on tasks which we can deal with the true loss directly.

heading: 7 EXPERIMENTS
text: In this section, we compare WSGD with the conventional stochastic gradient descent, which is regarded as the optimizer with the best generalization. We carry out the experiment on the CIFAR10 and CIFAR-100 datasets (Krizhevsky & Hinton, 2009), which both have 50k training data and 10k test data. On CIFAR-10, we apply our algorithm and SGD to two small networks and two large networks: ResNet-44, ResNet-56, VGG-16 and ResNet-34 (He et al., 2016; Simonyan & Zisserman, 2014). On CIFAR-100, we apply our algorithm and SGD to three large networks: VGG-16, ResNet34 and DenseNet-121 (Huang et al., 2017). We train the networks following the data augmentation in Lee et al. (2015): 4 pixels are padded on each side, and a 32 × 32 crop is randomly sampled from the padded image or its horizontal flip. We only do evaluation on the original 32 × 32 image while testing. In order to compare the generalization ability of optimization algorithms, we train the models for 600 epochs to ensure convergence. We set the initial learning rate to 0.1, and divide it by 10 at 300 and 450 epochs. The choice of hyper-parameters in WSGD(q, r) is listed in the experiment results. Empirically speaking, the hyper-parameters are mainly related to the dataset and network architecture, and a smaller r is suitable for larger networks.
We summarize the experiment results on CIFAR-10 and CIFAR-100 in Table 1. For each network and algorithm, we report the test accuracy, the L2 ball and L∞ ball robust loss. Lp ball robust loss is the robust loss with P as the intersection of Σ+ and the Lp ball centered at p0 with radius radp(Σ+). In other words, we consider the largest Lp ball in Σ centered p0 that has positive entries.
From table 1 we can see that WSGD outperforms SGD in all experiments. On CIFAR-10, WSGD decreases the misclassification rate by≈ 5‰ except VGG-16. On CIFAR-100, WSGD decreases the misclassification rate by 5‰ to 15‰ . Moreover, the models trained by WSGD have a lower L2 ball and L∞ ball robust loss than SGD. The good performance shows WSGD improves the distributional robustness of the models.

heading: 8 DISCUSSION AND CONCLUSION
text: In this paper we have theoretically analyzed the good property of DRO, and revealed the profound connection between SGD and DRO. Accordingly we have proposed a practical algorithm that can utilize the non-isotropic noise of the stochastic gradient. We have tested WSGD algorithm on CIFAR-10 and CIFAR-100, achieving significant improvements compared to SGD.
We hope this paper can inspire works that theoretically study the generalization ability of optimization algorithms. We think one should pay attention to the loss of each individual instance, rather than the average of the entire sample. We expect future works on SGD without simplification to isotropic noise, and algorithms that take advantage of the non-isotropic noise of the stochastic gradient.

heading: A SOME PROPERTIES OF ROBUST LOSS
text: Lemma 1. Let `1, · · · , `m be arbitrary real numbers. Let conv(P) denote the convex hull of P , then
sup p∈P m∑ i=1 pi`i = sup p∈conv(P) m∑ i=1 pi`i.
Lemma 2. Let `1, · · · , `m be arbitrary positive numbers. Assume P ∈ Σ and is symmetric. Then the following holds for all p ⊂ P:
| m∑ i=1 pi`i| ≤ sup p′∈P m∑ i=1 p′i`i.
Lemma 3. Assume the loss function `(z, θ) is L-Lipschitz with respect to θ for all z, then for any fixed sample S of size m, R̂S,P(θ) is ‖P‖0‖P‖∞L-Lipschitz
Proof.
|R̂S,P(θ)− R̂S,P(θ′)|
= | sup p∈P
∑ pi`(zi, θ)− sup
p∈P
∑ pi`(zi, θ ′)|
≤ max {
sup p∈P
∑ pi (` (zi, θ)− `(zi, θ′)) , sup
p∈P
∑ pi (` (zi, θ ′)− `(zi, θ)) } .
Notice that |`(zi, θ)− `(zi, θ′)| ≤ L‖θ − θ′‖ for all zi, we have: | ∑ pi (` (zi, θ)− `(zi, θ′)) |
≤ ∑ |pi|L‖θ − θ′‖
≤ ‖P‖0‖P‖∞L‖θ − θ′‖. One can develop various bounds on the Lipschitz constant of R̂S,P(θ).
Lemma 4 (Convexity). Assume that all entries of p ∈ P is non-negative. Futher assume that `(z, θ) is convex with respect to θ for all z, then for a fix sample of sizem, the robust loss R̂S,P(θ) is convex.
And denote p∗ = arg sup p∈P R̂S,p(θ), then for any gi ∈ ∂θ`(zi, θ), m∑ i=1 p∗i gi is in the subgradient of
R̂S,P(θ) at point θ.
Proof. For any p ∈ P , R̂S,p(θ) is a convex combination of convex functions, thus it is convex. R̂S,P is supremum over a family of convex functions, thus it is also convex. To complete the proof, by the definition of subgradient, we have the following holds for all i = 1, · · · ,m:
`(zi, θ ′) ≥ `(zi, θ) + 〈gi, θ′ − θ〉.
Then,
R̂S,p∗(θ ′) = m∑ i=1 p∗i `(zi, θ ′)
≥ m∑ i=1 p∗i `(zi, θ) + 〈 m∑ i=1 p∗i gi, θ ′ − θ〉.
So, R̂S,P(θ ′) ≥ R̂S,p∗(θ′)
≥ m∑ i=1 p∗i `(zi, θ) + 〈 m∑ i=1 p∗i gi, θ ′ − θ〉
= R̂S,P(θ) + 〈 m∑ i=1 p∗i gi, θ ′ − θ〉.
Lemma 5. Assume the loss function `(z, θ) is µ-strongly convex (convex), then for any P ⊂ Σ+, R̂S,P is µ-strongly convex (convex).
Lemma 6. Assume the loss function `(z, θ) takes values in [0, 1], then the robust loss R̂S,P is bounded by 1+‖P‖12

heading: B PROOF OF THE THEOREMS
text: Theorem 1. Let S be a fixed sample of size m. For a fixed θ ∈ Θ, define a matrix G = (∇θ`(z1, θ), · · · ,∇θ`(zm, θ)), whose i-th column is the gradient of the i-th loss function w.r.t. θ. For SGD with batch size k, denote p as the random weight vector uniformly sampled from P(k). Then the expected squared update amount can be calculated by:
E p ‖∇R̂S,p(θ)‖22 = ‖∇R̂S,p0(θ)‖22 + m− k k(m− 1)
( tr(G>G)
m − ‖∇R̂S,p0(θ)‖22
) .
Proof.
E p ‖∇R̂S,p(θ)‖22 = E p ‖Gp‖22
= E p ‖G(p− p0)‖22 + ‖Gp0‖22
= E p (p− p0)>G>G(p− p0) + ‖Gp0‖22
= E p
tr [ (p− p0)>G>G(p− p0) ] + ‖Gp0‖22
= E p
tr [ (p− p0)(p− p0)>G>G ] + ‖Gp0‖22
= tr [ E p (p− p0)(p− p0)>G>G ] + ‖Gp0‖22
E p (p− p0)(p− p0)> = E p
[ pp> + p0p > 0 − pp>0 − p0p> ] = E p [ pp> ] − p0p>0
One can easily calculate that
E p
[ pp> ] =  1 mk
k−1 m(m−1)k · · ·
k−1 m(m−1)k
k−1 m(m−1)k
1 mk · · ·
k−1 m(m−1)k
... . . . ... k−1
m(m−1)k · · · k−1 m(m−1)k 1 mk
 ,
when p is sampled from P(k), namely k entries of p are randomly set to 1k , while other entries are 0. And
p0p > 0 =  1 m2 1 m2 · · · 1 m2 1 m2 1 m2 · · · 1 m2 ... . . .
... 1 m2 · · · 1 m2 1 m2
 ,
Then
tr [ E p (p− p0)(p− p0)>G>G ]
=
( 1
mk − 1 m2
) tr(G>G) + ( k − 1
m(m− 1)k − 1 m2
)( 1>G>G1− tr(G>G) ) = ( 1
mk − 1 m2
)( tr(G>G)− 1 >G>G1
m ) +(
k − 1 m(m− 1)k − 1 m2
)( 1>G>G1
m − tr(G>G) ) =
m− k k(m− 1)
( tr(G>G)
m − p>0 G>Gp0
)
Theorem 2. Assume the loss function is L-Lipschitz. Suppose the SGD algorithm with constant learning rate η converges to and stays in B, where B = {θ : ‖θ− θ0‖2 ≤ B} is a ball that contains a minimum of ERM loss R̂S,p0 . Suppose that ∀p ∈ P(k), B contains a local minimum of R̂S,p and R̂S,p is µ-strongly convex on B. Further assume that the minimum value of the batch average loss R̂S,p in B is not too large, i.e., min
θ∈B R̂S,p(θ) ≤ min θ∈B R̂S,p0(θ) + δ. Then the robust loss R̂S,P(k)(θ)
can be bounded by:
R̂S,P(k)(θ) ≤ min θ′∈B R̂S,p0(θ ′) +
2B2 µη2 + 2LB + δ, ∀θ ∈ B.
Proof. We count the iteration after SGD converges to B and stays in B. We denote θt the parameter at iteration t and p(t) the weight at iteration t.
Due to the µ-strongly convexity of R̂S,p(θ),∀p ∈ P(k) we have: 1
2µ ‖∇R̂S,p(t)(θt)‖22 ≥ R̂S,p(t)(θt)− min θ′∈B R̂S,p(t)(θ ′)
≥ R̂S,p(t)(θt)− min θ′∈B R̂S,p0(θ ′)− δ.
Since the algorithm stays in B, the norm of the stochastic gradient can be bounded by 2Bη . Thus:
R̂S,p(t)(θt) ≤ min θ′∈B R̂S,p0(θ ′) +
2B2 µη2 + δ.
Plus, the weighted average loss R̂S,p(t) is L-Lipschitz, so R̂S,p(t)(θ) can be bounded by:
R̂S,p(t)(θ) ≤ min θ′∈B R̂S,p0(θ ′) +
2B2 µη2 + 2LB + δ, ∀θ ∈ B.
As SGD algorithm keeps stable in B, the above equality holds for all p(t) ∈ P(k). Finally we have:
R̂S,P(k)(θ) ≤ min θ′∈B R̂S,p0(θ ′) +
2B2 µη2 + 2LB + δ, ∀θ ∈ B.
Definition 1 (Robust Rademacher Complexity). For any m ≥ 1, the robust Rademacher complexity of the parameter space Θ and the weight set P is defined as:
Rm(Θ,P) = E S
[ R̂S(Θ,P) ] .
Theorem 3. For any δ > 0, with probability at least 1− δ, the following holds for all θ ∈ Θ:
E S′ R̂S′,P(θ) ≤ R̂S,P(θ) + 2R̂S(Θ,P) + 3m‖P‖∞
√ 1
2m log
2 δ .
Proof. For simplicity, let Rm,P(θ) = E S′
[ R̂S′,P(θ) ] , where S′ is m examples i.i.d. drawn from the
underlying distribution. For a sample S = (z1, · · · , zm) of size m, denote
Φ(S) = sup θ∈Θ
[ Rm,P(θ)− R̂S,P(θ) ] .
If we change zi to z̃i and get S̃ = (z1, · · · , z̃i, · · · , zm), then
Φ(S)− Φ(S̃) ≤ sup θ∈Θ
[ R̂S̃,P(θ)− R̂S,P(θ) ] ≤ sup θ∈Θ sup p∈P pi(`(zi, θ)− `(z̃i, θ))
≤ ‖P‖∞.
Similarly we have |Φ(S) − Φ(S̃)| ≤ ‖P‖∞. Then by McDiarmid’s inequality(McDiarmid, 1989), for any δ > 0, the following holds with probability at least 1− δ2 :
Φ(S) ≤ E S [Φ(S)] +m‖P‖∞
√ 1
2m log
2 δ .
Next we are to bound E S [Φ(S)]. By symmetrization, we have:
E S [Φ(S)] = E S sup θ∈Θ
[ Rm,P(θ)− R̂S,P(θ) ] = E
S sup θ∈Θ [ E S′ R̂S′,P(θ)− R̂S,P(θ) ] ≤ E S,S′ sup θ∈Θ [ R̂S′,P(θ)− R̂S,P(θ)
] ≤ E S,S′ sup θ∈Θ sup p∈P [ m∑ i=0 pi(`(zi, θ)− `(z′i, θ)) ]
= E σ,S,S′ sup θ∈Θ sup p∈P [ m∑ i=0 σipi(`(zi, θ)− `(z′i, θ)) ]
≤ E σ,S,S′ sup θ∈Θ sup p∈P [ m∑ i=0 σipi`(zi, θ) ] + E σ,S,S′ sup θ∈Θ sup p∈P [ m∑ i=0 −σipi`(z′i, θ) ] = 2Rm(Θ,P).
This shows that
sup θ∈Θ
[ Rm,P − R̂S,P(θ) ] ≤ 2Rm(Θ,P) +m‖P‖∞ √ 1
2m log
2
δ
holds with probability at least 1− δ2 . To complete the proof, next we show that
Rm(Θ,P) ≤ R̂S(Θ,P) +m‖P‖∞
√ 1
2m log
2
δ
holds with probability at least 1− δ2 .
If the sample S and S̃ differs only at the j-th instance, i.e., zi = z̃i for all i 6= j, then
R̂S(Θ,P)− R̂S̃(Θ,P) = Eσ supθ∈Θ sup p∈P [ m∑ i=0 σipi`(zi, θ) ] − E σ sup θ∈Θ sup p∈P [ m∑ i=0 σipi`(z̃i, θ) ] ≤ E σ sup θ∈Θ sup p∈P [σjpj(`(zj , θ)− `(z̃j , θ))]
≤ ‖P‖∞.
Similarly, we have |R̂S(Θ,P) − R̂S̃(Θ,P)| ≤ ‖P‖∞. Then the result holds immediately from McDiarmid’s inequality.
Theorem 4. For any > 0, the empirical Rademacher complexity and the empirical robust Rademacher complexity can be respectively bounded by:
R̂S(Θ) ≤ √
2 logN (Θ(S), ‖ · ‖, ) m + d∗(m) m ,
and
R̂S(Θ,P) ≤ (1 + RAD2(P)) √
2 logN (Θ(S,P), ‖ · ‖, ) m + d∗(m) m .
Proof. First, we introduce a lemma.
Lemma 7 (Massart 2000). Let A be a finite set in Rm, with r = maxx∈A ‖x‖2, then the following holds:
E σ [ sup x∈A m∑ i=1 σixi m ] ≤ r √ 2 log |A| m .
Denote N as the -net of Θ(S) where |N | = N (Θ(S), ‖ · ‖, ). Notice that `(zi, θ) ∈ [0, 1], then r ≤ √ m. Applying Massart’s lemma to N , we get:
E σ [ sup x∈N m∑ i=1 σixi m ] ≤ √ 2 logN (Θ(S), ‖ · ‖, ) m .
For any x ∈ Θ(S), there exists a x′ ∈ N , such that ‖x− x′‖ ≤ . Then m∑ i=1 σixi m = m∑ i=1 σix ′ i m + m∑ i=1 σi(xi − x′i) m
≤ m∑ i=1 σix ′ i m + d∗(m) m .
Then
R̂S(Θ) = E σ
[ sup
x∈Θ(S) m∑ i=1 σixi m
]
≤ E σ
[ sup
x∈Θ(S) m∑ i=1 σix ′ i m + d∗(m) m
]
≤ E σ [ sup x′∈N m∑ i=1 σix ′ i m ] + d∗(m) m
≤ √
2 logN (Θ(S), ‖ · ‖, ) m + d∗(m) m ,
which completes the first part of the theorem. The second part is similar, but one should notice that in this case r = max
x∈Θ(S,P) ‖x‖2 ≤
√ m‖P‖2 = √ m(1 + RAD2(P)).
Theorem 5. Assume the loss function is L-Lipschitz w.r.t. θ for all z. Consider the following hypothesis set:
Θc = { θ : ESR̂S,P(θ) ≤ c, ||θ||2 < r } .
For any > 0, let N (Θc, ) denote the covering number of Θc. Then for all 0 < δ < 1, with probability ≥ 1− δ, the localized empirical Rademacher complexity can be bounded as:
R̂S(Θc) ≤ 1√ m
( 1 +
4√ mrad∞(conv(P))
)( c+ ‖P‖0‖P‖∞L+‖P‖∞ √ m
2 log N (Θc, ) δ
) .
Proof. Notice that `(z, θ) ∈ [0, 1] for all z and θ, then if we change a single sample zi of S = (z1, · · · , zi, · · · , zm) to form S′ = (z1, · · · , z′i, · · · , zm), we have
R̂S,P(θ)− R̂S′,P(θ) = sup p∈P R̂S,p(θ)− sup p∈P R̂S′,p(θ)
≤ sup p∈P (R̂S,p(θ)− R̂S′,p(θ))
= sup p∈P
pi(`(zi, θ)− `(z′i, θ))
≤ ||P||∞.
By the same argument, we can get |R̂S,P(θ) − R̂S′,P(θ)| ≤ ||P||∞. By McDiarmid’s inequality(McDiarmid, 1989), we have:
P { R̂S,P(θ) ≥ ER̂S,P(θ) + t } ≤ exp { −2t2
m||P||2∞
} .
Let N denotes the minimum -net of Θc, where |N | = N (Θc, ). Then we have P { ∃θ0 ∈ N, R̂S,P(θ0) ≥ c+ t } ≤ N (Θc, )P { R̂S,P(θ) ≥ c+ t
} ≤ N (Θc, )P { R̂S,P(θ) ≥ ESR̂S,P(θ) + t
} ≤ N (Θc, ) exp { −2t2
m||P||2∞
} .
In other words, with probability larger than 1−N (Θc, ) exp { −2t2
m||P||2∞
} , we have
sup θ∈N R̂S,P(θ) < c+ t.
By the definition of -net, for any θ ∈ Θc, there exists a θ′ ∈ N , such that ||θ − θ′|| < . Let LS denotes an upper bound of the Lipschitz constant of R̂S(θ), then we have
sup θ∈Θc
R̂S,P(θ) < c+ t+ L S , (5)
with probability larger than 1−N (Θc, ) exp { −2t2
m||P||2∞
} .
Recall that assuming the loss function `(z, θ) is L-Lipschitz with respect to θ for all z, then the Lipschitz constant of R̂S(θ) can be bounded by ‖P‖0‖P‖∞L. For those samples S = (z1, · · · , zm) satisfying the universal bound (5), the empirical Rademacher complexity can be bounded as well. The idea of the proof involves the representation of σ as a convex combination of several vectors in P .
R̂S(Θc) = E σ [ sup θ∈Θc ∑ σi`(zi, θ) m ] ∆ = E σ [ sup θ∈Θc σ · `(S, θ) m ] .
Recall that p0 = ( 1
m , · · · , 1 m
)> and by our assumption of the symmetry of P , we have p0 ∈
conv(P). We hope that σ
m = ap0 + b(p1 − p0),
where p1 ∈ conv(P). Summing up all the coordinates of both sides of the equation, we have:∑ σi m = a1 · p0 + b (1 · p1 − 1 · p0) = a.
Define σ̄ = ∑ σi m , Then
b(p1 − p0) = (σ1 m − σ̄ m , · · · , σm m − σ̄ m )> . (6)
To derive the optimal bound, take the L∞ norm of both sides of (6), we get |b|‖p1 − p0‖∞ = max { |σ1 m − σ̄ m |, · · · , |σm m − σ̄ m | } ≤ 2 m .
Notice that ‖p1 − p0‖∞ ≥ rad∞(conv(P)), so |b| ≤ 2
mrad∞(conv(P)) .
Summarizing our result, we have: σ
m · `(S, θ) = [σ̄p0 + b(p1 − p0)] · `(S, θ)
≤ |σ̄ − b||p0 · `(S, θ)|+ |b||p1 · `(S, θ)| ≤ (|σ̄|+ 2|b|)R̂S,P(Θc).
Then
R̂S(Θc) ≤ E σ (|σ̄|+ 2|b|) R̂S,P(Θc) ≤ (√
E σ |σ̄|2 + 4
mrad∞(conv(P))
) R̂S,P(Θc)
= 1√ m
( 1 +
4√ mrad∞(conv(P))
) R̂S,P(Θc).

heading: C THE DETAILS OF WSGD
text: We denote the vector v(q, r) = (v1, v2, . . . , vm):
vi = 1
W̃ r ∑ q|Bt|>j≥0 Cjm−iC |Bt|−j−1 i−1 + ∑ |Bt|>j≥q|Bt| Cjm−iC |Bt|−j−1 i−1  where W̃ is the normalizer s.t. ∑m i=1 vi = 1. Then the stochastic gradient calculated in WSGD(q, r) is an unbiased estimation of the gradient of R̂S,P(v(q,r))(θ) at differentiable point. Recall that σ(v) is the set of all the vectors whose entries are permutations of v’s and P(v) denote the convex hull of σ(v).

Review summaries:
This paper consider the connections between SGD and distributionally robust optimization. There has long been observed a connection between robust optimization and generalization. Recently, this has been explored through the lens of distributionally robust optimization. e.g., in the papers of Namkoong and Duchi, but also many others, e.g., Farnia and Tse, etc. Primarily, this paper appears to build off the work of Namkoong.