D ISTRIBUTIONALLY R OBUST O PTIMIZATION L EADS TO B ETTER G ENERALIZATION : ON SGD AND B EYOND
Anonymous authors
Paper under double-blind review
A BSTRACT
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 signiﬁcant improvements over the conventional SGD.
I NTRODUCTION
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 (cid:96) ( 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 ﬁxed but unknown distribution over Z (denoted P ).
We attempt to ﬁnd a θ ∈ Θ , which gives a small generalization error.
That is, we address the following optimization problem:
[Equation]
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 deﬁned as (cid:98) P m ( z )
=m (cid:80) mi
=
I
[ z = z i ] , 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-ﬁtting.
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 ﬁxed 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 (cid:96) ( 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 deﬁned as:
[Equation]
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 unbiased 1 estimation of the expected loss E P λ (cid:96) ( z, θ ) :
[Equation]
This can be viewed as the expected loss with respect to a perturbed empirical generalized distribution 2 (cid:98)
[Equation]
I [ z = z i ] , so the DRO problem becomes minimizing the empirical distributionally robust loss:
[Equation]
Keeping the motivations above in mind, we focus our attention on the robust loss , deﬁned as the maximal possible weighted loss over a pre-speciﬁed weight set P for simplicity:
[Equation]
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:
[Equation]
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 beneﬁcial 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.
R ELATED W ORK
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 difﬁculty 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 beneﬁt 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 deﬁnitions of sharpness and ﬂatness 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.
N OTATION AND P RELIMINARIES
In our work, we assume the loss function (cid:96) ( z, θ ) is bounded between 0 and 1 , and is continuously differentiable.
Let Σ denote the hyperplane { p : (cid:80) mi =1 p i = 1 } and Σ + denote the simplex { p : (cid:80) m i =1 p
i
= 1 , p ≥ 0 } .
Denote p
= ( 1 m , . .
, 1 m ) (cid:62) .
For p = ( p 1 , p 2 , . . .
, p m ) (cid:62) ∈ Σ , let (cid:98) R S, p ( θ ) denote (cid:80) mi =1 p i (cid:96) ( z i , θ ) .
Then the ERM loss can be written as (cid:98) R S, p 0 ( θ ) , and the robust loss w.r.t.
the weight set P can be viewed as:
[Equation]
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 deﬁne P ( k ) as the set of all the vectors that have k 1 k ’s and ( m − k ) 0’s:
[Equation]
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 (cid:107)P(cid:107) p = sup p ∈P (cid:107) p (cid:107) p , and (cid:107)P(cid:107)
as the maximum number of non-zero entries of p ∈ P .
Deﬁnition 1.
A function f : Θ → R is L -Lipschitz, if:
[Equation]
Deﬁnition 2.
A function f : Θ → R is µ -strongly convex if f ( θ ) − µ 2 (cid:107) θ (cid:107) 2 is convex.
Some basic properties of the robust loss (cid:98) R S, P ( θ ) are summarized in Appendix A.
C ONNECTION 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 (cid:98) R S, p ( θ ) , where p is the weight vector corresponding to the current batch, and performs the update θ ← θ − η ∇ (cid:98) 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 ﬁxed sample of size m .
For a ﬁxed θ ∈ Θ , deﬁne a matrix G = ( ∇ θ (cid:96) ( z 1 , θ ) , · · · , ∇ θ (cid:96) ( 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:
[Equation]
When tr( G (cid:62) G ) m
(cid:29) (cid:107)∇ (cid:98) R S, p 0 ( θ ) (cid:107) 22 , 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 speciﬁc instance and fails to learn the intrinsic structure of the data, which can be viewed as a signal of over-ﬁtting.
We call this a bad stationary point.
So the scale of tr ( G (cid:62) G ) m and (cid:107)∇ (cid:98) R S, p 0 ( θ ) (cid:107) 22 captures the over-ﬁtting tendency of a model by measuring the heterogeneity of gradients of the individual instance.
Furthermore, we can know from Theorem 1 that for any ﬁxed θ , E p (cid:107)∇ (cid:98) R S, p 0 ( θ ) (cid:107) 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 (cid:98) 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 reﬂect 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 = { θ : (cid:107) θ − θ 0 (cid:107) 2 ≤ B } is a ball that contains a minimum of ERM loss (cid:98) R S, p 0 .
Suppose that ∀ p ∈ P ( k ) , B contains a local minimum of (cid:98) R S, p and (cid:98) R S, p is µ -strongly convex on B .
Further assume that the minimum value of the batch average loss (cid:98) R S, p in B is not too large, i.e., min θ ∈B (cid:98) R S, p ( θ ) ≤ min θ ∈B (cid:98) R S, p 0 ( θ ) + δ .
Then the robust loss (cid:98) R S, P ( k ) ( θ ) can be bounded by:
[Equation]
T HEORETICAL A NALYSIS 5. G ENERALIZATION G UARANTEE 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 justiﬁcations 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.
Deﬁnition 3 (Empirical Robust Rademacher Complexity) .
Let S =
( z 1 , · · · , z m ) be a ﬁxed 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:
[Equation]
where σ = ( σ 1 , . .
, σ m ) (cid:62) 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 θ ∈ Θ :
[Equation]
For a ﬁxed sample S of size m , denote Θ( S ) =
(cid:110) ( (cid:96) ( z 1 , θ ) , · · · , (cid:96) ( z m , θ ))
(cid:62) : θ ∈ Θ (cid:111) and Θ( S, P ) =
(cid:110) ( mp 1 (cid:96) ( z 1 , θ ) , · · · , mp m (cid:96) ( z m , θ ))
(cid:62) : θ ∈ Θ , p ∈ P (cid:111) .
For a norm (cid:107) · (cid:107) in R m , denote the dual norm by (cid:107) · (cid:107) ∗ .
For (cid:15) > 0 , let N ( A, (cid:107) · (cid:107) , (cid:15) ) denote the covering number of the set
A ∈ R m , which is the minimum cardinality of the (cid:15) -net w.r.t.
the norm (cid:107) · (cid:107) .
Denote by d ∗ ( m ) an upper bound of (cid:107) σ (cid:107) ∗ = (cid:107) ( σ 1 , · · · , σ m ) (cid:62) (cid:107) ∗ .
Theorem 4.
For any (cid:15) > 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.
Speciﬁcally, if we deﬁne
[Equation]
as the maximum distance of a point in Θ( S, P ) deviating from the set Θ( S ) , then it follows that an (cid:15) -net of Θ( S ) is also an ( (cid:15) + (cid:101) d ) -net of Θ( S, P ) .
To compare the right-hand sides of the covering number guarantees (3) and (4), we change (cid:15) to (cid:15) + (cid:101) d in (4), and we have:
[Equation]
When we use the L p norm, one can easily show that (cid:101) d ≤ m RAD p
( P ) , and (cid:107) σ (cid:107) ∗ can also be bounded as (cid:107) σ (cid:107) ∗ ≤ m
[Equation]
q = m p
[Equation]
1 p , so d ∗ ( m ) = m p − 1 p .
Then the extra term (cid:101) dd ∗ ( m ) m can be bounded by m 1 − 1 p RAD p
( P ) .
Those bounds quantify how the size of P will inﬂuence the generalization ability, and give a sufﬁcient condition on the size of P to ensure that the covering number bounds are close.
Conceivably, when P is sufﬁciently small, the generalization of robust loss is not too difﬁcult.
It is worth noticing that the inequality N (Θ( S, P ) , (cid:107) · (cid:107) , (cid:15) + (cid:101) d ) ≤ N (Θ( S ) , (cid:107) · (cid:107) , (cid:15) ) 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. L OCALIZED R ADEMACHER C OMPLEXITY
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.
DRO can be viewed as a variance-based regularization technique (Namkoong & Duchi, 2017).
A low robust loss means a low variance of the losses ( (cid:96) ( z 1 , θ ) , · · · , (cid:96) ( 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 ﬁt 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:
[Equation]
For any (cid:15) > 0 , let N (Θ c , (cid:15) ) denote the covering number of Θ c .
Then for all 0 < δ < 1 , with probability ≥ 1 − δ , the localized empirical Rademacher complexity can be bounded as:
[Equation]
0 (cid:107)P(cid:107) ∞ L + (cid:107)P(cid:107) ∞ (cid:114) m 2 log
N (Θ c , (cid:15) ) δ (cid:33) .
Conceivably, when m and c are ﬁxed, 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 (cid:107)P(cid:107) ∞ too much, but 1 rad ∞ ( conv ( P )) is inversely proportional to ρ .
Therefore, increasing ρ can reduce the scale of [Equation] )) dramatically.
W EIGHTED 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 signiﬁcance of distributional robustness before.
The critical factors that inﬂuence distributional robustness are the size of P and the value of the corresponding robust loss (cid:98) 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.
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 j 1 , z j 2 , . . .
z j | Bt | } 3:
Generate weight: ( w 1 , w 2 , . .
, w | B t | ) ← G 4:
Calculate normalization factor: W = (cid:80) | B t | i =1 w i 5:
Calculate stochastic gradient: g t = 1 W (cid:80) |
Bt | i =1 w i ∇ θ (cid:96) ( z j i , θ 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 (cid:96) ( z j i , θ t − 1 ) is among the q | B t | smallest elements of the losses ( (cid:96) ( z j 1 , θ t − 1 ) , · · · , (cid:96) ( z j Bt , θ 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-speciﬁed 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 classiﬁcation 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 (cid:98) 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 (cid:96) 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 (cid:96) ( z, θ ) is µ -strongly convex and
L -Lipshitz.
Consider WSGD ( q, r ) with step size η t = 1 µt and r is non-negative.
Let (cid:98) R ∗ S, P ( v ( q,r )) be the optimal value of (cid:98) R S, P ( v ( q,r ))
( θ ) .
Then for all T > 1 , it holds that:
[Equation]
µ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 classiﬁcation 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.
E XPERIMENTS
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 ﬂip.
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 misclassiﬁcation rate by ≈ 5‰ except VGG-16.
On CIFAR-100, WSGD decreases the misclassiﬁcation 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.
[Table]
D ISCUSSION AND C ONCLUSION
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 signiﬁcant 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 simpliﬁcation to isotropic noise, and algorithms that take advantage of the non-isotropic noise of the stochastic gradient.
A S OME P ROPERTIES OF R OBUST LOSS
Lemma 1.
Let (cid:96) 1 , · · · , (cid:96) m be arbitrary real numbers.
Let conv( P ) denote the convex hull of P , then
[Equation]
Lemma 2.
Let (cid:96) 1 , · · · , (cid:96) m be arbitrary positive numbers.
Assume P ∈ Σ and is symmetric.
Then the following holds for all p
⊂ P :
[Equation]
Lemma 3.
Assume the loss function (cid:96) ( z, θ ) is L -Lipschitz with respect to θ for all z , then for any ﬁxed sample S of size m , (cid:98) R S, P ( θ ) is (cid:107)P(cid:107) 0 (cid:107)P(cid:107)
∞ L -Lipschitz
Proof.
[Equation]
Notice that | (cid:96) ( z i , θ ) − (cid:96) ( z i , θ (cid:48) ) | ≤ L (cid:107) θ − θ (cid:48) [Equation] for all z i , we have:
[Equation]
One can develop various bounds on the Lipschitz constant of (cid:98) R S, P ( θ ) .
Lemma 4 (Convexity) .
Assume that all entries of p ∈ P is non-negative.
Futher assume that (cid:96) ( z, θ ) is convex with respect to θ for all z , then for a ﬁx sample of size m , the robust loss (cid:98) R S, P ( θ ) is convex.
And denote p ∗ = arg sup p ∈P (cid:98) R S, p ( θ ) , then for any g i ∈
∂ θ (cid:96) ( z i , θ ) , m (cid:88) i = p ∗
i g i is in the subgradient of (cid:98) R S, P ( θ ) at point θ .
Proof.
For any p ∈ P , (cid:98) R S, p ( θ ) is a convex combination of convex functions, thus it is convex.
(cid:98) R S, P is supremum over a family of convex functions, thus it is also convex.
To complete the proof, by the deﬁnition of subgradient, we have the following holds for all i = 1 , · · · , m :
[Equation]
− θ (cid:105) .
Then, (cid:98) R S, p ∗ ( θ (cid:48) )
= m (cid:88) i
= p ∗ i (cid:96) ( z i , θ (cid:48) )
≥ m (cid:88) i =
p ∗ i (cid:96) ( z i , θ ) + (cid:104) m (cid:88) i
= p ∗ i g i , θ (cid:48) − θ (cid:105) .
So, (cid:98) R S, P ( θ (cid:48) ) ≥ (cid:98) R S, p ∗ ( θ (cid:48) )
≥ m (cid:88) i =
p ∗ i (cid:96) ( z i , θ ) + (cid:104) m (cid:88) i
= p ∗ i g i , θ (cid:48) − θ (cid:105) =
(cid:98) R S, P ( θ ) + (cid:104) m (cid:88)
i
= p ∗ i g i , θ (cid:48) − θ (cid:105) .
Lemma 5.
Assume the loss function (cid:96) ( z, θ ) is µ -strongly convex (convex), then for any P ⊂ Σ + , (cid:98) R S, P is µ -strongly convex (convex).
Lemma 6.
Assume the loss function (cid:96) ( z, θ ) takes values in [0 , 1] , then the robust loss (cid:98) R S, P is bounded by 1+ (cid:107)P(cid:107)
B P ROOF OF THE T HEOREMS
Theorem 1.
Let S be a ﬁxed sample of size m .
For a ﬁxed θ ∈ Θ , deﬁne a matrix G = ( ∇ θ (cid:96) ( z 1 , θ ) , · · · , ∇ θ (cid:96) ( 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:
[Equation]
Proof.
[Equation]
One can easily calculate that
[Equation]

[Equation]
when p is sampled from P ( k )
[Equation]
namely k entries of p are randomly set to 1 k
[Equation]
while other entries are 0 .
And
[Equation]
Then
[Equation]
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 = { θ : (cid:107) θ − θ 0 (cid:107) 2 ≤ B } is a ball that contains a minimum of ERM loss (cid:98) R S, p 0 .
Suppose that ∀ p ∈ P ( k ) , B contains a local minimum of (cid:98) R S, p and (cid:98) R S, p is µ -strongly convex on B .
Further assume that the minimum value of the batch average loss (cid:98) R S, p in B is not too large, i.e., min θ ∈B (cid:98) R S, p ( θ ) ≤ min θ ∈B (cid:98) R S, p 0 ( θ ) + δ .
Then the robust loss (cid:98) R S, P
[Equation]
k ) ( θ ) can be bounded by:
[Equation]
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 (cid:98) R S, p ( θ ) , ∀ p ∈ P ( k ) we have:
[Equation]
Since the algorithm stays in B , the norm of the stochastic gradient can be bounded by 2 Bη .
Thus:
[Equation]
Plus, the weighted average loss (cid:98) R S, p ( t ) is L - Lipschitz , so (cid:98) R S, p ( t ) ( θ ) can be bounded by:
[Equation]
As SGD algorithm keeps stable in B , the above equality holds for all p ( t )
∈ P ( k ) .
Finally we have:
[Equation]
Deﬁnition 1 (Robust Rademacher Complexity) .
For any m ≥ 1 , the robust Rademacher complexity of the parameter space Θ and the weight set P is deﬁned as:
[Equation]
Theorem 3.
For any δ >
[Equation]
with probability at least 1 − δ
[Equation]
the following holds for all θ ∈ Θ :
[Equation]
δ .
Proof.
For simplicity, let R m, P ( θ ) =
E S (cid:48) (cid:104) (cid:98) R S (cid:48) , P ( θ ) (cid:105) , where S (cid:48) is m examples i.i.d.
drawn from the underlying distribution.
For a sample S
= ( z 1 , · · · , z m ) of size m , denote
[Equation]
If we change z i to (cid:101) z i and get (cid:101) S
= ( z 1 , · · · , (cid:101)
z i , · · · , z m ) , then
[Equation]
∞ .
Similarly we have | Φ( S ) − Φ( (cid:101) S ) | ≤ (cid:107)P(cid:107) ∞ .
Then by McDiarmid’s inequality(McDiarmid, 1989), for any δ > 0 , the following holds with probability at least 1 − δ 2 :
[Equation]
Next we are to bound E S
[Φ( S )] .
By symmetrization, we have:
[Equation]
This shows that
[Equation]
holds with probability at least 1 − δ 2 .
To complete the proof
[Equation]
next we show that
[Equation]
δ holds with probability at least 1 − δ 2 .
If the sample S and (cid:101) S differs only at the j -th instance, i.e., z i = (cid:101)
z i for all i (cid:54) =
j , then
[Equation]
≤ (cid:107)P(cid:107) ∞ .
Similarly, we have | (cid:98) R S (Θ , P ) − (cid:98) R (cid:101) S (Θ , P ) | ≤ (cid:107)P(cid:107) ∞ .
Then the result holds immediately from
McDiarmid’s inequality.
Theorem 4.
For any (cid:15) > 0 , the empirical Rademacher complexity and the empirical robust
Rademacher complexity can be respectively bounded by:
[Equation]
and
[Equation]
Proof.
First, we introduce a lemma.
Lemma 7 (Massart 2000) .
Let A be a ﬁnite set in R m , with r = max x ∈ A (cid:107) x (cid:107) 2 , then the following holds:
[Equation]
m
[Equation]
Denote N as the (cid:15) -net of Θ( S ) where | N | =
N (Θ( S ) , (cid:107) · (cid:107) , (cid:15) ) .
Notice that (cid:96) ( z i , θ ) ∈ [0 , 1] , then r ≤ √ m .
Applying Massart’s lemma to N , we get:
[Equation]
For any x ∈ Θ( S ) , there exists a x (cid:48) ∈ N , such that (cid:107) x − x (cid:48) (cid:107) ≤ (cid:15) .
Then m (cid:88) i
= σ i x i m = m (cid:88)
i =
σ i x (cid:48)
i m + m (cid:88)
i =
σ i ( x i − x (cid:48) i ) m ≤ m (cid:88) i
=
σ i x (cid:48)
i m + (cid:15)d ∗ ( m ) m .
Then
[Equation]
which completes the ﬁrst part of the theorem.
The second part is similar, but one should notice that in this case r =
[Equation]
Theorem 5.
Assume the loss function is L -Lipschitz w.r.t.
θ for all z .
Consider the following hypothesis set:
[Equation]
For any (cid:15) > 0 , let N (Θ c , (cid:15) ) denote the covering number of Θ c .
Then for all 0 < δ < 1 , with probability ≥ 1 − δ , the localized empirical Rademacher complexity can be bounded as:
[Equation]
Proof.
Notice that (cid:96) ( z, θ ) ∈ [0 , 1] for all z and θ , then if we change a single sample
z i of S =
( z 1 , · · · , z i , · · · , z m ) to form S (cid:48)
= ( z 1 , · · · , z (cid:48)
i , · · · , z m ) , we have
[Equation]
By the same argument, we can get | (cid:98) R S, P ( θ ) − (cid:98)
[Equation]
P ( θ ) | ≤ ||P|| ∞ .
By McDiarmid’s inequality(McDiarmid, 1989), we have:
[Equation]
Let N denotes the minimum (cid:15) -net of Θ c , where | N | =
N (Θ c , (cid:15) ) .
Then we have
[Equation]
In other words, with probability larger than 1 − N
(Θ c , (cid:15) ) exp (cid:26) − 2 tm ||P|| 2 ∞ (cid:27) , we have sup
[Equation]
, P
[Equation]
By the deﬁnition of (cid:15) -net, for any θ ∈ Θ c , there exists a θ (cid:48) ∈ N , such that || θ − θ (cid:48)
|| < (cid:15) .
Let L S denotes an upper bound of the Lipschitz constant of (cid:98) R S ( θ ) , then we have
[Equation]
with probability larger than 1 − N (Θ c , (cid:15) ) exp (cid:26) −
[Equation] [Equation]
Recall that assuming the loss function (cid:96) ( z, θ ) is L - Lipschitz with respect to θ for all z , then the
Lipschitz constant of (cid:98) R S ( θ ) can be bounded by (cid:107)P(cid:107) 0 (cid:107)P(cid:107)
∞ 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 .
[Equation]
Recall that p
= (cid:18)m, · · · , 1 m (cid:19) (cid:62) and by our assumption of the symmetry of P , we have p 0 ∈ conv( P ) .
We hope that
[Equation]
where p 1 ∈ conv( P )
Summing up all the coordinates of both sides of the equation, we have:
[Equation]
Deﬁne ¯ σ = (cid:80) σ i m , Then
[Equation]
To derive the optimal bound, take the L ∞ norm of both sides of (6), we get
[Equation]
Notice that (cid:107) p 1 − p 0 (cid:107) ∞ ≥ rad ∞ (conv( P )) , so | b | ≤ m rad ∞ (conv( P )) .
Summarizing our result, we have:
[Equation]
Then
[Equation]
C T HE D ETAILS OF WSGD
We denote the vector v ( q, r ) =
( v 1 , v 2 , . .
. , v m ) :
[Equation]

 where (cid:102) W is the normalizer s.t.
(cid:80) mi
=1 v i = 1 .
Then the stochastic gradient calculated in WSGD ( q, r ) is an unbiased estimation of the gradient of (cid:98) 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 ) .