SGD WITH H ARDNESS W EIGHTED S AMPLING FOR D ISTRIBUTIONALLY R OBUST D EEP L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Distributionally Robust Optimization (DRO) has been proposed as an alternative to Empirical Risk Minimization (ERM) in order to account for potential biases in the training data distribution.
However, its use in deep learning has been severely restricted due to the relative inefﬁciency of the optimizers available for DRO in comparison to the wide-spread Stochastic Gradient Descent (SGD) based optimizers for deep learning with ERM.
In this work, we demonstrate that SGD with hardness weighted sampling is a principled and efﬁcient optimization method for
DRO in machine learning and is particularly suited in the context of deep learning.
Similar to a hard example mining strategy in essence and in practice, the proposed algorithm is straightforward to implement and computationally as efﬁcient as SGD-based optimizers used for deep learning.
It only requires adding a softmax layer and maintaining an history of the loss values for each training example to compute adaptive sampling probabilities.
In contrast to typical ad hoc hard mining approaches, and exploiting recent theoretical results in deep learning optimization, we prove the convergence of our DRO algorithm for over-parameterized deep learning networks with ReLU activation and ﬁnite number of layers and parameters.
Preliminary results demonstrate the feasibility and usefulness of our approach.
I NTRODUCTION
In standard deep learning pipelines a neural network h with parameters θ is trained by minimizing the mean of a per-example loss L over a training dataset { ( x i , y i ) } ni =1 .
This corresponds to the empirical risk minimization optimization problem
[Equation]
Since the empirical risk is equal to the expectation of the per-example loss over the empirical training data distribution, an approximate solution of (1) can be obtained efﬁciently by Stochastic Gradient
Descent (SGD) with a uniform sampling over the training data (Bottou et al., 2018).
This approach has led to spectacular results in term of average performance, but may generate outliers with high loss values compared to the average loss.
Such cases can even be observed for elements belonging to the training data set.
This is because approximate solutions obtained by SGD are prone to ignore a few hard examples in order to obtain a low mean per-example loss.
Far from being only of academic interest, outlier results have, for example, been consistently reported in the context of deep learning for brain tumor segmentation, as illustrated in the recent annual BRATS challenges (Bakas et al., 2018).
For safety-critical systems, such as those used in healthcare, where outliers must be avoided, this is not satisfactory.
Efﬁcient biased sampling methods, including hard example mining (Shrivastava et al., 2016;
Loshchilov & Hutter, 2015; Chang et al., 2017) and weighted sampling (Bouchard et al., 2015;
Berger et al., 2018; Gibson et al., 2018), have been proposed to mitigate this issue.
However, even though these works typically start from an Empirical Risk Minimization formulation, it is not clear how those heuristics actually relate to Empirical Risk Minimization in theory.
Distributionally Robust Optimization (DRO) is an alternative to Empirical Risk Minimization (1) that takes into account uncertainty in the empirical data distribution.
The deep neural network h is now trained by accounting for potential deviations from the empirical training data distribution.
Formally, DRO corresponds to the min-max non-convex-concave optimization problem
[Equation]
where φ is a convex function that deﬁnes a φ -divergence (Csisz´ar et al., 2004), q = ( q i ) ni
=1 corresponds to arbitrary weighted sampling distributions over the training data, and β > 0 is a robustness parameter.
Instead of minimizing the mean per-example loss on the training set, DRO seeks the hardest weighted empirical training data distribution around the (uniform) empirical training data distribution.
This suggests a link between DRO and hard example mining.
The parameter β allows DRO to interpolate between ERM ( β → 0 ) and the minimization of maximum per-example loss ( β → + ∞ ).
Motivations for using the minimization of maximum perexample loss for safety-critical applications has been discussed in (Shalev-Shwartz & Wexler, 2016).
DRO as a generalization of empirical risk minimization for machine learning has been studied in (Duchi et al., 2016; Raﬁque et al., 2018; Namkoong & Duchi, 2016;
Chouzenoux et al., 2019), but still lacks optmization method as efﬁcient as SGD in the non-convex setting of deep learning.
If one could solve the max problem in (2) for a given θ , DRO could be addressed by alternating between this max problem and a minimisation scheme akin to the standard Empirical Risk Minimization (1), but over an adaptively weighted empirical distribution.
However, solving the max problem naively would require performing a forward pass over the entire training dataset.
This can not be done at each iteration efﬁciently for large dataset.
Previously proposed optimization methods for large-scale non-convex-concave problem of the form of (2) are based on the min-max structure of the problem, and consist in alternating between approximate maximization and minimization steps (Raﬁque et al., 2018; Lin et al., 2019; Jin et al., 2019).
However, they differ from
SGD methods for Empirical Risk Minimization by the introduction of additional hyperparameters for the optimizer such as a second learning rate and a ratio between the number of minimization and maximization steps.
As a result, those alternate optimization methods are difﬁcult to use as a replacement of Empirical Risk Minimization in practice.
In addition, from a practical perspective, those min-max methods do not use the link between DRO and adaptive weighted sampling, therefore departing from efﬁcient heuristics used in hard example mining.
From a theoretical perspective, they further make the assumption that the model is either smooth or weakly-convex, but none of those properties are true for the deep neural networks with
ReLU activation that are largely used in practice.
In this work, we propose SGD with hardness weighted sampling, a novel, principled optimization method for training deep neural networks with Distributionally Robust Optimization inspired by hard example mining.
Compared to SGD, our method only requires introducing an additional softmax layer and maintaining an history of the stale per-example loss to compute sampling probabilities over the training data.
Since the loss is already computed at each iteration for SGD, our SGD with hardness weighted sampling is computationally as efﬁcient as SGD methods for Empirical Risk
Minimization.
In practice, we show that our method performs favorably to SGD in the case of class imbalance despite using hyperparameters previously tuned for SGD.
We also formally link DRO in our method with hard example mining.
As a result our method can be seen as a principled hard example mining approach.
In this context, the robustness parameter β controls the trade-off between exploitation and exploration in the hard example mining process.
Last but not least, we generalize recent results in the convergence theory of SGD with ERM and over-parameterized deep learning networks with ReLU activation (Allen-Zhu et al., 2019; 2018;
Cao & Gu, 2019; Zou & Gu, 2019) to our SGD with hardness weighted sampling for DRO.
This is, to the best of our knowledge, the ﬁrst convergence result for deep learning network with ReLU trained with DRO.
R ELATED W ORK IN DRO WITH A W ASSERSTEIN D ISTANCE
In this work, we focus on DRO with a φ -divergence (Csisz´ar et al., 2004).
In this case, the data distributions that are considered in the DRO problem (2) are restricted to sharing the support of the empirical training distribution.
In other words the weights assigned to the training data can change, but the training data itself remains unchanged.
Another popular formulation for DRO is DRO with a Wasserstein distance (Sinha et al., 2017;
Duchi et al., 2016; Staib & Jegelka, 2017; Chouzenoux et al., 2019).
In contrast to DRO with a φ -divergence, using a Wasserstein distance in DRO seeks to apply small data augmentation to the training data to make the deep learning model robust to small deformation of the data, but the sampling weights of the training data distribution typically remains unchanged.
In this sense, DRO with a φ -divergence and DRO with a Wasserstein distance can be considered as orthogonal endeavours.
While we show that DRO with φ -divergence can be seen as a principled hard example mining method, it has been shown that DRO with a Wasserstein distance can be seen as a principled adversarial training method (Sinha et al., 2017; Staib & Jegelka, 2017).
Contrary to our SGD with hardness weighted sampling, the optimization methods of (Sinha et al., 2017; Staib & Jegelka, 2017) exploit alternate minimization maximization strategies that are orders of magnitude slower than SGD even for deep neural networks with smooth activation functions.
In addition, the method of (Sinha et al., 2017) is NP-hard in the case of deep neural networks with
ReLU activation functions.
M ACHINE L EARNING WITH D ISTRIBUTIONALLY R OBUST O PTIMISATION AND φ -D IVERGENCE
In machine learning based on Empirical Risk Minimization (ERM), a predictor h is trained using a training dataset { ( x i , y i ) } ni =1 to perform well on average on a task for which the performance is measured on a per-example basis by a smooth criteria L .
Note that parameter regularization terms can easily be embedded in L since they are independent of the example.
For ease of presentation, we focus on the supervised machine learning setting ,where h : x (cid:55)→ y , and omit explicitly mentioning any parameter regularisation term.
Let ∆ n ⊂ R n be the set of empirical weighted training data distribution deﬁned according to a given training dataset
[Equation]
and let ˆ p data be the corresponding uniform empirical training data distribution.
In other words, ∆ n corresponds to all sampling weights that can be applied to the training dataset.
Let θ be the set of parameters of the predictor h ( . ; θ ) : x (cid:55)→ y we want to train.
We assume L is a smooth and potentially non-convex function, and let h : θ (cid:55)→ ( h ( x i ; θ ))
ni
=1 be the vector of inferred outputs from the training data.
By abuse of notation, we denote L ( h ( θ ))
=
( L ( h ( x i ; θ ) , y i ))
ni =1 .
Deﬁnition 3.1 (Mean Loss) .
[Equation]
The ERM predictor, as used in most learning settings, is obtained by minimizing the mean loss (4).
Deﬁnition 3.2 (Empirical Risk Minimization (ERM) predictor) .
[Equation]
However, ˆ p data is typically biased compared to the true data distribution.
Predictors trained with
ERM are prone to fail on new examples that are not well represented in the training data.
Distributionally Robust Optimization (DRO) is an alternative to ERM that mitigates this issue by encouraging robustness to uncertainty in the empirical training data distribution.
DRO, in its simplest form, is based on the notion of φ -divergence that we use to induce robustness with respect to all the empirical distributions ∆ n of the training dataset.
Deﬁnition 3.3 ( φ -Divergence) .
Let φ : R + → R ∪ { + ∞} be a closed, convex, lower semicontinuous function such that ∀ z ∈ R + , φ ( z ) ≥ φ (1) = 0 .
The φ -Divergence D φ is deﬁned as, for all p =
[Equation]
i
[Equation] [Equation]
Example 3.1.
For φ : z (cid:55)→ z log( z ) , D φ is the Kullback-Leibler (KL) divergence:
[Equation]
[Equation]
And, for φ : z (cid:55)→ ( z − 1) 2 , D φ is the Pearson χ 2 divergence:
[Equation]
Deﬁnition 3.4 (Distributionally Robust Loss) .
[Equation]
where β > 0 is a hyperparameter that controls the amount of robustness.
For a given φ -divergence, we deﬁne the DRO predictor, that is obtained by minimizing the distributionally robust loss (9) instead of the mean loss (4).
Deﬁnition 3.5 (Distributionally Robust Optimization (DRO) predictor) .
[Equation]
It is worth noting that DRO interpolates between ERM as β −→ 0 and the minimization of the maximum loss as β −→ ∞ , and is equivalent to a mean-variance trade-off for β small (Gotoh et al., 2018). 
[Equation]
where V ˆ p data is the empirical variance.
Furthermore, we observe that the distributionally robust loss (9) is an upper bound to the mean loss (4) (independently to the choice of φ and β ), i.e. for all φ -divergence and all β >∀ θ , M ( L ( [Equation] [Equation]
≤ R ( L ( h ( θ ))
[Equation]
We now make assumptions for the φ -divergence to simplify the derivations in the remainder of the paper.
Let D φ
Assumption 3.1 (Regularity of φ ) .
φ : R + → R is two times continuously differentiable on [0 , n ] , ρ -strongly convex, for ρ > 0 , i.e.:
[Equation]
and φ satisﬁes:
[Equation]
These assumptions are veriﬁed by most of the φ -divergence used in practice (e.g. the KL divergence and Pearson χ 2 divergence).
SGD WITH H ARDNESS W EIGHTED S AMPLING 4. D ISTRIBUTIONALLY R OBUST O PTIMIZATION WITH SGD AND A DAPTIVE S AMPLING
Existing optimization methods for DRO with a non-convex predictor h alternate between approximate minimization and maximization steps (Raﬁque et al., 2018; Jin et al., 2019; Lin et al., 2019), requiring the introduction of additional hyperparameters compared to SGD.
These are difﬁcult to tune in practice and convergence has not been proven for deep networks with ReLU activations.
In this section, we highlight properties that allows to link Distributionally Robust Optimization with
SGD combined with adaptive sampling.
Our analysis relies mainly on Fenchel duality (Moreau, 1965) and the notion of Fenchel conjugate (Fenchel, 1949) that we now deﬁne.
Deﬁnition 4.1 (Fenchel Conjugate Function) .
Let f : R m → R ∪ { + ∞} a proper function.
The
Fenchel conjugate of f is deﬁned as ∀ v ∈ R m , f ∗ ( v ) =
max x ∈ R m (cid:104) v , x (cid:105) − f ( x )
Let ∀ p ∈ R n ,
[Equation]
where δ ∆ n is the characteristic function of the closed convex set ∆ n , i.e.
[Equation]
One can remark that the distributionally robust loss R (9) can be rewritten using the Fenchel conjugate function of G .
This allows to obtain regularity properties for R .
Lemma 4.1 (Regularity of R ) .
If φ satisﬁes Assumption 3.1, then G and R satisfy the following:
[Equation]
Equation (16) follows from Deﬁnition 4.1.
The proof of (15) and (17) is found in Appendix C.2.
According to (15), the optimization problem (16) is strictly convex and admit a unique solution in ∆ n .
Let us denote this solution
[Equation]
The following lemma shows that the gradient, with respect to θ , of the distributionally robust loss (9) at a given θ can be rewritten as the expectation, with respect to the weighted empirical distribution ¯ p ( L ( h ( θ ))) , of the per-example loss gradient.
We further show that straightforward analytical formulas exist for ¯ p when relying on classical φ -divergences.
This result motivates our Algorithm 4.
for efﬁcient training with the distributionally robust loss.
Lemma 4.2 (Stochastic Gradient of the Distributionally Robust Loss) .
For all θ , we have
[Equation]
where ∇ v R is the gradient of R with respect to its input.
The proof is found in Appendix C.3.
It is apparent from (19) that, given ¯ p , an estimate of ∇ θ R could easily be provided by sampling a batch according to ¯ p and estimating the per-example loss gradients in the batch as per standard practice.
We now provided closed-form formulas for ¯ p given L ( h ( θ )) .
Example 4.1.
For the KL divergence (i.e. φ : z
[Equation]
z log( z ) − z + 1 ), we have (see C.1 for a proof)
[Equation]
And for the Pearson χ 2 divergence (i.e. φ : z (cid:55)→ ( z − 1) 2 ), we have:
[Equation]
In both cases, we can verify consistency with (11) as
[Equation] [Equation]
0 , learning rate η > 0 , batch size b ∈ { 1 , . . . , n } .
2: Initialization: 3: Initialise θ randomly 4: Initialise the loss history ˜
L = −5:
Warm start : 6: // Split the training data into batches B and run 1 epoch with classic SGD 7: for { ( x i , y i ) } i ∈ I in B do 8: // Run forward pass and store losses for all the samples in the batch 9: for i ∈
I do 10: ˜ L
i ← L ( h ( x i ; θ ) , y i ) 11: // Run backward pass and update the parameters of the model 12: θ ← θ − η 1 b (cid:80) i ∈ I ∇ θ L ( h ( x i ; θ ) , y i ) 13:
SGD with dynamic hardness weighted sampling: 14: for epoch = 2 , . .
, T do 15: for iteration i = 1 , . . .
, (cid:16)(cid:106) n b (cid:107) + 1 (cid:17) do 16: //
Run softmax to update the sampling probabilities of the samples 17: ˆ p = softmax( β ˜ L ) 18: // Draw a batch with replacement using the probability distribution ˆ p 19: { ( x i , y i ) }
i ∈ I such that I i.i.d.
∼ ˜ p and | I | =
b 20: // Run forward pass and update losses for all the samples in the batch 21: for i ∈
I do 22:
˜ L
i ← L ( h ( x i ; θ ) , y i ) 23: // Run backward pass and update the parameters of the model 24: θ ← θ − η 1 b (cid:80) i ∈ I ∇ θ L ( h ( x i ; θ ) , y i ) 25:
Output: θ
4. E FFICIENT A LGORITHM FOR D ISTRIBUTIONALLY R OBUST D EEP L EARNING
The second equality in (19) implies that ∇ θ L ( h i ( θ ) , y i ) is an unbiased estimator of the distributionally robust loss
R ( L ( h ( θ )))
when i is sampled with respect to ¯ p ( L ( h ( θ ))) .
This suggests that the distributionally robust loss can be minimized efﬁciently by Stochastic Gradient Descent by sampling mini-batches with respect to ¯ p ( L ( h ( θ ))) at each iteration.
However, even though closed-form formulas were provided for ¯ p , evaluating exactly L ( h ( θ )) , i.e. doing one forward pass on the whole training set at each iteration, is computationally prohibitive for large dataset.
In practice, we propose to use a stale version of L ( h ( θ )) by maintaining an online history of the loss values of the training examples during training (cid:0) L ( h ( x i ; θ ( t i ) ) , y i ) (cid:1) .
Where for all i , t i is the last iteration at which the per-example loss of example i has been computed.
Using the the Kullback-Leibler divergence as φ -divergence, this leads to the Stochastic Gradient Descent with KL hardness weighted sampling algorithm proposed in Algorithm 4.1.
This would also apply to the χdivergence mutatis mutandis.
In contrast to alternate min-max optimization methods, our SGD with a adaptive sampling strategy is similar to the SGD-based optimizers used by the vast majority of deep learning practitioners (e.g.
SGD, SGD with momentum, ADAM).
Compared to standard SGD-based training optimizers for the mean loss, our algorithm requires only an additional softmax operation per iteration and to store an additional vector of size n (number of training examples), thereby making it ideally suited for deep learning applications.
4. DRO AS P RINCIPLED O NLINE H ARD E XAMPLE M INING
In this section, we discuss the relationship between DRO and Online Hard Example Mining (OHEM) (Shrivastava et al., 2016).
SGD with an ad hoc adaptive sampling strategy is already used in practice while starting from a mean loss optimization formulation in the online hard example mining literature (Loshchilov & Hutter, 2015; Shrivastava et al., 2016).
Similarly to our algorithm, in OHEM heuristics, the hard examples , those training examples with relatively high values of the loss, are sampled more often.
We formalize this in the following deﬁnition for OHEM sampling strategies.
Deﬁnition 4.2 (Online Hard Example Mining Sampling) .
Any adaptive sampling method such that the probability p i of sampling example
x i is an non-decreasing function of the (potentially stale) loss value associated with x i .
Theorem 4.1.
The proposed hardness weighted sampling is a hard example mining sampling for any φ -Divergence that satisﬁes Assumption 3.1.
In addition, the probability p i of sampling example x i is an non-increasing function of the loss value associated with x j for all j (cid:54)
=
i .
See Appendix C.4 for the proof.
The second part of Theorem 4.1 implies that as the loss of an example diminishes, the sampling probabilities of all the other examples increase.
As a result,
Distributionally Robust Optimization balances exploration and exploitation.
C ONVERGENCE OF SGD WITH H ARDNESS W EIGHTED S AMPLING FOR O VER - PARAMETERIZED D EEP N EURAL N ETWORKS WITH ReLU
Convergence results for over-parameterized deep learning has recently been proposed in (AllenZhu et al., 2019).
Their work gives convergence guarantees for deep neural networks h with any activation function (including ReLU), and with any (ﬁnite) number of layers L and parameters m , under the assumption that m is large enough.
Although some results suggest that this theory cannot explain all the properties of deep learning observed in practice (Chizat et al., 2019), at the time of writing, this is the most realistic setting for which a convergence theory of deep learning exists.
In this section, we demonstrate the ﬁrst convergence guarantees for deep neural networks with ReLU trained with DRO.
Our analysis is based on the results developed in (Allen-Zhu et al., 2019) which is a simpliﬁed version of (Allen-Zhu et al., 2018).
Improving on those theoretical results would automatically improve our results as well.
We focus at providing theoretical tools that could be used to generalize any convergence result for ERM using SGD to DRO using Algorithm 4.1.
Let us ﬁrst state our assumptions on the neural network h , and the per-example loss function L .
Assumption 5.1 (Deep Neural Network) .
In this section, we use the following notations and assumptions similar to (Allen-Zhu et al., 2019):
• h is a fully connected neural network with L + 2 layers, ReLU activation function, and m nodes in each hidden layers • For all i ∈ { 1 , . . . , n } , we denote h i : θ (cid:55)→ h
i ( x i ; θ ) the output d dimensional scores of h applied to example x i of dimension d . • θ =
( θ l ) L +1 l =0 is the set of parameters of the neural network h, where θ l is the set of weights for layer l with θ 0 ∈ R d × m , θ
L +1 ∈ R m × d , and θ l ∈ R m × m for any other l . • (Data separation)
It exists δ > 0 such that for all i, j ∈ { 1 , . . . , n } , if i (cid:54) =
j, (cid:107) x
i − x j (cid:107) ≥ δ .
• We assume
m ≥ Ω
( d × poly ( n, L, δ − 1 )) for some sufﬁciently large polynomial poly , and δ ≥ O (cid:0) 1 L (cid:1) .
We refer the reader to (Allen-Zhu et al., 2019) for details about poly .
• The parameters θ =
( θ l ) L +1 l =0 are initialized at random such that: – [ θ 0 ] i,j ∼ N (cid:0) 0 , 2 m (cid:1) for every ( i, j ) ∈ { 1 , . . .
, m } × { 1 , . . .
, d } – [ θ l ] i,j ∼ N (cid:0) 0 , 2 m (cid:1) for every ( i, j ) ∈ { 1 , . . .
, m } 2 and l ∈ { 1 , . . .
, L } – [ θ L +1 ] i,j ∼ N (cid:0) 0 , 1 d (cid:1) for every ( i, j ) ∈ { 1 , . . .
, d } × { 1 , . . .
, m }
Assumption 5.2 (Regularity of L ) .
For all i, L i is a C ( ∇ L ) -gradient
Lipschitz continuous, C ( L ) -
Lipschitz continuous, and bounded (potentially non-convex) function.
We ﬁrst generalize the convergence of SGD in (Allen-Zhu et al., 2019, Theorem 2) to the minimization of the distributionally robust loss using SGD and an exact hardness weighted sampling (19), i.e. with an exact non-stale loss history.
Theorem 5.1 (Convergence of Robust SGD with exact Loss History) .
Let batch size 1 ≤ b ≤ n , and (cid:15) > 0 .
Suppose there exists constants C 1 , C 2 , C 3 > 0 such that the number of hidden units satisﬁes m ≥ C 1 ( d(cid:15) − 1 × poly ( n, L, δ − 1 )) , δ ≥ (cid:0) C 2 L (cid:1) , and the learning rate be η exact =
C 3 (cid:16) min (cid:16) 1 ,
[Equation]
(cid:17) .
There exists constants C 4 , C 5 > 0 such that with probability at least 1 − exp (cid:0) − C 4 (log 2 ( m )) (cid:1) over the randomness of the initialization and the mini-batches, Robust SGD with exact loss vector ﬁnds (cid:107)∇ θ
( R ◦ f ◦ h )( θ ) (cid:107) ≤ (cid:15) after T =
C 5 (cid:16)
[Equation]
where α = min θ min i ¯ p
i ( L ( θ )) is lower bound on the sampling probabilities.
For the KullbackLeibler φ -divergence, and for any φ -divergence satisfying assumption 3.1 with a robustness parameter β small enough, we have α > 0 .
We refer the reader to (Allen-Zhu et al., 2019, Theorem 2) for the values of the constants C 1 , C 2 , C 3 , C 4 , C 5 and the deﬁnitions of the polynomials.
Compared to (Allen-Zhu et al., 2019, Theorem 2) only the learning rate differs.
The min(1 , . ) operation in the formula for η exact allows to guarantee that η exact ≤ η (cid:48) where η (cid:48) is the learning rate of (Allen-Zhu et al., 2019, Theorem 2).
The proof can be found in Appendix C.6.3.
It is worth noting that for the KL φ -divergence, ρ = 1 n .
In addition, in the limit β → 0 , which corresponds to ERM, we have α → 1 n .
As a result, we recover exactly Theorem 2 of (AllenZhu et al., 2019) as extended in their Appendix A for any L that satisﬁes assumption 5 . 2 with
C ( ∇ L ) = 1 .
When the amount of distributionally robustness increases the sampling differs more and more from the uniform sampling and becomes more sensitive to changes of the loss distribution.
One way to mitigate this issue is to reduce the learning rate.
The conditions of Theorem 5.1 are consistent with this observation since when β increases, α and η exact decreases.
In practice in algorithm 4.1, we have access only to a stale loss history.
We know restate the convergence of Robust SGD with a stale loss history and a warm-up as in Algorithm 4.1.
Theorem 5.2 (Convergence of Robust SGD with Stale Loss History and warm-up) .
Let batch size 1 ≤ b ≤ n , and (cid:15) > 0 .
Under the conditions of Theorem 5.1, the same notations, and with the learning rate η stale = C 6 min (cid:18) 1 ,
[Equation]
− α )
[Equation]
(cid:19) × η exact for a constant C 6 > 0 .
With probability at least 1 − exp (cid:0) − C 4 (log 2 ( m )) (cid:1) over the randomness of the initialization and the mini-batches, Robust SGD with exact loss vector ﬁnds (cid:107)∇ θ
R ◦ f ◦ h )( θ ) (cid:107) ≤ (cid:15) after T =
C [Equation] (cid:16)
[Figure]
iterations. Where C ( L ) > 0 is a constant such that L is C ( L ) -Lipschitz continuous, and A ( ∇ L ) > 0 is a constant that bound the gradient of L with respect to its input.
C ( L ) and A ( ∇ L ) are guaranteed to exist under assumptions 5.1.
The proof can be found in Appendix C.7.
Compared to Theorem 5.1 only the learning rate differs.
Similarly to Theorem 5.1, when β tends to zero we recover Theorem 2 of (Allen-Zhu et al., 2019).
It is worth noting that when β increases,
[Equation]
decreases.
This implies that η stale decreases faster than η exact when β increases.
This was to be expected since the error that is made by using the stale loss history instead of the exact loss increases when β increases.
E XPERIMENTS
We now illustrate the properties of our SGD with hardness weighted sampling described in Algorithm 4.1 for training deep neural networks with ReLU activation functions with DRO (10).
6. E XPERIMENTS ON MNIST
We create a bias between training and testing data distribution of MNIST (LeCun, 1998) by keeping only 1% of the digits 3 in the training dataset, while the testing dataset remains unchanged.
Implementation details can be found in Appendix A.1.
Figure 2 suggests that DRO and ERM lead to different optima. Indeed, the optima of DRO for β = 10 outperforms ERM by more than 10% of accuracy on the under-represented class, as illustrated on the top right panel of Figure 2. This suggests that DRO leads to better generalization than ERM. Especially, it appears that DRO is more robust than ERM to domain gaps between the training and the testing dataset.
In addition, the top left panel of Figure 2 suggests that DRO can convergence faster than ERM.
Furthermore, the variations of learning curves with β are consistent with our theoretical insight in
Theorem 5.2.
As β decreases to 0 , the learning curves of DRO with our method converges to the learning curve of ERM with SGD.
6. E XPERIMENTS ON CIFAR
We now show that our method for training deep neural network with DRO is not only more robust to bias in the training dataset, but can also be used with a higher learning rate than ERM.
We create a bias between training and testing data distribution in CIFAR10 (Krizhevsky et al., 2010) by keeping only 10% of the cats in the CIFAR10 training dataset, while the CIFAR10 testing dataset remains unchanged.
We used the state-of-the-art WRN- 28 - 10 deep neural network architecture proposed in (Zagoruyko & Komodakis, 2016).
We kept the hyperparameters and data augmentation, except that we used a higher learning rate equal to 1 (see Appendix A.2 for more details).
The learning curves of ﬁgure 1, computed on the CIFAR10 testing dataset, shows that our methods outperforms the baseline.
This suggests that replacing SGD-based optimizer by our Algorithm 4.
in an existing deep learning pipeline is straightforward, and that our method is robust to the choice of the learning rate.
C ONCLUSION AND D ISCUSSION
We have shown that efﬁcient training of deep neural networks with Distributionally Robust Optimization (DRO) with a φ -divergence (10) is possible.
Our Stochastic Gradient Descent (SGD) with hardness weighted sampling is a principled hard example mining method.
It is as straightforward to implement, and as computationally efﬁcient as SGD for Empirical Risk Minimization.
It can be used for deep neural networks with any activation function (including ReLU), and with any perexample loss function.
In addition, we prove the convergence of our method for over-parameterized deep neural networks.
This is, to the best of our knowledge, the ﬁrst convergence result for training a deep neural network based on DRO.
Our experiments on MNIST illustrate several behaviours that were predicted by our Theorem 5.2.
Our experiments on CIFAR10 illustrate that our SGD with hardness weighted sampling can be used successfully as a replacement of SGD in a state-of-the-art deep learning pipeline that was tuned for
SGD.
In addition, our results on CIFAR10 suggest that our method is more robust than SGD to bias in the training data distribution and to the choice of the learning rate.
The fact that our algorithm can be used with a higher learning rate than SGD goes beyond the prediction of our theoretical results.
Investigating this advantage is left for future work.
A M ORE ON OUR E XPERIMENTS A. E XPERIMENTS ON MNIST
[Figure]
A.1. I MPLEMENTATION D ETAILS
For our experiments on MNIST, we used a Wide Residual Network (WRN) (Zagoruyko & Komodakis, 2016).
The familly of WRN models has proved to be very efﬁcient and ﬂexible, achieving state-of-the-art accuracy on several dataset.
More speciﬁcally, we used WRN- 16 - 1 (see Zagoruyko & Komodakis, 2016, section 2.3).
For the optimization we used a learning rate of 0 . 01 .
No momentum or weight decay were used.
No data augmentation was used.
A.1. C OMMENT ON THE STABILITY OF THE LEARNING CURVES
For large values of β (here β ≥ 10 ), instabilities appear in the testing learning curves , as illustrated in the top panels of ﬁgure 2.
However, the bottom left panel of ﬁgure 2 shows that the training loss curves for β ≥ 10 were actually stable there. We also observe that during the iterations for which instabilities appears on the testing set , the standard deviation of the per-sample loss on the training set is relatively high (i.e. the hardness weighted probability is further away from the uniform distribution). This suggests that the apparent instabilities on the testing set are related to differences between the DRO loss (9) and the Mean loss (4).
A.1. C OMMENT ON E ARLY -S TOPPING WITH ERM AND DRO
It is worth noting that we do not use the knowledge of the training dataset bias in our method.
For a less obvious bias, we would have access to the accuracy on the under-representated subset of the training set.
As a result, in general, convergence criteria can be based only on the global measure of accuracy (in our case the top left panel of Figure 2).
Early-stopping is a widely used heuristic for selecting the optimal number of epochs to prevent overﬁtting (Bengio, 2012).
Early-stopping crucially depends on the patience parameter. In the top left panel of Figure 2, there is no improvement of the accuracy of ERM of more than 0 . 04% between epoch 20 and 30 . As a result, if the patience parameter is too low, ERM might be considered to plateau at an epoch at which it ERM achieves an accuracy of 0% on the under-represented class.
This suggests two things.
First, the mean accuracy is not a good criteria to decide when to stop the training of ERM for safety-critical systems.
Second, our SGD with hardness weighted sampling converge faster than SGD to a safe solution, and is more robust to the hyperparameters of earlystopping.
A. E XPERIMENTS ON CIFAR
For our experiments on the CIFAR10 dataset (Krizhevsky et al., 2010), we used the WideResNet
WRN- 28 - 10 as described in (Zagoruyko & Komodakis, 2016).
We used the same pipeline as in (Zagoruyko & Komodakis, 2016), that we summarized here:
• we used SGD with momentum β 1 = 0 .•
we trained for 200 epochs and multiply the learning rate by 0 .
2 at the beginning of epochs 60 , 120 and• we used a weight decay of 0 .•
we used a batch size of• dropout was not used • we normalized the images with ZCA whitening as in (Goodfellow et al., 2013)
• we used horizontal ﬂip with probability 0 .
5 and random crops applied to the image after padded by 4 pixels
The only difference is that we used a larger initial learning rate lr = 1 , rather than lr = 0 .in
(Zagoruyko & Komodakis, 2016).
We based our implementation on the code provided by the authors of (Zagoruyko & Komodakis, 2016) which can be found at https://github.com/szagoruyko/ wide-residual-networks .
B N OTATIONS
For the ease of following the proofs we ﬁrst summarize our notations.
B. P ROBABILITY T HEORY N OTATIONS
• ∆ n = { p =
( p i ) ni =1 ∈ [0 , 1] n , (cid:80) i p i = 1 } • Let q = ( q i ) ∈ ∆ n , and f a function, we denote E q [ f ( x )] :
= (cid:80) ni
=1 q i f ( x i ) .
• Let q ∈ ∆ n , and f a function, we denote V q [ f ( x )] := (cid:80) ni
=1 q i (cid:107) f ( x i ) − E
q
[ f ( x )] (cid:107) 2 .
• ˆ p data is the uniform training data distribution,
i
e. ˆ p data = (cid:0) 1 n (cid:1) n i =1 ∈ ∆ n
B. M ACHINE L EARNING NOTATIONS
• n is the number of training examples • d is the dimension of the output • d is the dimension of the input • training data: { ( x i , y i ) } ni =1 , where for all i ∈ { 1 , . . . , n } , x
i ∈ R d and y i ∈
R d • h : x (cid:55)→ y is the predictor • θ is the set of parameters of the predictor • For all i , h
i : θ (cid:55)→ h
( x i ; θ ) is the output of the network for example
i as a function of θ • L is the objective function • L i : h
i (cid:55)→ L ( h i , y i ) is the objective function for example i .
• By abuse of notation we also denote by L the function L : ( h i ) ni
=
(cid:55)→ ( L i ( h i )) ni = • b ∈ { 1 , . . .
, n } is the batch size • η > 0 is the learning rate • ERM is short for Empirical Risk Minimization
B. D ISTRIBUTIONALLY R OBUST O PTIMISATION N OTATIONS
• Forall θ , R ( L ( h ( θ )))
=
max q ∈ ∆ n E q [ L ( h ( x ; θ ) , y )] − 1 β D φ ( q (cid:107) ˆ p data ) is the Robust Loss evaluated at θ , where β > 0 is the parameter that adjusts the distributionally robustness (see (9) for more details).
• DRO is short for Distributionally Robust Optimisation
B. M ISCELLANEOUS
• By abuse of notation, and similarly to (Allen-Zhu et al., 2019), we use the BachmannLandau notations to hide constants that do not depend on our main hyper-parameters.
Let f and g be two scalars, we note:
[Equation]
f ≥ cg f =
{+Θ( g ) ⇐⇒ ∃ c 1 > 
+}
and ∃ c 2 > cs.t.
c 1 g ≤ f ≤ c 2 g
C P ROOFS C. P ROOF OF E XAMPLE 4.1: FORMULA OF THE SAMPLING PROBABILITIES FOR THE KL DIVERGENCE
We give here a simple proof of the formula of the sampling probabilities for the KL divergence as φ -divergence (i.e. φ : z
(cid:55)→ z log( z ) − z + 1 )
[Equation]
For any θ , the distributionally robust loss (9) for the KL divergence at θ is given by
[Equation]
To simplify the notations, let us denote v
= ( v i ) {+n i =
+}
[-=
-]=
L ◦ h ( θ ) =
( L i ◦ h
i ( θ ))
n i =1 , and ¯ p = (¯ p i ) ni
=
= ¯ p ( L ( h ( θ ))) .
Thus ¯ p ( L ( h ( θ ))) is, by deﬁnition, solution of the optimization problem [Equation]
[Equation]
First, let us remark that the function q (cid:55)→ (cid:80) ni
=1 q i log ( nq i ) is strictly convex on the non empty closed convex set ∆ n as a sum of strictly convex functions.
This implies that the optimization (23) has a unique solution and as a result ¯ p ( L ( h ( θ ))) is well deﬁned.
We now reformulate the optimization problem (23) as a convex smooth constrained optimization problem by writing the condition q ∈ ∆ n as constraints.
[Equation]
There exists a Lagrange multiplier λ ∈ R [Equation] such that the solution ¯ p of (24) is characterized by
[Equation]
Which we can rewrite as
[Equation]
The last equality gives
[Equation]
And by replacing in the formula of the ¯ p i
[Equation]
Which corresponds exactly to
[Equation]
C.
P ROOF OF LEMMA 4.1: REGULARITY PROPERTIES OF R
For the ease of reading, let us ﬁrst recall that given a φ -Divergence that satisﬁes assumptions 3.1, we have deﬁned in (9) [Equation]
[Equation]
And in (13)
[Equation] [Equation]
where δ ∆ n is the characteristic function of the closed convex set ∆ n , i.e.
[Equation]
[Equation]
We now prove Lemma 4.1 on the regularity of R.
Lemma C.1 (Regularity of R – Restated from Lemma 4.1) .
Let φ that satisﬁes Assumption 3.1, G and R satisfy
[Equation]
φ is ρ -strongly convex on [0 , n ] so
[Equation]
Let p =
( p i ) ni =1 , q = ( q i ) ni
=1 ∈ ∆ n , and λ ∈ [0 , 1] , using (33) and the convexity of δ ∆ n , we obtain:
[Equation]
This proves that G is nρβ -strongly convex.
Since G is convex, R
= G ∗ is also convex, and R ∗ =
( G ∗ ) ∗ =
G (Hiriart-Urruty & Lemar´echal, 2013).
We obtain (31) using Deﬁnition 4.1.
We now show that R is Frechet differentiable on R n .
Let v ∈ R n .
G is strongly-convex, so in particular G is strictly convex.
This implies that the following optimization problem has a unique solution that we denote ˆ p ( v ) .
[Equation]
In addition ˆ p ∈ ∆ n solution
[Equation]
where we have used (Hiriart-Urruty & Lemar´echal, 2013, Proposition 6.1.2 p.39) for the third equivalence, and (31) for the last equivalence.
As a result, ∂R ( v ) =
{ ˆ p ( v ) } .
this implies that R admit a gradient at v , and
[Equation]
Since this holds for any v ∈ R n , we deduce that R is Frechet differentiable on R n .
We are now ready to show that R is β nρ -gradient Lipchitz continuous by using the following lemma (Hiriart-Urruty & Lemar´echal, 2013, Theorem 6.1.2 p.280).
Lemma C.2.
A necessary and sufﬁcient condition for a convex function f : R n → R to be cstrongly convex on a convex set C is that for all x 1 , x 2 ∈ C (cid:104) s 2 − s 1 , x 2 − x 1 (cid:105) ≥ c (cid:107) x
2 − x 1 (cid:107)for all s i ∈ ∂f ( x i ) , i = 1 , 2 .
Using this lemma for f = G , c = nρβ , and C = ∆ n , we obtain:
For all p 1 , p 2 ∈ ∆ n , for all v 1 ∈ ∂G ( p 1 ) , v 2 ∈ ∂G
[Equation]
p 2 ) ,
[Equation]
In addition, for i ∈ { 1 , 2 } , v i ∈ ∂G
( p i )
⇐⇒ p
i ∈ ∂R ( v 1 )
= {∇ v R ( v i ) } .
And using Cauchy Schwarz inequality
[Equation]
We conclude that
[Equation]
Which implies that R is βnρ -gradient Lipchitz continuous.
C. P ROOF OF L EMMA 4.2: F ORMULA OF THE DISTRIBUTIONALLY ROBUST LOSS GRADIENT
We prove Lemma 4.2 that we restate here for the ease of reading.
Lemma C.3 (Stochastic Gradient of the Distributionally Robust Loss – Restated from Lemma 4.2) .
For all θ , we have
[Equation] [Equation]
38) where ∇ v R is the gradient of R with respect to its input.
For a given θ , equality (37) is a special case of (36) for v =
L ( h ( θ )) .
Then using the chain rule and (37),
[Equation]
C. P ROOF OF T HEOREM 4.1: D ISTRIBUTIONALLY R OBUST O PTIMIZATION AS P RINCIPLED H ARD E XAMPLE M INING
Let D φ an φ -divergence satisfying Assumption 3.1, and v = ( v i ) ni =1 ∈ R n .
v will play the role of a generic loss vector.
φ is strongly convex, and ∆ n is closed and convex, so the following optimization problem has one and only one solution:
[Equation]
Making the constraints associated with p ∈
∆ n explicit, this can be rewritten as
[Equation]
There exists KKT multipliers λ ∈ R and ∀ i, µ i ≥ 0 such that the solution ¯ p = (¯ p i ) ni
=1 satisﬁes:
[Equation]
Since φ is continuously differentiable and strongly convex, we have ( φ (cid:48) )
− 1 = ( φ ∗ ) (cid:48) , where φ ∗ is the Fenchel conjugate of φ (see Hiriart-Urruty & Lemar´echal, 2013, Proposition 6.1.2).
As a result, (41) can be rewritten has:
[Equation]
We now show that the KKT multipliers are uniquely deﬁned.
The µ i ’s are uniquely deﬁned by v and λ :
Since ∀ i ∈ { 1 , . . . , n } , µ i p i = 0 , p i ≥ 0 and µ i ≥ 0 , for all ∀ i ∈ { 1 , . . . , n } , either p i = 0 or µ i = 0 .
In the case p i = 0 , using (42) it comes ( φ ∗ ) (cid:48) ( β ( v i + λ − µ i ))
= 0 .
According to assumption 3.1, φ is strongly convex and continuously differentiable, so φ (cid:48) and ( φ ∗ ) (cid:48)
=
( φ (cid:48) )
− 1 are continuous and strictly increasing functions.
As a result, it exists a unique µ i (dependent to v and λ ) such that:
[Equation]
And (42) can be rewritten as:
[Equation]
λ is uniquely deﬁned by v and a continuous function of v :
Let λ ∈ R that satisﬁes (43).
We have 1 n (cid:80)
[Equation]
=1 ReLU [Equation] φ [Equation]
( β ( v i + λ )) (cid:1) = 1 .
So there exists at least one index i 0 such that
[Equation]
Since ( φ ∗ ) − 1 is continuous and striclty increasing, λ (cid:48)
(cid:55)→ ReLU (cid:0) ( φ ∗ ) (cid:48)
( β ( v i
+ λ (cid:48) ))
(cid:1) is continuous and strictly increasing on a neighborhood of λ .
In addition ReLU is continuous and increasing, so for all i ∈ { 1 , . .
. , n } , λ (cid:48)
(cid:55)→
ReLU (cid:0) ( φ ∗ ) (cid:48)
( β ( v i + λ (cid:48) )) (cid:1) is a continuous and increasing function.
As a result, λ (cid:48) (cid:55)→ 1 n (cid:80) ni =1 ReLU (cid:0) ( φ ∗ ) (cid:48)
( β ( v i + λ (cid:48) )) (cid:1) is a continous function that is increasing on R , and strictly increasing on a neighborhood of λ .
This implies that λ is uniquely deﬁned by v , and that v (cid:55)→ λ ( v ) is continuous.
Hard Example Mining Sampling
For any pseudo loss vector v = ( v i ) ni
=1 ∈ R n , there exists a unique λ and a unique ¯ p that satisﬁes (43), so we can deﬁne the mapping:
[Equation]
where for all v , λ ( v ) is the unique λ ∈ R satisfying (43).
We will now demonstrate that each ¯ p i 0 ( v ) for i 0 ∈ { 1 , . . . , n } is an increasing function of v i and a decreasing function of the v i for i (cid:54)
=
i 0 .
Without loss of generality we assume i
= 1 .
Let v = ( v i ) ni =1 ∈ R n , and (cid:15) > 0 .
Let us deﬁne v (cid:48)
= ( v (cid:48) i ) ni
=1 ∈ R n , such that v (cid:48) 1 = v 1 + (cid:15) and ∀ i ∈ { 2 , . . . , n } , v (cid:48)
i = v i .
Similarly as in the proof of the uniqueness of λ above, we can show that there exists η > 0 such that the function
[Equation]
is continuous and strictly increasing on [ λ ( v ) − η, λ ( v ) + η ] , and F ( λ ( v ))
= 1 .
v (cid:55)→ λ ( v ) is continuous, so for (cid:15) small enough λ ( v (cid:48) )
∈ [ λ ( v ) − η, λ ( v ) + η ] .
Let us now prove by contradiction
that λ ( v (cid:48) ) ≤ λ ( v ) .
Therefore, let us assume that λ ( v (cid:48) ) >
λ ( v ) .
Then, as ReLU ◦ ( φ ∗ ) (cid:48) is an increasing function and F is strictly increasing on [ λ ( v ) − η, λ ( v ) + η ] , and (cid:15) >
we obtain
[Equation]
which is a contradiction.
As a result
[Equation] [Equation]
Using (45), (43), and the fact that ReLU ◦ ( φ ∗ ) (cid:48) is an increasing function, we obtain for all i ∈ { 2 , . . . , n }
[Equation]
In addition n
[Equation]
So necessarily
[Equation]
This holds for any i 0 and any v , which concludes the proof.
C. P ROOF THAT R ◦ L IS ONE - SIDED GRADIENT L IPCHITZ
This property that R ◦ L is one-sided gradient Lipschitz is a key element for the proof of the semismoothness theorem for the distributionally robust loss Theorem C.1.
Under assumption 3.1, we have shown that R ∗ is βnρ -gradient Lipchitz continuous.
And under assumption 5.2, for all i , L i is C ( L ) -Lipschitz continuous and C ( ∇ L ) -gradient
Lipschitz continuous.
Let z =
( z i ) ni
=1 , z (cid:48) =
( z (cid:48) i )
ni =1 ∈ R dn .
We want to show that R ◦ L is one-sided gradient Lipschitz, i [Equation] e. we want to prove the existence of a constant C > 0 , independent to z and z (cid:48) , such that:
[Equation]
We have
[Equation]
Where for all i ∈ { 1 , . . . , n } we have used the chain rule
[Equation]
Let
A
[Equation]
For all i , L i is C ( ∇ L ) -gradient
Lipchitz continuous, so using Cauchy-Schwarz inequality
A
[Equation]
Let
B =
[Equation]
n (cid:88) i =
[Equation]
[Table]
[Equation]
Combining (48), (49) and (50) we ﬁnally obtain:
[Equation]
From there [Equation] we can obtain the following inequality that will be used for the proof of the semismoothness property in Theorem C [Equation]
[Equation]
C. P ROOF OF THE CONVERGENCE OF R OBUST SGD
In this part, we prove the results of Therem 5.1 and 5.2.
They are generalizations of the convergence result for SGD presented in Theorem 2 of (Allen-Zhu et al., 2019).
For the ease of reading the proof
[Equation]
we remind here the chain rules for the distributionally robust loss (9) that we are going to use intensively in the following proofs.
Chain rule for the derivative of R ◦ L with respect to the network outputs h : [Equation] n }
[Equation]
Chain rule for the derivative of R ◦ L ◦ h with respect to the network parameters θ :
[Equation]
where for all i ∈ { 1 , . . .
n } , ∇ θ h
i ( θ ) is the transpose of the Jacobian matrix of h i as a function of θ .
C.6. S EMI - SMOOTHNESS PROPERTY FOR THE D ISTRIBUTIONALLY R OBUST L OSS
We prove the following lemma which is a generalization of Theorem 4 in (Allen-Zhu et al., 2019) for the distributionally robust loss (9).
Theorem C.1 (Semi-smoothness of the Distributionally Robust Loss) .
Let ω ∈
[Equation]
(cid:17)(cid:105) , and the θ (0) being initialized randomly as described in assumption 5.1.
With probability as least 1 − exp ( − Ω( mω 3 / 2 L )) over the initialization, we have for all θ , θ (cid:48) ∈ ( R m × m )
L with (cid:13)(cid:13)
θ − θ (0) (cid:13)(cid:13) 2 ≤ ω , and (cid:107) θ − θ (cid:48)
(cid:107) 2 ≤ ω
R ( L (
[Equation]
[Equation]
(cid:48) ))
[Equation]
[Equation]
[Equation]
[Equation]
+ (cid:104)∇
[Equation]
( R
[Equation]
L ◦ h )( θ ) ,
[Equation]
(cid:48)
[Equation]
(cid:105)
+ (cid:107)∇
[Equation]
[Equation]
L )
[Equation]
[Equation]
[Equation]
+
[Equation]
where for all layer l ∈ { 1 , . . .
, L } , θ l is the vector of parameters for layer l , and
[Equation]
To compare this semi-smoothness result to the one in (Allen-Zhu et al., 2019, Theorem 4), let us ﬁrst remark that
[Equation]
As a result, our result is analogous to (Allen-Zhu et al., 2019, Theorem 4), up to an additional multiplicative factor (cid:16) βC
( L ) 2 nρ + 2 C ( ∇ L ) (cid:17) in the last term of the right-hand side.
It is worth noting that there is also implicitly an additional multiplicative factor C ( ∇ L ) in Theorem 3 of (Allen-Zhu et al., 2019) since (Allen-Zhu et al., 2019) make the assumption that C ( ∇ L ) =
1 (see Allen-Zhu et al., 2019, Appendix A).
Let θ , θ
(cid:48) ∈
( R m × m )
L verifying the conditions of Theorem C.1.
Let A =
R ( L ( h ( θ (cid:48) ))
− R ( L ( h ( θ )) − (cid:104)∇ θ ( R ◦ L ◦ h )( θ ) , θ (cid:48) − θ (cid:105) , the quantity we want to bound.
Using (52) for z = h ( θ ) and z (cid:48)
=
h ( θ (cid:48) ) , we obtain
[Equation]
Then using the chain rule (54)
[Equation]
+ [Equation] For all i ∈ { 1 , . . . , n } , let us denote ˘ loss i := ∇ h
i
R ◦ L )( h ( θ )) to match the notations used in (Allen-Zhu et al., 2019) for the derivative of the loss with respect to the output of the network for example i of the training set.
With this notation, we obtain exactly equation (11.3) in (Allen-Zhu et al., 2019) up to the multiplicative factor (cid:16) βC
( L ) 2 nρ + 2 C ( ∇ L ) (cid:17) for the distributionally robust loss.
From there the proof of Theorem 4 in (Allen-Zhu et al., 2019) being independent to the formula for ˘ loss i , we can conclude the proof of our Theorem C.1 (as in Allen-Zhu et al., 2019, Appendix A).
C.6. G RADIENT B OUNDS FOR THE D ISTRIBUTIONALLY R OBUST L OSS
We prove the following lemma which is a generalization of Theorem 3 in (Allen-Zhu et al., 2019) for the distributionally robust loss (9).
Theorem C.2 (Gradient Bounds for the Distributionally Robust Loss) .
Let ω ∈ O (cid:16) δ 3 / 2 n 9 / 2 L 6 log 3 ( m ) (cid:17) , and θ (0) being initialized randomly as described in assumption 5.1.
With probability as least 1 − exp ( − Ω( mω 3 / 2 L )) over the initialization, we have for all θ ∈
( R m × m )
L with (cid:13) (cid:13) θ − θ (0) (cid:13)(cid:13) 2 ≤ ω ∀ i ∈ { 1 , . . . , n } , ∀ l ∈ { 1 , . . .
, L } , ∀ ˆ L ∈ R n (cid:13)(cid:13)(cid:13) ¯ p
i (ˆ L ) ∇ θ l ( L i ◦ h i )
( θ ) (cid:13)(cid:13)(cid:13) 2 2 ≤ O (cid:18) [Equation] [Equation] [Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
It is worth noting that the loss vector ˆ L used for computing the robust probabilities ¯ p (ˆ L ) = (cid:16) ¯ p i
(ˆ L ) (cid:17) n i
=1 does not have to be equal to L ( h ( θ )) .
We will use this for the proof of the Robust SGD with stale loss history.
The adaptation of the proof of Theorem 3 in (Allen-Zhu et al., 2019) is straightforward.
Let θ ∈
( R m × m )
L satisfying the conditions of Theorem C.2, and ˆ L ∈ R n .
Let us denote v := (cid:16) ¯ p i
(ˆ L ) ∇ h
i L
i
( h i ( θ )) (cid:17) n i
=1 , applying the proof of Theorem 3 in (Allen-Zhu et al., 2019) to our v gives:
[Equation]
[Table]
This allows to conclude the proof of our Theorem C.2.
C.6. C ONVERGENCE OF R OBUST SGD WITH E XACT L OSS H ISTORY
We can now prove Theorem 5.1.
Theorem C.3 (Convergence of Robust SGD with exact Loss History – Restated from Theorem 5.1) .
Suppose batch size 1 ≤ b ≤ n , number of hidden units m ≥ Ω( d(cid:15) − 1 × poly ( n, L, δ − 1 )) , and δ ≥
O (cid:0) 1 L (cid:1) .
Let (cid:15) > 0 , and the learning rate be η exact =
Θ (cid:16)
[Equation]
with probability at least 1 − exp (cid:0) − Ω(log 2 ( m )) (cid:1) over the randomness of the initialization and the mini-batches, Robust SGD with exact loss vector ﬁnds (cid:107)∇ θ
R ◦ f ◦ h )( θ ) (cid:107) ≤ (cid:15) after T =
O (cid:16) Ln 3 ηδ(cid:15) 2 (cid:17) iterations.
Similarly to the proof of the convergence of SGD for the mean loss (4) (Theorem 2 in (Allen-Zhu et al., 2019)), the convergence of SGD for the distributionally robust loss (9) will mainly rely on the semi-smoothness property (Theorem C.1) and the gradient bound (Theorem C.2) that we have proved previously for the distributionally robust loss.
Let θ ∈
( R m × m )
L satisfying the conditions of Theorem 5.1, and ˆ L be the exact loss history at θ , i.e.
[Equation]
For the batch size b ∈ { 1 , . . . , n } , let S =
{ i j } bj =1 a batch of indices drawn from ¯ p (ˆ L ) without replacement, i.e.
[Equation]
Let θ (cid:48) ∈
( R m × m )
L be the values of the parameters after a stochastic gradient descent step at θ for the batch S , i.e.
[Equation]
where η > 0 is the learning rate.
Assuming that θ and θ (cid:48) satisﬁes the conditions of Theorem C.1, we obtain
[Equation]
∞ (62) where we refer to (54) for the form of ∇ θ ( R ◦L ◦ h )( θ ) and to (53) for the form of ∇ h
( R ◦L )( h ( θ )) .
In addition, we make the assumption that for the set of values of θ considered the hardness weighted sampling probabilities admit an upper-bound
[Equation]
Which is always satisﬁed under assumption 5.2 for Kullback-Leibler φ -divergence, and for any φdivergence satisfying assumption 3.1 with a robustness parameter β small enough.
Let E S be the expectation with respect to S .
Applying E S to (62), we obtain
E S [ R ( L ( h ( θ [Equation]
[Equation]
where we have used the following results:
• For any integer k ≥ 1 , and all ( a i ) ni
=1 ∈ (cid:0) R k (cid:1) n , we have (see the proof in C.6.4)
[Equation]
• Using (65) for ( a i ) ni
=
= ( ∇ θ ( L i ◦ h i )( θ ))
ni =1 , and the chain rule (54)
[Equation]
• Using the triangular inequality
[Equation]
And
[Table]
for
[Table]
a
i ) ni
=1 =
[Equation]
[Table]
[Equation] where we have used Cauchy-Schwarz inequality for the last inequality. • Using (67) and the convexity of the function x (cid:55)→ x
[Equation]
And using (65) for ( a i ) ni
=1 = (cid:16) (cid:107)∇ θ
[Equation]
L i
[Equation]
E S 
[Equation]
Important Remark:
It is worth noting the apparition of α (63) in (70).
If we were using a uniform sampling as for ERM (i.e. for DRO in the limit β → 0 ), we would have α = 1 n .
So although our inequality (70) may seem brutal, it is consistent with equation (13.2) in (Allen-Zhu et al., 2019) and the corresponding inequality in the case of ERM.
The rest of the proof of convergence will consist in proving that η (cid:107)∇ θ
( R ◦ L ◦ h )( θ ) (cid:107) 22 , 2 dominates the two last terms in (62).
As a result, we can already state that either the robustness parameter β , or the learning rate η will have to be small enough to control α .
This is consistent with what we observed in our experiments.
Indeed, combining (62) with the chain rule (54), and the gradient bound Theorem C.2 where we use our ˆ L deﬁned in (59)
[Equation]
) where we have used
[Equation]
There are only two differences with equation (13.2) in (Allen-Zhu et al., 2019):
• in the last fraction we have n/α instead of n 2 (see remark C.6.3 for more details), and an additional multiplicative term K .
So in total, this term differs by a multiplicative factor αnK from the analogous term in the proof of (Allen-Zhu et al., 2019).
• we have (cid:80) ni
=1 (cid:13)(cid:13)(cid:13) ¯ p
i (ˆ L ) ∇ h i
L i ( h i ( θ ))
(cid:13)(cid:13)(cid:13) 2 2 instead of F ( W ( t ) ) .
In fact they are analogous since in equation (13.2) in (Allen-Zhu et al., 2019), F ( W ( t ) ) is the squared norm of the mean loss for the L 2 loss.
We don’t make such a strong assumption on the choice of L (see assumption 5.2).
It is worth noting that the same analogy is used in (Allen-Zhu et al., 2019,
Appendix A) where they extend their result to the mean loss with other objective function than the L 2 loss.
Our choice of learning rate in Theorem 5.2 can be rewritten as
[Equation]
And we also have
[Equation]
where η (cid:48) is the learning rate chosen in the proof of Theorem 2 in (Allen-Zhu et al., 2019).
We refer the reader to (Allen-Zhu et al., 2019) for the details of the constant in ” Θ ” and the exact form of the polynom poly ( n, L ) .
As a result, for η = η exact , the term Ω (cid:16) ηmδdn 2 (cid:17) dominates the other term of the right-hand side of inequality (71) as in the proof of Theorem 2 in (Allen-Zhu et al., 2019).
This implies that the conditions of Theorem C.2 are satisﬁed for all θ ( t ) , and that we have for all iteration t >
[Equation]
And using a result in Appendix A of (Allen-Zhu et al., 2019), since under assumption 5.2 the distributionally robust loss is non-convex and bounded, we obtain for all (cid:15) (cid:48) >
[Equation]
where according to (53)
[Equation]
However, we are interested in a bound on (cid:13)(cid:13)
[Equation]
rather than
[Equation]
Using the gradient bound of Theorem C.2 and the chain rules (54) and (53)
[Equation]
where c 1 > 0 is the constant hidden in O (cid:16)(cid:113) Lmnd (cid:17) .
So with (cid:15) (cid:48) = 1 c 1 (cid:113) d Lmn (cid:15) , we ﬁnally obtain
[Equation]
[Equation]
If
T
[Equation]
O
[Equation]
which concludes the proof.
C.6. P ROOF OF TECHNICAL LEMMA
For any integer k ≥ 1 , and all ( a i ) ni
=1 ∈ (cid:0) R k (cid:1) n , we have
E S (cid:34)b (cid:88) i ∈
S a i (cid:35) =
(cid:88) 1 ≤ i 1 ,...,i b ≤ n   (cid:32) n (cid:89) k
= ¯ p
i k
(ˆ L ) (cid:33)b b (cid:88)
j = a i j   =
b (cid:88) 1 ≤ i 1 ,...,i b ≤ n   b (cid:88)
j = ¯ p
i j (ˆ L ) a i j   n (cid:89) k
=1 k (cid:54) =
j ¯ p i
k (ˆ L ) 
 
 =b b (cid:88) j =   (cid:88) 1 ≤ i 1 ,...,i b ≤ n ¯ p i
j (ˆ L ) a i j   n (cid:89) k
=1 k (cid:54) =
j ¯ p i
k (ˆ L ) 
 
 =b b (cid:88)
j =     n (cid:88) i j = ¯ p i j (ˆ L ) a i j 
 n (cid:89) k =1 k (cid:54) =
j (cid:32) n (cid:88) i k
= ¯ p
i k
(ˆ L ) (cid:33)
 =b b (cid:88)
j = (cid:32) n (cid:88) i =
¯ p
i (ˆ L ) a i (cid:33) =
n (cid:88)
[Equation]
C. C ONVERGENCE OF R OBUST SGD WITH S TALE L OSS H ISTORY
The proof of the convergence of Algorithm 4.1 under the conditions of Theorem 5.2 follows the same structure as the proof of the convergence of Robust SGD with exact loss history C.6.3.
We will reuse the intermediate results of C.6.3 when possible and focus on the differences between the two proofs due to the inexactness of the loss history.
Let an iteration number t , so that the warm-up of Algorithm 4.1 is already over at t .
Let θ ( t ) ∈
( R m × m )
L the parameters of the deep neural network at iteration t .
We deﬁne the stale loss history at iteration t as
[Equation]
where for all i , t i ( t ) < t corresponds to the latest iteration before t at which the loss for example i has been updated.
Or equivalently, it corresponds to the last iteration before t when example i was drawn to be part of a mini-batch.
Thanks to the warm-up stage of Algorithm 4.1, it is guaranteed that the loss value of every example has been computed at least once before we start using the adaptive sampling.
As a result, for all iteration after the warm-up, the stale loss history ˆ L is well deﬁned.
We also deﬁne the exact loss history that is unknown in Algorithm 4.1, as
[Equation]
Remark on the warm-up stage of Algorithm 4.1: The iterations performed during the warm-up stage amounts to classic SGD to minimize the mean loss (4).
As a result, the convergence results of (Allen-Zhu et al., 2019, Theorem 2) apply during the warm-up.
This guarantees that the condition on θ of Theorem 5.2 remains satisﬁed during the warm-up if it was satisﬁed by the initial parameters.
Similarly to (61) we deﬁne
[Equation]
and using Theorem C.1, similarly to (62), we obtain
[Equation]
We can still deﬁne α as in (63)
[Equation]
86) where we are guaranteed that α > 0 under assumptions 5.1.
Since Theorem C.2 is independent to the choice of ˆ L , taking the expectation with respect to S , similarly to (71)
[Equation]
we obtain
[Equation]
where the differences with respect to (71) comes from the fact that ˆ L is not the exact loss history here,
i
[Equation]
e. ˆ L (cid:54) =
˘ L , which leads to
[Equation]
And
[Equation]
Let
[Equation]

 (90)
Where C ( L ) > 0 is a constant such that L is C ( L ) -Lipschitz continuous, and A ( ∇ L ) > 0 is a constant that bound the gradient of L with respect to its input.
We can prove that, with probability at least 1 − exp (cid:0) − Ω (cid:0) log 2 ( m ) (cid:1)(cid:1) ,
• according to lemma C
[Equation]
• according to lemma C.7.
[Equation]
• according to lemma C.7.
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Equation]
Combining those three inequalities with (87) we obtain
E S (cid:104) R
( L ( h ( θ ( t +1) ))
(cid:105)
[Equation]
R
[Equation]
L
[Equation]
h ( θ ( t ) )) ≤
[Equation]
One can see that compared to (71), there is only the additional term (cid:0)
+ md (cid:1) K (cid:48) .
Using our choice of η , η = η stale ≤ O (cid:18)
δ n 2 K (cid:48)
η exact (cid:19) (95) where η exact is the learning rate of Theorem 5.1, we have
[Equation]
As a result, η 2 (cid:0) 1 + md (cid:1) K (cid:48) is dominated by the term Ω (cid:16) ηmδdn 2 (cid:17)
In addition, since η stale ≤ η exact , Ω (cid:16) ηmδdn 2 (cid:17) still dominates also the ther terms as in the proof of
Theorem 5.1.
As a consequence, we obtain as in (75) that for any iteration t > 0 (after the end of the warm-up)
[Equation]
This concludes the proof using the same arguments as in the end of the proof of Theorem 5.1 starting from (75).
C.7. P ROOF OF TECHNICAL LEMMA
Using Lemma
[Equation]
2 and Lemma
[Equation]
1 we obtain
[Equation]
Using assumptions 5.2 and (Allen-Zhu et al., 2019, Claim 11.2)
[Equation]
Where C ( L ) is the constant of Lipschitz continuity of the per-example loss L (see assumptions 5.2) and C ( h ) is the constant of Lipschitz continuity of the deep neural network h with respect to its parameters θ .
By developing the recurrence formula of θ ( t ) (84), we obtain
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:118)(cid:117)(cid:117)(cid:117)(cid:116)
[Equation]
(cid:88)
i
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:88)
[Equation]
( t )
[Equation]
)( θ ( τ ) ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 2 ,
Let A ( ∇ L ) a bound on the gradient of the per-example loss function.
Using Theorem C.2 and the chain rule
[Equation]
And using the triangular inequality
[Equation]
As a result, we obtain
[Equation]
Therefore, for any k ≥ 1 and for any t ,
[Equation]
For k ≥ 1 b Ω (cid:18)
[Equation]
, we have (1 − α ) kb ≤ exp (cid:0) − Ω (cid:0) log 2 ( m ) (cid:1)(cid:1) , and thus with probability at least 1 − exp (cid:0) − Ω (cid:0) log
[Equation]
m )
[Equation] [Equation]
As a result, we ﬁnally obtain that with probability at least 1 − exp (cid:0) − Ω (cid:0) log 2 ( m ) (cid:1)(cid:1) ,
[Equation]
[Equation]
C.7. P ROOF OF TECHNICAL LEMMA
Let us ﬁrst denote
[Equation]
Using Cauchy-Schwarz inequality
[Equation]
( L i ◦ h i )( θ ( t ) )) , n (cid:88)
j = ¯ p j
(ˆ L ) ∇ θ ( L j ◦ h j )
( θ ( t ) ))
(cid:105) 
 (107)
