Title: DISTRIBUTIONALLY ROBUST OPTIMIZATION LEADS TO BETTER GENERALIZATION: ON SGD AND BEYOND

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 CIFAR-100, and WSGD achieves significant improvements over the conventional SGD.

1. INTRODUCTION:

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(θ) E P (z, θ).

(1) However, the real distribution P is unknown, so we consider an empirical distribution. Given the training samples S = (z 1 , z 2 , • • • , z m ), 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

EQUATION

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 family M = {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(θ) = E P λ 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 E P λ (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 E P λ (z, θ):

1 m m i=1 P λ (z i ) P λ0 (z i ) (z i , θ).

This can be viewed as the expected loss with respect to a perturbed empirical generalized distribution2 P λ = 1 m 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 E P λ (z, θ) = sup λ∈K m i=1 P λ (z i ) mP λ0 (z i ) (z i , θ).

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 p i (z i , θ).

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 |B t | zi∈Bt ∇ θ (z i , θ t ) ,

where B t ⊂ S is the sampled mini-batch, η t is the learning rate, and |B t | 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.

2. RELATED WORK:

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( 1 m ) approximation to the generalization error R(θ), in contrast to the general O( 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.

3. NOTATION AND PRELIMINARIES:

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 p i = 1} and Σ + denote the simplex {p :

m i=1 p i = 1, p ≥ 0}. Denote p 0 = ( 1 m , . . . , 1 m ) . For p = (p 1 , p 2 , . . . , p m ) ∈ Σ, let R S,p (θ) denote m i=1 p i (z i , θ).

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 1 k 's and (m -k) 0's:

P(k) = p = 1 k v : v ∈ {0, 1} m , m i=1 v i = 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 p i = 1 k 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 p 0 ∈ 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 RAD p (P) to denote the minimum radius of the L p ball centered at p 0 that contains P, and we use rad p (P) to denote the maximum radius of the L p ball centered at p 0 whose intersection with Σ is contained within P. We denote P p = sup p∈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.

4. CONNECTION BETWEEN SGD AND DRO:

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 = (∇ θ (z 1 , θ), • • • , ∇ θ (z m , θ))

, 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 (θ) 2 2 = ∇ R S,p0 (θ) 2 2 + m -k k(m -1) tr(G G) m -∇ R S,p0 (θ) 2 2 . 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 g 1 , • • • , g m 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 1 m 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 (θ) 2 2 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 B t 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 (θ ) + 2B 2 µη 2 + 2LB + δ, ∀θ ∈ B.

5. THEORETICAL ANALYSIS:



5.1. GENERALIZATION GUARANTEE OF DRO:

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. 

R S (Θ, P) = E σ sup θ∈Θ sup p∈P m i=1 σ i p i (z i , θ) ,

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 (θ) + 2 R S (Θ, P) + 3m P ∞ 1 2m log 2 δ .

For a fixed sample S of size m, denote 

Θ(S) = ( (z 1 , θ), • • • , (z m , θ)) : θ ∈ Θ and Θ(S, P) = (mp 1 (z 1 , θ), • • • , mp m (z m , θ)) : θ ∈ Θ, p ∈ P . For a norm • in R m ,

= (σ 1 , • • • , σ m ) * .

Theorem 4. For any > 0, the empirical Rademacher complexity and the empirical robust Rademacher complexity can be respectively bounded by:

EQUATION

and

EQUATION

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)

xy 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 + RAD 2 (P)) 2 log N (Θ(S, P), • , + d) m + ( + d)d * (m) m ≤ (1 + RAD 2 (P)) 2 log N (Θ(S), • , ) m + d * (m) m + dd * (m) m .