Let
B = (cid:104)∇ θ
( L i ◦ h i )( θ ( t ) )) , n (cid:88)
j = ¯ p j
(ˆ L ) ∇ θ ( L j ◦ h j )
( θ ( t ) ))
(cid:105) (108)
Using again Cauchy-Schwarz inequality
[Equation]
n (cid:88)
[Equation]
[Table]
ˆ p (˘ L ) (cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) n (cid:88)
j = ¯ p
[Equation]
(ˆ L
[Equation]
( L j
[Equation]
)( θ ( t ) ))
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 2 ,(110)
Using the triangular inequality, Theorem C.2, and Lemma C.7.1, we ﬁnally obtain
[Equation]
(cid:13)n (cid:88)
j = (cid:13)(cid:13)
(cid:13)(cid:13) ¯
p j
(ˆ L ) ∇ h
[Equation]
j ( h j ( θ ( t ) ))
(cid:13)(cid:13) (cid:13)(cid:13) 2 ,≤ η m d K (cid:48) n (cid:88)
j = (cid:13)(cid:13)(cid:13)(cid:13)
[Equation] [Equation]
L
[Equation] [Equation]
( θ ( t ) ))
(cid:13)(cid:13)(cid:13)(cid:13) 2 ,(111)
C.7. P ROOF OF TECHNICAL LEMMA
We have
[Equation]
∇
[Equation]
j
[Equation]
( h
[Equation]
( θ ( t ) ))
(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
=
[Equation]
=
[Equation]
ˆ
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
h
[Equation] [Equation]
( t )
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
( h
[Equation]
( θ ( t ) ))
(cid:13)(cid:13)(cid:13)(cid:13) 2 ,(112)
Using Cauchy-Schwarz inequality
[Equation]
[Equation]
(cid:13)(cid:13) (cid:13) (cid:13)
[Equation]
Using Lemma C.7.
n
[Equation]
Therefore
[Equation]
we ﬁnally obtain
[Equation]