When we use the L p norm, one can easily show that d ≤ mRAD p (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 dd * (m) m can be bounded by m 1-1 p RAD p (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.

5.2. LOCALIZED RADEMACHER COMPLEXITY:

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 L 2 norm ball. 3DRO can be viewed as a variance-based regularization technique (Namkoong & Duchi, 2017) . A low robust loss means a low variance of the losses

( (z 1 , θ), • • • , (z m , θ)

), 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 = θ : E S R 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, when m 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 p 0 . When ρ is relatively small compared to 1 m , 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.

6. WEIGHTED SGD:

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.

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 1 k , 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. Calculate stochastic gradient:

g t = 1 W |Bt| i=1 w i ∇ θ (z ji , θ t-1 ) 6:

Update parameter: θ t = θ t-1 -η t g t 7: end for 8: return θ T We turn to the choices of the weight generator. In particular, we recommend two approaches:

• G 1 (q, r): w i = r if (z ji , θ t-1 ) is among the q|B t | smallest elements of the losses ( (z j1 , θ t-1 ), • • • , (z j B t , θ t-1 )), otherwise w i = 1;

• G 2 (q, r): Randomly select q|B t | of the indices {j i } 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 G 1 achieves a better performance on classification tasks than the conventional SGD (see Section 7). We refer to Algorithm 1 with weight generator G 1 (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)) ] ≤ 17L 2 (1 + log(T )) µT .

It is worth pointing out that G 2 usually yields a lower test loss in our trial. Strictly speaking, this means G 2 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 recommend G 2 because it may achieve a good performance on tasks which we can deal with the true loss directly.

7. EXPERIMENTS:

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 CIFAR-10 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, ResNet-34 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 L 2 ball and L ∞ ball robust loss. L p ball robust loss is the robust loss with P as the intersection of Σ + and the L p ball centered at p 0 with radius rad p (Σ + ).

In other words, we consider the largest L p ball in Σ centered p 0 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 L 2 ball and L ∞ ball robust loss than SGD. The good performance shows WSGD improves the distributional robustness of the models. 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.

A SOME PROPERTIES OF ROBUST LOSS 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 p i i = sup p∈conv(P) m i=1 p i 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 p i 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 p i (z i , θ) -sup p∈P p i (z i , θ )| ≤ max sup p∈P p i ( (z i , θ) -(z i , θ )) , sup p∈P p i ( (z i , θ ) -(z i , θ)) . Notice that | (z i , θ) -(z i , θ )| ≤ L θ -θ for all z i , we have: | p i ( (z i , θ) -(z i , θ )) | ≤ |p i |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 size m, the robust loss R S,P (θ) is convex.

And denote p * = arg sup p∈P R S,p (θ), then for any

g i ∈ ∂ θ (z i , θ), m i=1 p * i g i 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:

(z i , θ ) ≥ (z i , θ) + g i , θ -θ . Then, R S,p * (θ ) = m i=1 p * i (z i , θ ) ≥ m i=1 p * i (z i , θ) + m i=1 p * i g i , θ -θ . So, R S,P (θ ) ≥ R S,p * (θ ) ≥ m i=1 p * i (z i , θ) + m i=1 p * i g i , θ -θ = R S,P (θ) + m i=1 p * i g i , θ -θ .

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 1 2 B PROOF OF THE THEOREMS Theorem 1. Let S be a fixed sample of size m. For a fixed θ ∈ Θ, define a matrix G = (∇ θ (z 1 , θ), • • • , ∇ θ (z m , θ))

, 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 (θ) 2 2 = ∇ R S,p0 (θ) 2 2 + m -k k(m -1) tr(G G) m -∇ R S,p0 (θ) 2 2 .

Proof.

E p ∇ R S,p (θ) 2 2 = E p Gp 2 2 = E p G(p -p 0 ) 2 2 + Gp 0 2 2 = E p (p -p 0 ) G G(p -p 0 ) + Gp 0 2 2 = E p tr (p -p 0 ) G G(p -p 0 ) + Gp 0 2 2 = E p tr (p -p 0 )(p -p 0 ) G G + Gp 0 2 2 = tr E p (p -p 0 )(p -p 0 ) G G + Gp 0 2 2 E p (p -p 0 )(p -p 0 ) = E p pp + p 0 p 0 -pp 0 -p 0 p = E p pp -p 0 p 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 1 k , while other entries are 0. And 

p 0 p 0 =      1 m 2 1 m 2 • • • 1 m 2 1 m 2 1 m 2 • • • 1 m 2 . . . . . . . . . 1 m 2 • • • 1 m 2 1 m 2      , Then tr E p (p -p 0 )(p -p 0 ) G G = 1 mk - 1 m 2 tr(G G) + k -1 m(m -1)k - 1 m 2 1 G G1 -tr(G G) = 1 mk - 1 m 2 tr(G G) - 1 G G1 m + k -1 m(m -1)k - 1 m 2 1 G G1 m -tr(G G) = m -k k(m -1) tr(G G) m -p 0 G Gp 0 Theorem 2.

(θ) ≤ min θ ∈B R S,p0 (θ ) + 2B 2 µη 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 ) 2 2 ≥ 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 (θ ) + 2B 2 µη 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 (θ ) + 2B 2 µη 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 (θ ) + 2B 2 µη 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: R m (Θ, P) = E S R S (Θ, P) .

Theorem 3. For any δ > 0, with probability at least 1 -δ, the following holds for all θ ∈ Θ: If we change z i to z i and get

E S R S ,P (θ) ≤ R S,P (θ) + 2 R S (Θ, P) + 3m P ∞ 1 2m log 2 δ .

S = (z 1 , • • • , z i , • • • , z m ), then Φ(S) -Φ( S) ≤ sup θ∈Θ R S,P (θ) -R S,P (θ) ≤ sup θ∈Θ sup p∈P p i ( (z i , θ) -( 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 This shows that

sup θ∈Θ R m,P -R S,P (θ) ≤ 2R m (Θ, P) + m P ∞ 1 2m log 2 δ

holds with probability at least 1 -δ 2 . To complete the proof, next we show that

R m (Θ, 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., z i = z i for all i = j, then

R S (Θ, P) -R S (Θ, P) = E σ sup θ∈Θ sup p∈P m i=0 σ i p i (z i , θ) -E σ sup θ∈Θ sup p∈P m i=0 σ i p i ( z i , θ) ≤ E σ sup θ∈Θ sup p∈P [σ j p j ( (z j , θ) -( z j , θ))] ≤ P ∞ .

Similarly, we have | R S (Θ, P) -R S (Θ, P)| ≤ P ∞ . Then the result holds immediately from McDiarmid's inequality. Proof. First, we introduce a lemma.

Lemma 7 (Massart 2000) . Let A be a finite set in R m , with r = max x∈A x 2 , then the following holds:

E σ sup x∈A m i=1 σ i x i m ≤ r 2 log |A| m .

Denote N as the -net of Θ(S) where |N | = N (Θ(S), • , ). Notice that (z i , θ) ∈ [0, 1], then r ≤ √ m. Applying Massart's lemma to N , we get:

E σ sup x∈N m i=1 σ i x i m ≤ 2 log N (Θ(S), • , ) m .

For any x ∈ Θ(S), there exists a x ∈ N , such that xx ≤ . Then

m i=1 σ i x i m = m i=1 σ i x i m + m i=1 σ i (x i -x i ) m ≤ m i=1 σ i x i m + d * (m) m . Then R S (Θ) = E σ sup x∈Θ(S) m i=1 σ i x i m ≤ E σ sup x∈Θ(S) m i=1 σ i x i m + d * (m) m ≤ E σ sup x ∈N m i=1 σ i x i m + d * (m) m ≤ 2 log N (Θ(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 + RAD 2 (P)).

Theorem 5. Assume the loss function is L-Lipschitz w.r.t. θ for all z. Consider the following hypothesis set:

Θ c = θ : E S R 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: where p 1 ∈ conv(P). Summing up all the coordinates of both sides of the equation, we have:

R S (Θ c ) ≤ 1 √ m 1 + 4 √ mrad ∞ (

EQUATION

To derive the optimal bound, take the L ∞ norm of both sides of (6), we get

|b| p 1 -p 0 ∞ = max | σ 1 m - σ m |, • • • , | σ m m - σ m | ≤ 2 m .

Notice that p 1 -p 0 ∞ ≥ rad ∞ (conv(P)), so |b| ≤ 2 mrad ∞ (conv(P))

.

Summarizing our result, we have: 

C THE DETAILS OF WSGD:

We denote the vector v(q, r) = (v 1 , v 2 , . . . , v m ):

v i = 1 W   r q|Bt|>j≥0 C j m-i C |Bt|-j-1 i-1 + |Bt|>j≥q|Bt| C j m-i C |Bt|-j-1 i-1  

where W is the normalizer s.t. m i=1 v i = 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).

:

To make the fraction reasonable, we assume P λ 0 (z) > 0 whenever P λ (z) > 0.

Strictly 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.

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 θ 2 2 outputs θ , we generally haveL(θ ) ≤ L(0), hence, λ wd θ 2 2 ≤ f (θ ) + λ wd θ 2 2 ≤ f (0).

Figures and Tables:

Weighted Stochastic Gradient Descent (WSGD) Require: Initial parameter: θ 0 , Weight Generator: G, Learning Rate: η t , Total Iteration: T , Training Data: S = {z 1 , z 2 , . . . , z m }. 1: for t = 1 to T do 2: Select a mini-batch: B t = {z j1 , z j2 , . . . , z j |B t | } 3: Generate weight: (w 1 , w 2 , . . . , w |Bt| ) ← G

For simplicity, let R m,P (θ) = E S R S ,P (θ) , where S is m examples i.i.d. drawn from the underlying distribution. For a sample S = (z 1 , • • • , z m ) of size m, denote Φ(S) = sup θ∈Θ R m,P (θ) -R S,P (θ) .

i ( (z i , θ) -(z i , θ)) p i (z i , θ) = 2R m (Θ, P).

For any > 0, the empirical Rademacher complexity and the empirical robust Rademacher complexity can be respectively bounded by: R S (Θ) ≤ 2 log N (Θ(S), • , )

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 samplez i of S = (z 1 , • • • , z i , • • • , z m ) to form S = (z 1 , • • • , z i , • • • , z m ), we have R S,P (θ) -R S ,P ,p (θ) -R S ,p (θ)) = sup p∈P p i ( (z i , θ) -(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 (θ) ≥ E R S,P (θ) + t ≤ exp -2t 2 m||P|| 2 ∞ .Let N denotes the minimum -net of Θ c , where |N | = N (Θ c , ). Then we haveP ∃θ 0 ∈ N, R S,P (θ 0 ) ≥ c + t ≤ N (Θ c , )P R S,P (θ) ≥ c + ≤ N (Θ c , )P R S,P (θ) ≥ E S R S,P (θ) + t ≤ N (Θ c , ) exp -2t 2 m||P|| 2 ∞In other words, with probability larger than 1 -N (Θ c , P (θ) < c + t.By the definition of -net, for any θ ∈ Θ c , there exists a θ ∈ N , such that ||θ -θ || < . Let L S denotes an upper bound of the Lipschitz constant of R S (θ), then we havesup θ∈Θc R S,P (θ) < c + t + L S ,(5)with probability larger than1 -N (Θ c , ) exp -2t 2 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 = (z 1 , • • • , z m ) 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 Recall that p 0 = 1 m , • • • , 1 mand by our assumption of the symmetry of P, we have p 0 ∈ conv(P). We hope that σ m = ap 0 + b(p 1 -p 0 ),

σ m • (S, θ) = [σp 0 + b(p 1 -p 0 )] • (S, θ) ≤ |σ -b||p 0 • (S, θ)| + |b||p 1 • (S, θ)| ≤ (|σ| + 2|b|) R S,P (Θ c ). Then R S (Θ c ) ≤ E σ (|σ| + 2|b|) R S,P (Θ c )

Top 1 accuracy(%) and robust loss(×10 -3 ) on CIFAR 10 and CIFAR 100

Table content: <table><tr><td>Network</td><td colspan="2">Params Acc</td><td colspan="3">L 2 -ball L ∞ -ball Method</td><td>Dataset</td></tr><tr><td>ResNet-44</td><td>0.66M</td><td colspan="2">93.60 0.5902</td><td>1.116</td><td>SGD</td><td>CIFAR-10</td></tr><tr><td>ResNet-44</td><td>0.66M</td><td colspan="2">94.24 0.4892</td><td>0.9507</td><td>WSGD(0.5,0.3)</td><td>CIFAR-10</td></tr><tr><td>ResNet-56</td><td>0.86M</td><td colspan="2">93.86 0.3762</td><td>0.6745</td><td>SGD</td><td>CIFAR-10</td></tr><tr><td>ResNet-56</td><td>0.86M</td><td colspan="2">94.33 0.3070</td><td>0.5835</td><td>WSGD(0.5,0.4)</td><td>CIFAR-10</td></tr><tr><td>VGG-16</td><td colspan="3">14.74M 94.05 0.9617</td><td>1.112</td><td>SGD</td><td>CIFAR-10</td></tr><tr><td>VGG-16</td><td colspan="3">14.74M 94.30 0.8807</td><td>1.004</td><td>WSGD(0.5,0.2)</td><td>CIFAR-10</td></tr><tr><td>ResNet-34</td><td colspan="3">21.30M 95.13 0.7113</td><td>0.9379</td><td>SGD</td><td>CIFAR-10</td></tr><tr><td>ResNet-34</td><td colspan="3">21.30M 95.65 0.6422</td><td>0.8381</td><td>WSGD(0.5,0)</td><td>CIFAR-10</td></tr><tr><td>VGG-16</td><td colspan="3">14.78M 72.13 1.535</td><td>2.942</td><td>SGD</td><td>CIFAR-100</td></tr><tr><td>VGG-16</td><td colspan="3">14.78M 73.60 1.209</td><td>2.346</td><td>WSGD(0.5,0)</td><td>CIFAR-100</td></tr><tr><td>ResNet-34</td><td colspan="3">21.35M 77.89 1.473</td><td>2.562</td><td>SGD</td><td>CIFAR-100</td></tr><tr><td>ResNet-34</td><td colspan="3">21.35M 78.43 0.9876</td><td>1.871</td><td>WSGD(0.5,0)</td><td>CIFAR-100</td></tr><tr><td colspan="2">DenseNet-121 7.13M</td><td colspan="2">79.31 3.359</td><td>4.614</td><td>SGD</td><td>CIAFR-100</td></tr><tr><td colspan="2">DenseNet-121 7.13M</td><td colspan="2">80.35 1.196</td><td>2.152</td><td colspan="2">WSGD(0.5,-0.1) CIAFR-100</td></tr><tr><td colspan="4">8 DISCUSSION AND CONCLUSION</td><td/><td/></tr></table>

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

Table content: <table><tr><td>θ∈B</td><td>R S,p (θ) ≤ min θ∈B</td><td>R S,p0 (θ) + δ. Then the robust loss R S,P(k) (θ)</td></tr><tr><td>can be bounded by:</td><td/><td/></tr><tr><td>R S,P(k)</td><td/><td/></tr></table>

Peer review (iclr2019):
 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. The key connection this paper tries to make is between SGD and DRO, since SGD in sampling a minibatch, can be considered a small perturbation to the distribution. Therefore the authors use this intuition to propose a weighted version of SGD (WSGD) whereby high variance weights are assigned to mini batch, thus making the training accomplish a higher level of distributional robustness. This idea is tested on a few data sets including CIFAR-10 and -100. The results compare WSGD with SGD, and they show that the WSGD-trained models have a lower robust loss, and also have a higher (testing) accuracy. This is an interesting paper. There has been much discussion of the role of batch size, and considering it from a different perspective seems to be of interest. But the connection of the empirical results to the theoretical results seems tenuous. It’s not clear how predictions of the theory match up. This would be useful to understand better. More generally, a simpler presentation of the key results would be useful, so as to allow the reader to better appreciate what are the main claims and if they are as substantial as claimed. Overall the writing needs significant polishing, though this is only at a local level, i.e, it doesn’t obscure the flow of the paper.