Year: 2019

Abstract:
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 inefficiency 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 propose SGD with hardness weighted sampling, a principled and efficient optimization method for DRO in machine learning that is particularly suited in the context of deep learning. Similar to an online hard example mining strategy in essence and in practice, the proposed algorithm is straightforward to implement and computationally as efficient as SGD-based optimizers used for deep learning. It only requires adding a softmax layer and maintaining a 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 finite number of layers and parameters. Preliminary results demonstrate the feasibility and usefulness of our approach.

heading: 1 INTRODUCTION
text: 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 {(xi, yi)} n i=1, where xi are the inputs and yi are the labels. This corresponds to Empirical Risk Minimization (ERM), defined as the non-convex optimization problem
arg min θ
1
n n∑ i=1 L (h(xi;θ), yi) (1)
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 efficiently by Stochastic Gradient Descent (SGD) with a uniform sampling over the training data (Bottou et al., 2018).
This approach has led to remarkable results in terms of average performance, but may lead to 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 solutions of ERM are prone to ignore a few hard examples in order to obtain a low mean per-example loss.
In practice, 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.
Efficient biased sampling methods, including Online Hard Example Mining (OHEM) (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 ERM formulation, it is not clear how those heuristics actually relate to ERM in theory.
Distributionally Robust Optimization (DRO) is an alternative to ERM (1) that takes into account uncertainty in the empirical data distribution. Formally, training a deep neural network with DRO corresponds to the min-max non-convex-concave optimization problem
arg min θ max q ( n∑ i=1 qi L (h(xi;θ), yi)− 1 β n∑ i=1 1 n φ (nqi) ) (2)
where φ is a convex function that defines a φ-divergence (Csiszár et al., 2004), q = (qi) n i=1 corresponds to an arbitrary sampling distribution over the training data, and β > 0 is a robustness parameter. Instead of minimizing the mean per-example loss on the training dataset, DRO seeks the hardest weighted empirical training data distribution around the (uniform) empirical training data distribution. This suggests a link between DRO and OHEM.
The parameter β allows DRO to interpolate between ERM (β → 0) and the minimization of the maximum per-example loss (β → +∞). Motivations for using the minimization of the maximum per-example loss for safety-critical applications have been discussed in (Shalev-Shwartz & Wexler, 2016).
DRO as a generalization of ERM for machine learning has been studied in (Duchi et al., 2016; Rafique et al., 2018; Namkoong & Duchi, 2016; Chouzenoux et al., 2019), but still lacks optimization methods that are computationally as efficient as SGD in the non-convex setting of deep learning.
If one could efficiently solve the inner maximization problem in (2) for a given θ, DRO could be addressed by alternating between this maximization problem and a minimization scheme akin to the standard ERM (1), but over an adaptively weighted empirical distribution. However, even when a closed-form solution is available for the inner maximization problem, it requires performing a forward pass over the entire training dataset at each iteration. This cannot be done efficiently 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 (Rafique et al., 2018; Lin et al., 2019; Jin et al., 2019). However, they differ from SGD methods for ERM 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, DRO is difficult to use as a replacement of ERM in practice. In addition, those min-max methods do not use the link between DRO and adaptive weighted sampling, therefore departing from efficient 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 deep neural networks with ReLU activation functions 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 DRO (2) and inspired by OHEM. Compared to SGD, our method only requires introducing an additional softmax layer and maintaining a 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 for DRO is computationally as efficient as SGD for ERM. In practice, we show that our method performs favorably to SGD in the case of class imbalance.
We also formally link DRO in our method with OHEM (Shrivastava et al., 2016). As a result, our method can be seen as a principled OHEM approach. In this context, the robustness parameter β controls the trade-off between exploitation and exploration in the OHEM 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 functions (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 first convergence result for deep learning networks with ReLU trained with DRO.

heading: 2 RELATED WORK IN DRO WITH A WASSERSTEIN DISTANCE
text: In this work, we focus on DRO with a φ-divergence (Csiszár 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 is DRO with a Wasserstein distance (Sinha et al., 2017; Duchi et al., 2016; Staib & Jegelka, 2017; Chouzenoux et al., 2019). In contrast to φ-divergences, 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 OHEM 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).
3 MACHINE LEARNING WITH DRO AND φ-DIVERGENCE
In machine learning based on Empirical Risk Minimization (ERM), a predictor h is trained using a training dataset {(xi,yi)}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 7→ y, and omit explicitly mentioning any parameter regularisation term.
Let ∆n ⊂ Rn be the set of empirical weighted training data distributions defined according to a given training dataset
∆n = { p = (pi) n i=1 ∈ [0, 1] n | n∑ i=1 pi = 1 } (3)
and let p̂data be the uniform empirical training data distribution, i.e. for all i, p̂data,i = 1n .
Let θ be the set of parameters of the predictor h( . ;θ) : x 7→ y we want to train, and h : θ 7→ (h(xi;θ)) n i=1 be the vector of inferred outputs from the training data. We assume L is a smooth and potentially non-convex function. We also denote L(h(θ)) = (L(h(xi;θ),yi))ni=1. Definition 3.1 (Mean Loss).
M(L(h(θ))) = Ep̂data [L (h(x;θ),y)] = 1
n n∑ i=1 L (h(xi;θ),yi) (4)
The ERM predictor, as used in most learning settings, is obtained by minimizing the mean loss (4). Definition 3.2 (Empirical Risk Minimization (ERM) predictor).
θ̃ = arg min θ
M(L(h(θ))) (5)
However, p̂data is typically biased compared to the true data distribution. Therefore, predictors trained with ERM are prone to fail on new examples that are not well represented in the training dataset.
Distributionally Robust Optimization (DRO) is an alternative to ERM that mitigates this issue by encouraging robustness to bias 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 the set of all the empirical distributions of the training dataset ∆n. Definition 3.3 (φ-divergence). Let φ : R+ → R∪{+∞} be a closed, convex, lower semi-continuous function such that ∀z ∈ R+, φ(z) ≥ φ(1) = 0. The φ-divergence Dφ is defined as, for all p = (pi) n i=1, q = (qi) n i=1 ∈ ∆n
Dφ (q‖p) = n∑ i=1 piφ ( qi pi ) (6)
Example 3.1. For φ : z 7→ z log(z), Dφ is the Kullback-Leibler (KL) divergence:
Dφ(q‖p) = DKL(q‖p) = n∑ i=1 qi log ( qi pi ) (7)
And, for φ : z 7→ (z − 1)2, Dφ is the Pearson χ2 divergence:
Dφ(q‖p) = χ2(q‖p) = n∑ i=1 (qi − pi)2 pi (8)
Definition 3.4 (Distributionally Robust Loss).
R(L(h(θ))) = max q∈∆n
Eq [L (h(x;θ),y)]− 1
β Dφ(q‖p̂data)
= max q∈∆n n∑ i=1 qi L (h(xi;θ),yi)− 1 nβ n∑ i=1 φ (nqi) (9)
where β > 0 is a hyperparameter that controls the amount of robustness.
For a given φ-divergence, we define the DRO predictor, that is obtained by minimizing the distributionally robust loss (9) instead of the mean loss (4).
Definition 3.5 (Distributionally Robust Optimization (DRO) predictor).
θ̃ = arg min θ
R(L(h(θ))) (10)
DRO interpolates between ERM as β −→ 0 and the minimization of the maximum loss as β −→ ∞, and is equivalent to a mean-variance trade-off when β −→ 0 small (Gotoh et al., 2018), i.e. max q∈∆n ( Eq [L(h (x;θ) , y)]− 1 β Dφ(q‖p̂data) ) = Ep̂data [L(h (x;θ) , y)] + β 2φ′′(1) Vp̂data [L(h (x;θ) ,y)] + o(β)
max q∈∆n
( Eq [L(h (x;θ) , y)]− 1
β Dφ(q‖p̂data) ) −−−−−→ β−→+∞ maxi L(h (xi;θ) , yi)
(11)
where Vp̂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 β > 0
∀θ, M(L(h(θ))) ≤ R(L(h(θ))) (12)
We now make assumptions for the φ-divergence to simplify the derivations of our optimization method for DRO (10) in the remainder of the paper.
Assumption 3.1 (Regularity of φ). φ : R+ → R is two times continuously differentiable on [0, n], ρ-strongly convex on [0, n], i.e.: ∃ρ > 0,∀z, z′ ∈ [0, n], φ(z′) ≥ φ(z) +φ′(z)(z′− z) + ρ2 (z− z ′)2 and satisfies (see D.1 for a justification): ∀z ∈ R, φ(z) ≥ φ(1) = 0, φ′(1) = 0.
These assumptions are verified by most φ-divergences used in practice (e.g. the KL divergence).

heading: 4 SGD WITH HARDNESS WEIGHTED SAMPLING

heading: 4.1 DISTRIBUTIONALLY ROBUST OPTIMIZATION WITH SGD AND ADAPTIVE SAMPLING
text: Existing optimization methods for DRO with a non-convex predictor h alternate between approximate minimization and maximization steps (Rafique et al., 2018; Jin et al., 2019; Lin et al., 2019), requiring the introduction of additional hyperparameters compared to SGD. These are difficult to
tune in practice and convergence has not been proven for deep neural networks with ReLU activation functions.
In this section, we highlight properties that allows us to link DRO with SGD combined with adaptive sampling. Our analysis relies on Fenchel duality (Moreau, 1965) and the notion of Fenchel conjugate (Fenchel, 1949) that we now define. Definition 4.1 (Fenchel Conjugate Function). Let f : Rm → R ∪ {+∞} be a proper function. The Fenchel conjugate of f is defined as ∀v ∈ Rm, f∗(v) = maxx∈Rm〈v,x〉 − f(x).
Let ∀p ∈ Rn, G(p) = 1
β Dφ(p‖ptrain) + δ∆n(p) (13)
where δ∆n is the characteristic function of the closed convex set ∆n, i.e. ∀p ∈ Rn, δ∆n(p) = { 0 if p ∈ ∆n +∞ otherwise (14)
One can remark that the distributionally robust loss R (9) can be rewritten using the Fenchel conjugate function of G. This allows us to obtain regularity properties for R. Lemma 4.1 (Regularity of R). If φ satisfies Assumption 3.1, then G and R satisfy the following:
G is ( nρ
β
) -strongly convex (15)
∀θ, R(L(h(θ))) = max q∈Rn
(〈L(h(θ)), q〉 −G(q)) = G∗ (L(h(θ))) (16)
R is ( β
nρ
) -gradient Lipschitz continuous. (17)
Equation (16) follows from Definition 4.1. Proofs of (15) and (17) can be found in Appendix D.3.
According to (15), the optimization problem (16) is strictly convex and admit a unique solution in ∆n. Let us denote this solution
p̄(L(h(θ))) = arg max q∈Rn (〈L(h(θ)), q〉 −G(q)) (18)
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.1 for efficient training with the distributionally robust loss. Lemma 4.2 (Stochastic Gradient of the Distributionally Robust Loss). For all θ, we have
p̄(L(h(θ))) = ∇vR(L(h(θ))) ∇θ(R ◦ L ◦ h)(θ) = Ep̄(L(h(θ))) [∇θ L (h(x;θ), y)]
(19)
where∇vR is the gradient of R with respect to its input.
The proof is found in Appendix D.4. It is apparent from (19) that, given p̄, an estimate of ∇θ(R ◦ L ◦h) 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 provide closed-form formulas for p̄ given L(h(θ)) for the KL divergence and the Pearson χ2 divergence. Example 4.1. For the KL divergence (i.e. φ : z 7→ z log(z)− z + 1), we have (see D.2 for a proof)
p̄(L(h(θ))) = softmax (β L(h(θ))) (20) And for the Pearson χ2 divergence (i.e. φ : z 7→ (z − 1)2), we have:
∀i, p̄i(L(h(θ))) = ReLU  1 n 1 + β 2 L(h(θ))i − 1 n n∑ j=1 L(h(θ))j  (21) In both cases, we can verify consistency with (11) as
∀i ∈ {1, . . . , n}, p̄i(L(h(θ))) −−−→ β−→0
1 n (22)
Algorithm 1 SGD-HWS: SGD with Hardness Weighted Sampling for Kullback-Leibler DRO 1: Input: Training data {(xi, yi)}ni=1, number of epochs T > 1, robustness parameter β > 0,
learning rate η > 0, batch size b ∈ {1, . . . , n}. 2: Initialization: 3: Initialise θ randomly 4: Initialise the loss history L̃ = −1 5: Warm start: 6: // Split the training data into batches B and run one epoch with classic SGD 7: for {(xi, yi)}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(xi;θ), yi) 11: // Run backward pass and update the parameters of the model 12: θ ← θ − η 1b ∑ i∈I ∇θ L(h(xi;θ), yi) 13: SGD with dynamic hardness weighted sampling: 14: for epoch = 2, . . . , T do 15: for iteration i = 1, . . . , (⌊n b ⌋ + 1 )
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: {(xi, yi)}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(xi;θ), yi) 23: // Run backward pass and update the parameters of the model 24: θ ← θ − η 1b ∑ i∈I ∇θ L(h(xi;θ), yi) 25: Output: θ

heading: 4.2 EFFICIENT ALGORITHM FOR DISTRIBUTIONALLY ROBUST DEEP LEARNING
text: The second equality in (19) implies that ∇θ L(hi(θ),yi) is an unbiased estimator of the distributionally robust loss gradient when i is sampled with respect to p̄(L(h(θ))). This suggests that the distributionally robust loss can be minimized efficiently by SGD 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 dataset at each iteration, is computationally prohibitive for large training 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 ( L(h(xi;θ(ti)),yi) ) . Where for all i, ti is the last iteration at which the per-example loss of example i has been computed. Using the Kullback-Leibler divergence as φ-divergence, this leads to the SGD with hardness weighted sampling algorithm proposed in Algorithm 4.1. This would also apply to the Pearson χ2 divergence mutatis mutandis.
In contrast to alternate min-max optimization methods, our SGD with an adaptive sampling strategy is similar to the SGD-based optimizers used by the vast majority of deep learning practitioners. 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.

heading: 4.3 CONVERGENCE OF SGD WITH HARDNESS WEIGHTED SAMPLING FOR
text: OVER-PARAMETERIZED DEEP NEURAL NETWORKS 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 functions (including ReLU), and with any (finite) number of layers L and parameters m,
under the assumption that m is large enough. At the time of writing, this is the most realistic setting for which a convergence theory of deep learning exists.
In this section, we extend the convergence theory developed by (Allen-Zhu et al., 2019) for ERM and SGD to DRO and the proposed SGD with hardness weighted sampling (as stated in Algorithm 4.1). Theorem 4.1 (Convergence of Algorithm 4.1 for over-parameterized neural networks with ReLU). Let L be a smooth per-sample loss function, b ∈ {1, . . . , n} be the batch size, and > 0. If m is large enough, and the learning rate is small enough, then, with high probability over the randomness of the initialization and the mini-batches, Algorithm 4.1 finds ‖∇θ(R ◦ L ◦h)(θ)‖ ≤ after a finite number of iterations.
A more detailed version of this theorem is described in B.2 and the proof can be found in D.8.

heading: 4.4 DRO AS PRINCIPLED ONLINE HARD EXAMPLE MINING
text: 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 OHEM 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 definition for OHEM sampling. Definition 4.2 (Online Hard Example Mining Sampling). Any adaptive sampling method such that the probability pi of sampling example xi is an non-decreasing function of the (potentially stale) loss value associated with xi. Theorem 4.2. The proposed hardness weighted sampling is a hard example mining sampling for any φ-divergence that satisfies Assumption 3.1. In addition, the probability pi of sampling example xi is an non-increasing function of the loss value associated with xj for all j 6= i.
See Appendix D.5 for the proof. The second part of Theorem 4.2 implies that as the loss of an example diminishes, the sampling probabilities of all the other examples increase. As a result, the proposed SGD with hardness weighted sampling balances exploitation (i.e. sampling the identified hard examples) and exploration (i.e. sampling any example to keep the record of hard examples up to date).

heading: 5 EXPERIMENTS
text: 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 for DRO (10).

heading: 5.1 ROBUSTNESS TO DOMAIN GAP
text: 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.
Comparison of ERM with SGD and DRO with our SGD with hardness weighted sampling for β = 10 at testing can be found in Figure 1. More values of β and the learning curves during training can be found in Figure 2.
Our experiment suggests that DRO and ERM lead to different optima. Indeed, DRO for β = 10 outperforms ERM by more than 10% of accuracy on the under-represented class, as illustrated in Figure 1. 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, Figure 1 suggests that DRO with our SGD with hardness weighted sampling can convergence faster than ERM with SGD.
Furthermore, the variations of learning curves with β that can be found in Figure 2 are consistent with our theoretical insight. As β decreases to 0, the learning curve of DRO with our Algorithm 4.1 converges to the learning curve of ERM with SGD.

heading: 5.2 STABILITY OF DRO
text: For large values of β (here β ≥ 10), instabilities appear in the testing learning curves at the begining of training, as illustrated in Figure 1 and in the top panels of Figure 2. For ERM this usually suggests that the learning rate is too high.
However, we observed that reducing the learning rate does not reduce those instabilities. The bottom left panel of Figure 2 shows that the training loss curves for β ≥ 10 were actually stable there. We also observe that during the iterations for which instabilities appear on the testing set, the standard deviation of the per-sample loss on the training set increases and then decreases. Following (20) the higher the standard deviation of the per-sample loss history, the more our weighted sampler focuses on hard examples. Therefore, instabilities on the testing set during training with DRO are due to the sampler focusing on hard examples.

heading: 6 CONCLUSION AND DISCUSSION
text: We have shown that efficient 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 Online Hard Example Mining (OHEM) method. It is as straightforward to implement, and as computationally efficient as SGD for Empirical Risk Minimization (ERM). It can be used for deep neural networks with any activation function (including ReLU), and with any per-example loss function. We have shown that the proposed approach can formally be described as a principled online hard example mining strategy. In addition, we prove the convergence of our method for over-parameterized deep neural networks. Thereby, extending the convergence theory of deep learning of (Allen-Zhu et al., 2019). This is, to the best of our knowledge, the first convergence result for training a deep neural network based on DRO.
Our experiments on an imbalanced MNIST dataset illustrate the practical usefulness and feasibility of our methods. SGD with hardness weighted sampling is more robust to domain gaps between the training and the testing dataset and converges faster than SGD.
However, adapting acceleration methods for ERM with SGD, like momentum updates, to DRO remains non-trivial because of the inner maximization of DRO. Investigating accelerated extension of our method is left for future work.

heading: A MORE ON OUR EXPERIMENTS

heading: A.1 EXPERIMENTS ON MNIST
text: A.1.1 IMPLEMENTATION DETAILS
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 efficient and flexible, achieving state-of-the-art accuracy on several dataset. More specifically, 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.

heading: A.1.2 COMMENT ON EARLY-STOPPING WITH ERM AND DRO
text: 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 overfitting (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.

heading: B CONVERGENCE OF SGD WITH HARDNESS WEIGHTED SAMPLING FOR
text: OVER-PARAMETERIZED DEEP NEURAL NETWORKS WITH ReLU (DETAILED STATEMENT)
In this section, we give a detailed statement of the convergence results presented in 4.3.
Our analysis is based on the results developed in (Allen-Zhu et al., 2019) which is a simplified version of (Allen-Zhu et al., 2018). Improving on those theoretical results would automatically improve our results as well. We focus on providing theoretical tools that could be used to generalize any convergence result for ERM using SGD to DRO using Algorithm 4.1. In addition, we compared our conditions and results the one obtained for ERM and SGD in (Allen-Zhu et al., 2019) and discuss the implication of those differences.
Let us first state our assumptions on the neural network h, and the per-example loss function L. Assumption B.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 hi : θ 7→ hi(xi;θ) the output d dimensional scores of h applied to example xi of dimension d.
• θ = (θl)L+1l=0 is the set of parameters of the neural network h, where θl is the set of weights for layer l with θ0 ∈ Rd×m, θL+1 ∈ Rm×d, and θl ∈ Rm×m for any other l.
• (Data separation) It exists δ > 0 such that for all i, j ∈ {1, . . . , n}, if i 6= j, ‖xi − xj‖ ≥ δ.
• We assume m ≥ Ω(d × poly(n,L, δ−1)) for some sufficiently large polynomial poly, and δ ≥ O ( 1 L ) . We refer the reader to (Allen-Zhu et al., 2019) for details about poly.
• The parameters θ = (θl)L+1l=0 are initialized at random such that:
– [θ0]i,j ∼ N ( 0, 2m ) for every (i, j) ∈ {1, . . . ,m} × {1, . . . , d}
– [θl]i,j ∼ N ( 0, 2m ) for every (i, j) ∈ {1, . . . ,m}2 and l ∈ {1, . . . , L}
– [θL+1]i,j ∼ N ( 0, 1d ) for every (i, j) ∈ {1, . . . , d} × {1, . . . ,m}
Assumption B.2 (Regularity of L). For all i, Li is a C(∇L)-gradient Lipschitz continuous, C(L)Lipschitz continuous, and bounded (potentially non-convex) function.
We first 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 B.1 (Convergence of Robust SGD with exact Loss History). Let batch size 1 ≤ b ≤ n, and > 0. Suppose there exists constants C1, C2, C3 > 0 such that the number of hidden units satisfies m ≥ C1(d −1 × poly(n,L, δ−1)), δ ≥ ( C2 L ) , and the learning rate be ηexact =
C3
( min ( 1, αn
2ρ βC(L)2+2nρC(∇L) ) × bδdpoly(n,L)m log2(m) ) . There exists constants C4, C5 > 0 such
that with probability at least 1 − exp ( −C4(log2(m)) ) over the randomness of the initialization and the mini-batches, Robust SGD with exact loss vector finds ‖∇θ(R ◦ L ◦h)(θ)‖ ≤ after T = C5 ( Ln3
ηexactδ 2
) iterations.
where α = minθ mini 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 C1, C2, C3, C4, C5 and the definitions 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 us to guarantee that ηexact ≤ η′ where η′ is the learning rate of (Allen-Zhu et al., 2019, Theorem 2). The proof can be found in Appendix D.7.3.
It is worth noting that for the KL φ-divergence, ρ = 1n . In addition, in the limit β → 0, which corresponds to ERM, we have α → 1n . As a result, we recover exactly Theorem 2 of (AllenZhu et al., 2019) as extended in their Appendix A for any L that satisfies assumption B.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 B.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 B.2 (Convergence of Robust SGD with Stale Loss History and warm-up). Let batch size 1 ≤ b ≤ n, and > 0. Under the conditions of Theorem B.1, the same notations, and with the
learning rate ηstale = C6 min ( 1, αρd3/2δb log( 11−α )
βC(L)A(∇L)Lm3/2n3/2 log2(m)
) × ηexact for a constant C6 > 0.
With probability at least 1 − exp ( −C4(log2(m)) ) over the randomness of the initialization and the mini-batches, Robust SGD with exact loss vector finds ‖∇θ(R ◦ L ◦h)(θ)‖ ≤ after T = C5 ( Ln3
ηstaleδ 2
) 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 B.1. The proof can be found in Appendix D.8.
Compared to Theorem B.1 only the learning rate differs. Similarly to Theorem B.1, when β tends to zero we recover Theorem 2 of (Allen-Zhu et al., 2019).
It is worth noting that when β increases, αρd3/2δb log( 11−α )
βC(L)A(∇L)Lm3/2n3/2 log2(m) 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.

heading: C SUMMARY OF THE NOTATIONS USED IN THE PROOFS
text: For the ease of following the proofs we first summarize our notations.

heading: C.1 PROBABILITY THEORY NOTATIONS
text: • ∆n = {p = (pi)ni=1 ∈ [0, 1]n, ∑ i pi = 1}
• Let q = (qi) ∈ ∆n, and f a function, we denote Eq[f(x)] := ∑n i=1 qif(xi).
• Let q ∈ ∆n, and f a function, we denote Vq[f(x)] := ∑n i=1 qi ‖f(xi)− Eq[f(x)]‖ 2.
• p̂data is the uniform training data distribution, i.e. p̂data = ( 1 n )n i=1 ∈ ∆n

heading: C.2 MACHINE LEARNING NOTATIONS
text: • n is the number of training examples
• d is the dimension of the output • d is the dimension of the input • training data: {(xi, yi)}ni=1, where for all i ∈ {1, . . . , n}, xi ∈ Rd and yi ∈ Rd
• h : x 7→ y is the predictor • θ is the set of parameters of the predictor • For all i, hi : θ 7→ h(xi;θ) is the output of the network for example i as a function of θ • L is the objective function • Li : hi 7→ L(hi, yi) is the objective function for example i. • By abuse of notation we also denote by L the function L : (hi)ni=1 7→ (Li(hi))ni=1 • b ∈ {1, . . . , n} is the batch size • η > 0 is the learning rate • ERM is short for Empirical Risk Minimization

heading: C.3 DISTRIBUTIONALLY ROBUST OPTIMISATION NOTATIONS
text: • Forall θ, R(L(h(θ))) = maxq∈∆n Eq [L (h(x;θ),y)] − 1βDφ(q‖p̂data) is the Distributionally Robust Loss evaluated at θ, where β > 0 is the parameter that adjusts the distributionally robustness (see (9) for more details). For short, we also used the term Robust Loss for R(L(h(θ)))
• DRO is short for Distributionally Robust Optimisation

heading: C.4 MISCELLANEOUS
text: By abuse of notation, and similarly to (Allen-Zhu et al., 2019), we use the Bachmann-Landau notations to hide constants that do not depend on our main hyper-parameters. Let f and g be two scalars, we note: {
f ≤ O(g) ⇐⇒ ∃c > 0 s.t. f ≤ cg f ≥ Ω(g) ⇐⇒ ∃c > 0 s.t. f ≥ cg f = Θ(g) ⇐⇒ ∃c1 > 0 and ∃c2 > c1 s.t. c1g ≤ f ≤ c2g

heading: D PROOFS
text: D.1 ASSUMPTIONS FOR THE φ-DIVERGENCE
Let Dφ be a φ-Divergence, and c ∈ R. One can note that for φc : z 7→ φ(z) + c (1− z), we have for all p = (pi)ni=1, q = (qi) n i=1 ∈ ∆n,
Dφc (q‖p) = n∑ i=1 piφc ( qi pi )
= n∑ i=1 pi ( φ ( qi pi ) + c ( 1− qi pi ))
= Dφ (q‖p) + c ( n∑ i=1 pi − n∑ i=1 qi ) = Dφ (q‖p)
(23)
In other words, Dφ is not uniquely defined by φ under the definition 3.3. If φ is differentiable, one can assume without loss of generality that φ′(1) = 0.
For more detailed discussion, we refer the interested reader to (Csiszár et al., 2004).

heading: D.2 PROOF OF EXAMPLE 4.1: FORMULA OF THE SAMPLING PROBABILITIES FOR THE KL DIVERGENCE
text: We give here a simple proof of the formula of the sampling probabilities for the KL divergence as φ-divergence (i.e. φ : z 7→ z log(z)− z + 1)
∀θ, p̄(L(h(θ))) = softmax (β L(h(θ))) For any θ, the distributionally robust loss (9) for the KL divergence at θ is given by
R ◦ L ◦ h(θ) = max q∈∆n ( n∑ i=1 qi L i ◦ hi(θ)− 1 β n∑ i=1 qi log (nqi) )
= max q∈∆n n∑ i=1 ( qi L i ◦ hi(θ)− 1 β qi log (nqi) ) To simplify the notations, let us denote v = (vi)ni=1 = L◦ h(θ) = (Li ◦ hi(θ)) n i=1, and p̄ = (p̄i) n i=1 = p̄(L(h(θ))).
Thus p̄(L(h(θ))) is, by definition, solution of the optimization problem
arg max q∈∆n n∑ i=1 ( qivi − 1 β qi log (nqi) ) (24)
First, let us remark that the function q 7→ ∑n i=1 qi log (nqi) is strictly convex on the non empty closed convex set ∆n as a sum of strictly convex functions. This implies that the optimization (24) has a unique solution and as a result p̄(L(h(θ))) is well defined. We now reformulate the optimization problem (24) as a convex smooth constrained optimization problem by writing the condition q ∈ ∆n as constraints.
arg max q∈Rn+ n∑ i=1 ( qivi − 1 β qi log (nqi) )
s.t. n∑ i=1 qi = 1
(25)
There exists a Lagrange multiplier λ ∈ R, such that the solution p̄ of (25) is characterized by
∀i ∈ {1, . . . , n}, vi − 1
β (log (np̄i) + 1) + λ = 0
n∑ i=1 p̄i = 1 (26)
Which we can rewrite as
∀i ∈ {1, . . . , n}, p̄i = 1
n exp (β (vi + λ)− 1)
1
n n∑ i=1 exp (β (vi + λ)− 1) = 1 (27)
The last equality gives exp (βλ− 1) = n∑n
i=1 exp (βvi)
And by replacing in the formula of the p̄i
∀i ∈ {1, . . . , n}, p̄i = 1
n exp (βvi) exp (βλ− 1)
= exp (βvi)∑n j=1 exp (βvj)
Which corresponds exactly to p̄ = softmax (βv)

heading: D.3 PROOF OF LEMMA 4.1: REGULARITY PROPERTIES OF R
text: For the ease of reading, let us first recall that given a φ-Divergence that satisfies assumptions 3.1, we have defined in (9)
R : Rn → R
v 7→ max q∈∆n ∑ i qivi − 1 β Dφ(q‖ptrain) (28)
And in (13) G : Rn → R
p 7→ 1 β Dφ(p‖ptrain) + δ∆n(p)
(29)
where δ∆n is the characteristic function of the closed convex set ∆n, i.e. ∀p ∈ Rn, δ∆n(p) = { 0 if p ∈ ∆n +∞ otherwise (30)
We now prove Lemma 4.1 on the regularity of R. Lemma D.1 (Regularity of R – Restated from Lemma 4.1). Let φ that satisfies Assumption 3.1, G and R satisfy
G is ( nρ
β
) -strongly convex (31)
R(L(h(θ))) = max q∈Rn
(〈L(h(θ)), q〉 −G(q)) = G∗ (L(h(θ))) (32)
R is ( β
nρ
) -gradient Lipschitz continuous. (33)
φ is ρ-strongly convex on [0, n] so
∀x, y ∈ [0, n]2,∀λ ∈ [0, 1], φ (λx+ (1− λ)y) ≤ λφ(x) + (1−λ)φ(y)− ρλ(1− λ) 2 |y−x|2 (34)
Let p = (pi) n i=1, q = (qi) n i=1 ∈ ∆n, and λ ∈ [0, 1], using (34) and the convexity of δ∆n , we obtain:
G (λp+ (1− λ)q) = 1 βn n∑ i=1 φ (nλpi + n(1− λ)qi) + δ∆n (λp+ (1− λ)q)
≤ λG(p) + (1− λ)G(q)− 1 βn n∑ i=1 ρλ(1− λ) 2 |nqi − npi|2
≤ λG(p) + (1− λ)G(q)− nρ β λ(1− λ) 2 ‖q − p‖2
(35)
This proves that G is nρβ -strongly convex.
Since G is convex, R = G∗ is also convex, and R∗ = (G∗)∗ = G (Hiriart-Urruty & Lemaréchal, 2013).
We obtain (32) using Definition 4.1.
We now show that R is Frechet differentiable on Rn. Let v ∈ Rn. 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).
arg max q∈Rn
(〈v, q〉 −G(q)) (36)
In addition p̂ ∈ ∆n solution of (36) ⇐⇒ 0 ∈ v − ∂G(p̂)
⇐⇒ v ∈ ∂G(p̂) ⇐⇒ p̂ ∈ ∂G∗(v) ⇐⇒ p̂ ∈ ∂R(v)
where we have used (Hiriart-Urruty & Lemaréchal, 2013, Proposition 6.1.2 p.39) for the third equivalence, and (32) for the last equivalence.
As a result, ∂R(v) = {p̂(v)}. this implies that R admit a gradient at v, and
∇vR(v) = p̂(v) (37)
Since this holds for any v ∈ Rn, we deduce that R is Frechet differentiable on Rn.
We are now ready to show that R is βnρ -gradient Lipchitz continuous by using the following lemma (Hiriart-Urruty & Lemaréchal, 2013, Theorem 6.1.2 p.280).
Lemma D.2. A necessary and sufficient condition for a convex function f : Rn → R to be cstrongly convex on a convex set C is that for all x1, x2 ∈ C
〈s2 − s1, x2 − x1〉 ≥ c ‖x2 − x1‖2 for all si ∈ ∂f(xi), i = 1, 2.
Using this lemma for f = G, c = nρβ , and C = ∆n, we obtain:
For all p1, p2 ∈ ∆n, for all v1 ∈ ∂G(p1), v2 ∈ ∂G(p2),
〈v2 − v1, p2 − p1〉 ≥ nρ
β ‖p2 − p1‖2
In addition, for i ∈ {1, 2}, vi ∈ ∂G(pi)⇐⇒ pi ∈ ∂R(v1) = {∇vR(vi)}. And using Cauchy Schwarz inequality
‖v2 − v1‖ ‖p2 − p1‖ ≥ 〈v2 − v1, p2 − p1〉
We conclude that nρ
β ‖∇vR(v2)−∇vR(v1)‖ ≤ ‖v2 − v1‖
Which implies that R is βnρ -gradient Lipchitz continuous.

heading: D.4 PROOF OF LEMMA 4.2: FORMULA OF THE DISTRIBUTIONALLY ROBUST LOSS GRADIENT
text: We prove Lemma 4.2 that we restate here for the ease of reading.
Lemma D.3 (Stochastic Gradient of the Distributionally Robust Loss – Restated from Lemma 4.2). For all θ, we have
p̄(L(h(θ))) = ∇vR(L(h(θ))) (38) ∇θ(R ◦ L ◦ h)(θ) = Ep̄(L(h(θ))) [∇θ L (h(x;θ), y)] (39)
where∇vR is the gradient of R with respect to its input. For a given θ, equality (38) is a special case of (37) for v = L(h(θ)). Then using the chain rule and (38),
∇θ(R ◦ L ◦ h)(θ) = n∑ i=1 ∂R ∂vi (L◦ h(θ)))∇θ(L i ◦ hi)(θ)
= n∑ i=1 p̄i(L(h(θ)))∇θ(L i ◦ hi)(θ) = Ep̄(L(h(θ))) [∇θ L (h(x;θ), y)]

heading: D.5 PROOF OF THEOREM 4.2: DISTRIBUTIONALLY ROBUST OPTIMIZATION AS PRINCIPLED HARD EXAMPLE MINING
text: Let Dφ an φ-divergence satisfying Assumption 3.1, and v = (vi) n i=1 ∈ Rn. 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:
max p=(pi) n i=1∈∆n 〈v, p〉 − 1 βn n∑ i=1 φ(npi) (40)
Making the constraints associated with p ∈ ∆n explicit, this can be rewritten as
max p=(pi) n i=1∈Rn 〈v, p〉 − 1 βn n∑ i=1 φ(npi)
s.t. ∀i ∈ {1, . . . , n}, pi ≥ 0
s.t. n∑ i=1 pi = 1
(41)
There exists KKT multipliers λ ∈ R and ∀i, µi ≥ 0 such that the solution p̄ = (p̄i)ni=1 satisfies: ∀i ∈ {1, . . . , n}, vi − 1 β φ′(np̄i) + λ− µi = 0 ∀i ∈ {1, . . . , n}, µipi = 0 ∀i ∈ {1, . . . , n}, pi ≥ 0
n∑ i=1 p̄i = 1
(42)
Since φ is continuously differentiable and strongly convex, we have (φ′)−1 = (φ∗)′, where φ∗ is the Fenchel conjugate of φ (see Hiriart-Urruty & Lemaréchal, 2013, Proposition 6.1.2). As a result, (42) can be rewritten has: ∀i ∈ {1, . . . , n}, p̄i = 1 n (φ∗) ′ (β(vi + λ− µi)) ∀i ∈ {1, . . . , n}, µipi = 0 ∀i ∈ {1, . . . , n}, pi ≥ 0 1
n n∑ i=1 (φ∗) ′ (β(vi + λ− µi)) = 1
(43)
We now show that the KKT multipliers are uniquely defined.

heading: The µi’s are uniquely defined by v and λ:
text: Since ∀i ∈ {1, . . . , n}, µipi = 0, pi ≥ 0 and µi ≥ 0, for all ∀i ∈ {1, . . . , n}, either pi = 0 or µi = 0. In the case pi = 0, using (43) it comes (φ∗) ′ (β(vi + λ− µi)) = 0.
According to assumption 3.1, φ is strongly convex and continuously differentiable, so φ′ and (φ∗)′ = (φ′)−1 are continuous and strictly increasing functions. As a result, it exists a unique µi (dependent to v and λ) such that:
(φ∗) ′ (β(vi + λ− µi)) = 0
And (43) can be rewritten as: ∀i ∈ {1, . . . , n}, p̄i = ReLU ( 1 n (φ∗) ′ (β(vi + λ)) ) = 1 n ReLU ( (φ∗) ′ (β(vi + λ)) ) 1
n n∑ i=1 ReLU ( (φ∗) ′ (β(vi + λ)) ) = 1
(44)
λ is uniquely defined by v and a continuous function of v: Let λ ∈ R that satisfies (44). We have 1n ∑n i=1 ReLU ( (φ∗) ′ (β(vi + λ)) ) = 1. So there exists at least one index i0 such that
ReLU ( (φ∗)
′ (β(vi0 + λ))
) = (φ∗)
′ (β(vi0 + λ)) ≥ 1
Since (φ∗)−1 is continuous and striclty increasing, λ′ 7→ ReLU ( (φ∗)
′ (β(vi0 + λ
′)) )
is continuous and strictly increasing on a neighborhood of λ.
In addition ReLU is continuous and increasing, so for all i ∈ {1, . . . , n}, λ′ 7→ ReLU ( (φ∗) ′ (β(vi + λ ′)) ) is a continuous and increasing function.
As a result, λ′ 7→ 1n ∑n i=1 ReLU ( (φ∗) ′ (β(vi + λ ′)) )
is a continuous function that is increasing on R, and strictly increasing on a neighborhood of λ.
This implies that λ is uniquely defined by v, and that v 7→ λ(v) is continuous.

heading: Hard Example Mining Sampling:
text: For any pseudo loss vector v = (vi) n i=1 ∈ Rn, there exists a unique λ and a unique p̄ that satisfies (44), so we can define the mapping:
p̄ : Rn → ∆n v 7→ p̄(v;λ(v)) (45)
where for all v, λ(v) is the unique λ ∈ R satisfying (44). We will now demonstrate that each p̄i0(v) for i0 ∈ {1, . . . , n} is an increasing function of vi and a decreasing function of the vi for i 6= i0. Without loss of generality we assume i0 = 1. Let v = (vi) n i=1 ∈ Rn, and > 0.
Let us define v′ = (v′i) n i=1 ∈ Rn, such that v′1 = v1 + and ∀i ∈ {2, . . . , n}, v′i = vi.
Similarly as in the proof of the uniqueness of λ above, we can show that there exists η > 0 such that the function
F : λ′ 7→ 1 n n∑ i=1 ReLU ( (φ∗) ′ (β(vi + λ ′)) )
is continuous and strictly increasing on [λ(v)− η, λ(v) + η], and F (λ(v)) = 1. v 7→ λ(v) is continuous, so for small enough λ(v′) ∈ [λ(v)− η, λ(v) + η]. Let us now prove by contradiction that λ(v′) ≤ λ(v). Therefore, let us assume that λ(v′) > λ(v). Then, as ReLU ◦ (φ∗)′ is an increasing function and F is strictly increasing on [λ(v)− η, λ(v) + η], and > 0 we obtain
1 = 1
n n∑ i=1 ReLU ( (φ∗) ′ (β(v′i + λ(v ′))) )
≥ 1 n n∑ i=1 ReLU ( (φ∗) ′ (β(vi + λ(v ′))) ) ≥ F (λ(v′)) > F (λ(v))
> 1
which is a contradiction. As a result
λ(v′) ≤ λ(v) (46)
Using (46), (44), and the fact that ReLU ◦ (φ∗)′ is an increasing function, we obtain for all i ∈ {2, . . . , n}
p̄i(v ′) =
1 n ReLU
( (φ∗)
′ (β(v′i + λ(v
′))) )
= 1
n ReLU
( (φ∗) ′ (β(vi + λ(v ′))) )
≤ 1 n
ReLU ( (φ∗) ′ (β(vi + λ(v))) ) ≤ p̄i(v)
(47)
In addition n∑ i=1 p̄i(v ′) = 1 = n∑ i=1 p̄i(v)
So necessarily p̄1(v ′) ≥ p̄1(v) (48)
This holds for any i0 and any v, which concludes the proof.

heading: D.6 PROOF THAT R ◦ L IS ONE-SIDED GRADIENT LIPCHITZ
text: 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 D.1.
Under assumption 3.1, we have shown that R∗ is βnρ -gradient Lipchitz continuous. And under assumption B.2, for all i, Li isC(L)-Lipschitz continuous andC(∇L)-gradient Lipschitz continuous. Let z = (zi)ni=1, z ′ = (z′i) n i=1 ∈ Rdn.
We want to show that R ◦ L is one-sided gradient Lipschitz, i.e. we want to prove the existence of a constant C > 0, independent to z and z′, such that:
〈∇z(R ◦ L)(z)−∇z(R ◦ L)(z′), z − z′〉 ≤ C ‖z − z′‖ 2
We have
〈∇z(R ◦ L)(z)−∇z(R ◦ L)(z′), z − z′〉
= n∑ i=1 〈∇zi(R ◦ L)(z)−∇zi(R ◦ L)(z′), zi − z′i〉
= n∑ i=1 〈p̄i(L(z))∇zi L i (zi)− p̄i(L(z′))∇zi L i (z′i), zi − z′i〉
= n∑ i=1 p̄i(L(z))〈∇zi L i (zi)−∇zi L i (z′i), zi − z′i〉
+ n∑ i=1 (p̄i(L(z))− p̄i(L(z′))) 〈∇zi L i (z′i), zi − z′i〉
(49)
Where for all i ∈ {1, . . . , n} we have used the chain rule
∇zi(R ◦ L)(z) = n∑ j=1 ∂R∗ ∂vj (L(z))∇zi L j (zj) = p̄i(L(z))∇zi L i (zi)
Let
A = ∣∣∣∣∣ n∑ i=1 p̄i(L(z))〈∇zi L i (zi)−∇zi L i (z′i), zi − z′i〉 ∣∣∣∣∣
For all i, Li is C(∇L)-gradient Lipchitz continuous, so using Cauchy-Schwarz inequality
A ≤ n∑ i=1 C(∇L) ‖zi − z′i‖ 2 = C(∇L) ‖z − z′‖2 (50)
Let
B = ∣∣∣∣∣ n∑ i=1 (p̄i(L(z))− p̄i(L(z′))) 〈∇zi L i (z′i), zi − z′i〉 ∣∣∣∣∣ Using the triangular inequality:
B ≤ ∣∣∣∣∣ n∑ i=1 (p̄i(L(z))− p̄i(L(z′))) (L i (zi)− L i (z′i) ∣∣∣∣∣ +
∣∣∣∣∣ n∑ i=1 (p̄i(L(z))− p̄i(L(z′))) (L i (z′i) + 〈∇zi L i (z′i), zi − z′i〉 − L i (zi) ∣∣∣∣∣ ≤ 〈∇(R∗)(L(z))−∇(R∗)(L(z′)),L(z)− L(z′)〉
+ 2 n∑ i=1 ∣∣∣L i (z′i) + 〈∇zi L i (z′i), zi − z′i〉 − L i (zi) ∣∣∣
≤ β nρ ‖L(z)− L(z′)‖2 + 2C(∇L) 2 ‖z − z′‖2
≤ ( βC(L)2
nρ + C(∇L)
) ‖z − z′‖2
(51)
Combining (49), (50) and (51) we finally obtain: 〈∇z(R ◦ L)(z)−∇z(R ◦ L)(z′), z − z′〉 ≤ ( βC(L)2
nρ + 2C(∇L)
) ‖z − z′‖2 (52)
From there, we can obtain the following inequality that will be used for the proof of the semismoothness property in Theorem D.1:
R(L(z′))−R(L(z))− 〈∇z(R ◦ L)(z), z′ − z〉
= ∫ 1 t=0 〈∇z(R ◦ L) (z + t(z′ − z))−∇z(R ◦ L)(z), z′ − z〉dt
≤ 1 2
( βC(L)2
nρ + 2C(∇L)
) ‖z − z′‖2
(53)

heading: D.7 PROOF OF THE CONVERGENCE OF ROBUST SGD
text: In this part, we prove the results of Therem B.1 and B.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, 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:
∇h(R ◦ L)(h(θ)) = (∇hi(R ◦ L)(h(θ))) n i=1
∀i ∈ {1, . . . n}, ∇hi(R ◦ L)(h(θ)) = n∑ j=1 ∂R ∂vj (L(h(θ)))∇hi L j (hj(θ))
= p̄i(L(h(θ)))∇hi L i (hi(θ))
(54)
Chain rule for the derivative of R ◦ L ◦h with respect to the network parameters θ:
∇θ(R ◦ L ◦h)(θ) = n∑ i=1 ∇θhi(θ)∇hi(R ◦ L)(h(θ))
= n∑ i=1 p̄i(L(h(θ)))∇θhi(θ)∇hi L i (hi(θ))
= n∑ i=1 p̄i(L(h(θ))∇θ(L i ◦hi)(θ))
(55)
where for all i ∈ {1, . . . n}, ∇θhi(θ) is the transpose of the Jacobian matrix of hi as a function of θ.

heading: D.7.1 SEMI-SMOOTHNESS PROPERTY FOR THE DISTRIBUTIONALLY ROBUST LOSS
text: 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 D.1 (Semi-smoothness of the Distributionally Robust Loss). Let ω ∈ [ Ω ( d3/2
m3/2L3/2 log3/2(m)
) , O (
1 L4.5 log3(m)
)] , and the θ(0) being initialized randomly as de-
scribed in assumption B.1. With probability as least 1− exp (−Ω(mω3/2L)) over the initialization, we have for all θ,θ′ ∈ (Rm×m)L with ∥∥θ − θ(0)∥∥ 2 ≤ ω, and ‖θ − θ′‖2 ≤ ω
R(L(h(θ′)) ≤ R(L(h(θ)) + 〈∇θ(R ◦ L ◦h)(θ),θ′ − θ〉
+ ‖∇h(R ◦ L)(h(θ))‖2,1O
( L2ω1/3 √ m log(m)√ d ) ‖θ′ − θ‖2,∞
+O
(( βC(L)2
nρ + 2C(∇L)
) nL2m
d
) ‖θ′ − θ‖22,∞
(56)
where for all layer l ∈ {1, . . . , L}, θl is the vector of parameters for layer l, and
‖θ′ − θ‖2,∞ = max l ‖θ′l − θl‖2
‖θ′ − θ‖22,∞ = (
max l ‖θ′l − θl‖ 2 2
)2 = max
l ‖θ′l − θl‖ 2 2
‖∇h(R ◦ L)(h(θ))‖2,1 = n∑ i=1 ‖∇hi(R ◦ L)(h(θ))‖2
= n∑ i=1 ∥∥∥p̄i(L(h(θ)))∇hi L i (hi(θ)) ∥∥∥ 2 (chain rule (54))
To compare this semi-smoothness result to the one in (Allen-Zhu et al., 2019, Theorem 4), let us first remark that
‖∇h(R ◦ L)(h(θ))‖2,1 ≤ √ n ‖∇h(R ◦ L)(h(θ))‖2,2
As a result, our result is analogous to (Allen-Zhu et al., 2019, Theorem 4), up to an additional multiplicative factor ( βC(L)2 nρ + 2C(∇L) ) 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 θ,θ′ ∈ (Rm×m)L verifying the conditions of Theorem D.1. Let A = R(L(h(θ′))−R(L(h(θ))− 〈∇θ(R ◦ L ◦h)(θ),θ′ − θ〉 , the quantity we want to bound.
Using (53) for z = h(θ) and z′ = h(θ′), we obtain
A ≤ 1 2
( βC(L)2
nρ + 2C(∇L)
) ‖h(θ′)− h(θ)‖22
+ 〈∇h(R ◦ L)(h(θ)), h(θ′)− h(θ)〉 − 〈∇θ(R ◦ L ◦h)(θ),θ′ − θ〉
(57)
Then using the chain rule (55)
A ≤ 1 2
( βC(L)2
nρ + 2C(∇L)
) ‖h(θ′)− h(θ)‖22
+ n∑ i=1 〈∇hi(R ◦ L)(h(θ)), hi(θ′)− hi(θ)− (∇θhi(θ)) T (θ′ − θ)〉 (58)
For all i ∈ {1, . . . , n}, let us denote ˘lossi := ∇hi(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 ( βC(L)2 nρ + 2C(∇L) ) for the distributionally robust loss.
From there the proof of Theorem 4 in (Allen-Zhu et al., 2019) being independent to the formula for ˘lossi, we can conclude the proof of our Theorem D.1 (as in Allen-Zhu et al., 2019, Appendix A).

heading: D.7.2 GRADIENT BOUNDS FOR THE DISTRIBUTIONALLY ROBUST LOSS
text: 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 D.2 (Gradient Bounds for the Distributionally Robust Loss). Let ω ∈ O ( δ3/2
n9/2L6 log3(m)
) , and θ(0) being initialized randomly as described in assump-
tion B.1. With probability as least 1 − exp (−Ω(mω3/2L)) over the initialization, we have for all θ ∈ (Rm×m)L with ∥∥θ − θ(0)∥∥ 2 ≤ ω
∀i ∈ {1, . . . , n}, ∀l ∈ {1, . . . , L}, ∀L̂ ∈ Rn∥∥∥p̄i(L̂)∇θl(L i ◦hi)(θ) ∥∥∥2 2 ≤ O ( m d ∥∥∥p̄i(L̂)∇hi L i (hi(θ)) ∥∥∥2 2 ) ∀l ∈ {1, . . . , L}, ∀L̂ ∈ Rn∥∥∥∥∥ n∑ i=1 p̄i(L̂)∇θl(L i ◦hi)(θ) ∥∥∥∥∥ 2
2
≤ O
( mn
d n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ)) ∥∥∥2 2 ) ∥∥∥∥∥ n∑ i=1 p̄i(L̂)∇θL(L i ◦hi)(θ) ∥∥∥∥∥ 2
2
≥ Ω
( mδ
dn2 n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ)) ∥∥∥2 2
) (59)
It is worth noting that the loss vector L̂ used for computing the robust probabilities p̄(L̂) =( p̄i(L̂) )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 θ ∈ (Rm×m)L satisfying the conditions of Theorem D.2, and L̂ ∈ Rn.
Let us denote v := ( p̄i(L̂)∇hi Li(hi(θ)) )n i=1
, applying the proof of Theorem 3 in (Allen-Zhu et al., 2019) to our v gives:
∀i ∈ {1, . . . , n}, ∀l ∈ {1, . . . , L},∥∥∥p̄i(L̂)∇θl(L i ◦hi)(θ) ∥∥∥2 2 ≤ O ( m d ∥∥∥p̄i(L̂)∇hi L i (hi(θ)) ∥∥∥2 2 ) ∀l ∈ {1, . . . , L}, ∀L̂ ∈ Rn∥∥∥∥∥ n∑ i=1 p̄i(L̂)∇θl(L i ◦hi)(θ) ∥∥∥∥∥ 2
2
≤ O
( mn
d n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ)) ∥∥∥2 2 ) ∥∥∥∥∥ n∑ i=1 p̄i(L̂)∇θL(L i ◦hi)(θ) ∥∥∥∥∥ 2
2
≥ Ω ( mδ
dn max i (∥∥∥p̄i(L̂)∇hi L i (hi(θ)) ∥∥∥2 2 ))
In addition
max i (∥∥∥p̄i(L̂)∇hi L i (hi(θ)) ∥∥∥2 2 ) ≥ 1 n n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ)) ∥∥∥2 2
This allows us to conclude the proof of our Theorem D.2.

heading: D.7.3 CONVERGENCE OF ROBUST SGD WITH EXACT LOSS HISTORY
text: We can now prove Theorem B.1.
Theorem D.3 (Convergence of Robust SGD with exact Loss History – Restated from Theorem B.1). Suppose batch size 1 ≤ b ≤ n, number of hidden units m ≥ Ω(d −1 × poly(n,L, δ−1)), and δ ≥ O (
1 L ) . Let > 0, and the learning rate be ηexact = Θ ( αn2ρ βC(L)2+2nρC(∇L) × bδd poly(n,L)m log2(m) ) ,
with probability at least 1 − exp ( −Ω(log2(m)) ) over the randomness of the initialization and the
mini-batches, Robust SGD with exact loss vector finds ‖∇θ(R ◦ f ◦ h)(θ)‖ ≤ after T = O ( Ln3
ηδ 2 ) 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 D.1) and the gradient bound (Theorem D.2) that we have proved previously for the distributionally robust loss.
Let θ ∈ (Rm×m)L satisfying the conditions of Theorem B.1, and L̂ be the exact loss history at θ, i.e.
L̂ = ( L i (hi(θ)) )n i=1
(60)
For the batch size b ∈ {1, . . . , n}, let S = {ij}bj=1 a batch of indices drawn from p̄(L̂) without replacement, i.e.
∀j ∈ {1, . . . b}, ij i.i.d.∼ p̄(L̂) (61)
Let θ′ ∈ (Rm×m)L be the values of the parameters after a stochastic gradient descent step at θ for the batch S, i.e.
θ′ = θ − η 1 b ∑ i∈S ∇θ(L i ◦hi)(θ) (62)
where η > 0 is the learning rate.
Assuming that θ and θ′ satisfies the conditions of Theorem D.1, we obtain
R(L(h(θ′)) ≤R(L(h(θ))− η〈∇θ(R ◦ L ◦h)(θ), 1
b ∑ i∈S ∇θ(L i ◦hi)(θ)〉
+ η √ n ‖∇h(R ◦ L)(h(θ))‖2,2O
( L2ω1/3 √ m log(m)√ d )∥∥∥∥∥1b∑ i∈S ∇θ(L i ◦hi)(θ) ∥∥∥∥∥ 2,∞
+ η2O
(( βC(L)2
nρ + 2C(∇L)
) nL2m
d )∥∥∥∥∥1b∑ i∈S ∇θ(L i ◦hi)(θ) ∥∥∥∥∥ 2
2,∞ (63)
where we refer to (55) for the form of∇θ(R◦L ◦h)(θ) and to (54) 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
α = min θ min i p̄i(L(θ)) > 0 (64)
Which is always satisfied under assumption B.2 for Kullback-Leibler φ-divergence, and for any φ-divergence satisfying assumption 3.1 with a robustness parameter β small enough.
Let ES be the expectation with respect to S. Applying ES to (63), we obtain ES [R(L(h(θ′))] ≤R(L(h(θ))− η ‖∇θ(R ◦ L ◦h)(θ)‖22,2
+ η ‖∇h(R ◦ L)(h(θ))‖2,2O
( nL2ω1/3 √ m log(m)√ d )√√√√ n∑ i=1 max l ∥∥∥p̄i(L̂)∇θl(L i ◦hi)(θ) ∥∥∥2 + η2O (( βC(L)2
nρ + 2C(∇L)
) nL2m
d
) 1
α n∑ i=1 max l ∥∥∥p̄i(L̂)∇θl(L i ◦hi)(θ) ∥∥∥2 (65)
where we have used the following results: • For any integer k ≥ 1, and all (ai)ni=1 ∈ ( Rk )n , we have (see the proof in D.7.4)
ES
[ 1
b ∑ i∈S ai ] = Ep̄(L̂) [ai] (66)
• Using (66) for (ai)ni=1 = (∇θ(Li ◦hi)(θ)) n i=1, and the chain rule (55)
ES
[ 1
b ∑ i∈S ∇θ(L i ◦hi)(θ)
] =
n∑ i=1 p̄i(L̂)∇θ(L i ◦hi)(θ) = ∇θ(R ◦ L ◦h)(θ) (67)
• Using the triangular inequality∥∥∥∥∥1b∑ i∈S ∇θ(L i ◦hi)(θ) ∥∥∥∥∥ 2,∞ ≤ 1 b ∑ i∈S ∥∥∥∇θ(L i ◦hi)(θ) ∥∥∥ 2,∞
(68)
And using (66) for (ai) n i=1 = ( ‖∇θ(Li ◦hi)(θ)‖2,∞ )n i=1 ,
ES ∥∥∥∥∥1b∑ i∈S ∇θ(L i ◦hi)(θ) ∥∥∥∥∥ 2,∞  ≤ n∑ i=1 p̄i(L̂) ∥∥∥∇θ(L i ◦hi)(θ) ∥∥∥ 2,∞
≤ n∑ i=1 max l ∥∥∥∇θl(p̄i(L̂)L i ◦hi)(θ) ∥∥∥ 2
≤ √ n √√√√ n∑ i=1 max l ∥∥∥∇θl(p̄i(L̂)L i ◦hi)(θ) ∥∥∥2 2
(69)
where we have used Cauchy-Schwarz inequality for the last inequality. • Using (68) and the convexity of the function x 7→ x2∥∥∥∥∥1b∑
i∈S ∇θ(L i ◦hi)(θ)
∥∥∥∥∥ 2
2,∞
≤ 1 b ∑ i∈S ∥∥∥∇θ(L i ◦hi)(θ) ∥∥∥2 2,∞
(70)
And using (66) for (ai) n i=1 = ( ‖∇θ(Li ◦hi)(θ)‖22,∞ )n i=1 ,
ES ∥∥∥∥∥1b∑ i∈S ∇θ(L i ◦hi)(θ) ∥∥∥∥∥ 2
2,∞
 ≤ n∑ i=1 p̄i(L̂) ∥∥∥∇θ(L i ◦hi)(θ) ∥∥∥2 2,∞
≤ n∑ i=1 1 p̄i(L̂) max l ∥∥∥∇θl(p̄i(L̂)L i ◦hi)(θ) ∥∥∥2 2
≤ 1 α n∑ i=1 max l ∥∥∥∇θl(p̄i(L̂)L i ◦hi)(θ) ∥∥∥2 2
(71)
Important Remark: It is worth noting the apparition of α (64) in (71). If we were using a uniform sampling as for ERM (i.e. for DRO in the limit β → 0), we would have α = 1n . So although our inequality (71) 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 η ‖∇θ(R ◦ L ◦h)(θ)‖22,2 dominates the two last terms in (63). 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 (63) with the chain rule (55), and the gradient bound Theorem D.2 where we use our L̂ defined in (60)
ES [R(L(h(θ′))] ≤ R(L(h(θ))− Ω ( ηmδ
dn2 ) n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ)) ∥∥∥2 2
+ ηO
( nL2ω1/3 √ m log(m)√ d ) O (√ m d ) n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ)) ∥∥∥2 2
+ η2O
(( βC(L)2
nρ + 2C(∇L)
) nL2m
d ) O (m dα ) n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ)) ∥∥∥2 2
≤ R(L(h(θ))− Ω ( ηmδ
dn2 ) n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ)) ∥∥∥2 2
+O
( ηnL2mω1/3 √ log(m)
d +K
η2(n/α)L2m2
d2
) n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ)) ∥∥∥2 2
(72) where we have used
K := βC(L)2
nρ + 2C(∇L) (73)
There are only two differences with equation (13.2) in (Allen-Zhu et al., 2019):
• in the last fraction we have n/α instead of n2 (see remark D.7.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 ∑n i=1 ∥∥∥p̄i(L̂)∇hi Li(hi(θ))∥∥∥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 L2 loss. We don’t make such a strong assumption on the choice of L (see assumption B.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 L2 loss.
Our choice of learning rate in Theorem B.2 can be rewritten as
ηexact = Θ
( αn2ρ
βC(L)2 + 2nρC(∇L) × bδd poly(n,L)m log2(m) ) = Θ ( αn
K × bδd poly(n,L)m log2(m) ) ≤ αn
K × η′
(74)
And we also have ηexact ≤ η′ (75) where η′ 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 Ω ( ηmδ dn2 ) dominates the other term of the right-hand side of inequality (72) as in the proof of Theorem 2 in (Allen-Zhu et al., 2019).
This implies that the conditions of Theorem D.2 are satisfied for all θ(t), and that we have for all iteration t > 0
ESt [ R(L(h(θ(t+1))) ] ≤ R(L(h(θ(t)))− Ω ( ηmδ
dn2 ) n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ (t))) ∥∥∥2 2 (76)
And using a result in Appendix A of (Allen-Zhu et al., 2019), since under assumption B.2 the distributionally robust loss is non-convex and bounded, we obtain for all ′ > 0∥∥∥∇h(R ◦ L)(h(θ(T )))∥∥∥
2,2 ≤ ′ if T = O
( dn2
ηδm ′2
) (77)
where according to (54)∥∥∥∇h(R ◦ L)(h(θ(T )))∥∥∥ 2,2 = n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ (t))) ∥∥∥2 2 (78)
However, we are interested in a bound on ∥∥∇θ(R ◦ L ◦h)(θ(T ))∥∥2,2, rather than∥∥∇h(R ◦ L)(h(θ(T )))∥∥2,2.
Using the gradient bound of Theorem D.2 and the chain rules (55) and (54)∥∥∥∇θ(R ◦ L ◦h)(θ(T ))∥∥∥ 2,2 ≤ c1 √ Lmn d ∥∥∥∇h(R ◦ L)(h(θ(T )))∥∥∥ 2,2
(79) where c1 > 0 is the constant hidden in O (√ Lmn d ) .
So with ′ = 1c1
√ d Lmn , we finally obtain∥∥∥∇θ(R ◦ L ◦h)(θ(T ))∥∥∥ 2,2 ≤ c1 √ Lmn d ∥∥∥∇h(R ◦ L)(h(θ(T )))∥∥∥ 2,2
≤ c1
√ Lmn
d ′
≤
(80)
If
T = O
( dn2
ηδm ′2
) = O ( dn2
ηδm
Lmn
d 2
) = O ( Ln3
ηδ 2
) (81)
which concludes the proof.

heading: D.7.4 PROOF OF TECHNICAL LEMMA 1
text: For any integer k ≥ 1, and all (ai)ni=1 ∈ ( Rk )n , we have
ES
[ 1
b ∑ i∈S ai
] =
∑ 1≤i1,...,ib≤n ( n∏ k=1 p̄ik(L̂) ) 1 b b∑ j=1 aij  = 1
b ∑ 1≤i1,...,ib≤n  b∑ j=1 p̄ij (L̂) aij  n∏ k=1 k 6=j p̄ik(L̂)  
= 1
b b∑ j=1  ∑ 1≤i1,...,ib≤n p̄ij (L̂) aij  n∏ k=1 k 6=j p̄ik(L̂)  
= 1
b b∑ j=1   n∑ ij=1 p̄ij (L̂) aij  n∏ k=1 k 6=j ( n∑ ik=1 p̄ik(L̂) ) = 1
b b∑ j=1 ( n∑ i=1 p̄i(L̂) ai )
= n∑ i=1 p̄i(L̂) ai = Ep̄(L̂) [ai]
(82)

heading: D.8 CONVERGENCE OF ROBUST SGD WITH STALE LOSS HISTORY
text: The proof of the convergence of Algorithm 4.1 under the conditions of Theorem B.2 follows the same structure as the proof of the convergence of Robust SGD with exact loss history D.7.3. We will reuse the intermediate results of D.7.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) ∈ (Rm×m)L the parameters of the deep neural network at iteration t. We define the stale loss history at iteration t as
L̂ = ( L i (hi(θ (ti(t)))) )n i=1
(83)
where for all i, ti(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 defined. We also define the exact loss history that is unknown in Algorithm 4.1, as
L̆ = ( L i (hi(θ (t))) )n i=1
(84)
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 B.2 remains satisfied during the warm-up if it was satisfied by the initial parameters.
Similarly to (62) we define
θ(t+1) = θ(t) − η 1 b ∑ i∈S ∇θ(L i ◦hi)(θ(t)) (85)
and using Theorem D.1, similarly to (63), we obtain
R(L(h(θ(t+1))) ≤R(L(h(θ(t)))− η〈∇θ(R ◦ L ◦h)(θ(t)), 1
b ∑ i∈S ∇θ(L i ◦hi)(θ(t))〉
+ η ∥∥∥∇h(R ◦ L)(h(θ(t)))∥∥∥
1,2 O
( L2ω1/3 √ m log(m)√ d )∥∥∥∥∥1b∑ i∈S ∇θ(L i ◦hi)(θ(t)) ∥∥∥∥∥ 2,∞
+ η2O
(( βC(L)2
nρ + 2C(∇L)
) nL2m
d )∥∥∥∥∥1b∑ i∈S ∇θ(L i ◦hi)(θ(t)) ∥∥∥∥∥ 2
2,∞ (86)
We can still define α as in (64)
α = min θ min i p̄i(L(θ)) > 0 (87)
where we are guaranteed that α > 0 under assumptions B.1.
Since Theorem D.2 is independent to the choice of L̂, taking the expectation with respect to S, similarly to (72), we obtain
ES [ R(L(h(θ(t+1))) ] ≤ R(L(h(θ(t)))− η〈∇θ(R ◦ L ◦h)(θ(t)), n∑ i=1 p̄i(L̂)∇θ(L i ◦hi)(θ(t)))〉
+ η ∥∥∥∇h(R ◦ L)(h(θ(t)))∥∥∥
1,2 O
( L2ω1/3 √ nm log(m)√ d )√√√√ n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ(t))) ∥∥∥2 2
+ η2O
(( βC(L)2
nρ + 2C(∇L)
) nL2m
d ) O (m dα ) n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ (t))) ∥∥∥2 2
(88)
where the differences with respect to (72) comes from the fact that L̂ is not the exact loss history here, i.e. L̂ 6= L̆, which leads to
∇θ(R ◦ L ◦h)(θ(t)) = n∑ i=1 p̂i(L̆)∇θ(L i ◦hi)(θ(t)))
6= n∑ i=1 p̄i(L̂)∇θ(L i ◦hi)(θ(t)))
(89)
And ∥∥∥∇h(R ◦ L)(h(θ(t)))∥∥∥ 1,2 = n∑ i=1 ∥∥∥p̂i(L̆)∇hi L i (hi(θ (t)))) ∥∥∥ 2
6= n∑ i=1 ∥∥∥p̂i(L̂)∇hi L i (hi(θ (t)))) ∥∥∥ 2
(90)
Let
K ′ = C(L)A(∇L)O  βLm3/2 log2(m) αn1/2ρd3/2b log ( 1
1−α
)  (91)
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 B.1. We can prove that, with probability at least 1− exp ( −Ω ( log2(m) )) ,
• according to lemma D.8.1
∥∥∥p̂(L̂)− p̂(L̆)∥∥∥ 2 = √√√√ n∑ i=1 ( p̂i(L̂)− p̂i(L̆) )2 ≤ ηαK ′ (92)
• according to lemma D.8.2∣∣∣∣∣〈∇θ(R ◦ L ◦h)(θ(t))− n∑ i=1 p̄i(L̂)∇θ(L i ◦hi)(θ(t))), n∑ i=1 p̄i(L̂)∇θ(L i ◦hi)(θ(t)))〉 ∣∣∣∣∣ ≤ ηm
d K ′ n∑ i=1 ∥∥∥p̄i(L̂)∇θ(L i ◦hi)(θ(t))) ∥∥∥2 2 (93)
• according to lemma D.8.3
∥∥∥∇h(R ◦ L)(h(θ(t)))∥∥∥ 1,2 ≤ (√ n+ ηK ′ )√√√√ n∑ i=1 ∥∥∥p̄i(L̂)∇θ(L i ◦hi)(θ(t))) ∥∥∥2 2
(94)
Combining those three inequalities with (88) we obtain ES [ R(L(h(θ(t+1))) ] −R(L(h(θ(t))) ≤
η [ −Ω ( mδ
dn2
) +O ( nL2mω1/3 √ log(m)
d
)] n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ (t))) ∥∥∥2 2
η2O ( K (n/α)L2m2 d2 + ( 1 + m d ) K ′ ) n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ (t))) ∥∥∥2 2
(95) One can see that compared to (72), there is only the additional term ( 1 + md ) K ′.
Using our choice of η,
η = ηstale ≤ O ( δ
n2K ′ ηexact
) (96)
where ηexact is the learning rate of Theorem B.1, we have
Ω
( ηmδ
dn2
) ≥ O ( η2 ( 1 + m
d
) K ′ )
(97)
As a result, η2 ( 1 + md ) K ′ is dominated by the term Ω ( ηmδ dn2 ) In addition, since ηstale ≤ ηexact, Ω ( ηmδ dn2 ) still dominates also the ther terms as in the proof of Theorem B.1.
As a consequence, we obtain as in (76) that for any iteration t > 0 (after the end of the warm-up)
ESt [ R(L(h(θ(t+1))) ] ≤ R(L(h(θ(t)))− Ω ( ηmδ
dn2 ) n∑ i=1 ∥∥∥p̄i(L̂)∇hi L i (hi(θ (t))) ∥∥∥2 2 (98)
This concludes the proof using the same arguments as in the end of the proof of Theorem B.1 starting from (76).

heading: D.8.1 PROOF OF TECHNICAL LEMMA 2
text: Using Lemma 4.2 and Lemma 4.1 we obtain∥∥∥p̂(L̂)− p̂(L̆)∥∥∥ 2 = ∥∥∥∇vR(L̂)−∇vR(L̆)∥∥∥ 2
≤ β nρ ∥∥∥L̂ − L̆∥∥∥ 2
(99)
Using assumptions B.2 and (Allen-Zhu et al., 2019, Claim 11.2)∥∥∥p̂(L̂)− p̂(L̆)∥∥∥ 2 ≤ β nρ √√√√ n∑ i=1 ( L i ◦ hi(θ(t))− L i ◦ hi(θ(ti(t))) )2
≤ β nρ C(L)C(h) √√√√ n∑ i=1 ∥∥θ(t) − θ(ti(t))∥∥2 2,2
≤ C(L)O ( βLm1/2
nρd1/2 )√√√√ n∑ i=1 ∥∥θ(t) − θ(ti(t))∥∥2 2,2
(100)
Where C(L) is the constant of Lipschitz continuity of the per-example loss L (see assumptions B.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) (85), we obtain ∥∥∥p̂(L̂)− p̂(L̆)∥∥∥ 2 ≤ C(L)O ( βLm1/2 nρd1/2 )√√√√√ n∑ i=1 ∥∥∥∥∥∥θ(ti(t)) −  t−1∑ τ=ti(t) η b ∑ j∈Sτ ∇θ(L j ◦ hj)(θ(τ)) − θ(ti(t)) ∥∥∥∥∥∥ 2
2,2
≤ ηC(L)O ( βLm1/2
nρd1/2 )√√√√√ n∑ i=1 ∥∥∥∥∥∥ t−1∑
τ=ti(t)
1
b ∑ j∈Sτ ∇θ(L j ◦ hj)(θ(τ)) ∥∥∥∥∥∥ 2
2,2
Let A(∇L) a bound on the gradient of the per-example loss function. Using Theorem D.2 and the chain rule ∀j, ∀τ ∥∥∥∥∇θ(Lj ◦ hj)(θ(τ)) ∥∥∥∥ 2,2 ≤ A(∇L)O (m d ) (101)
And using the triangular inequality∥∥∥∥∥∥ t−1∑
τ=ti(t)
1
b ∑ j∈Sτ ∇θ(L j ◦ hj)(θ(τ)) ∥∥∥∥∥∥ 2,2 ≤ t−1∑ τ=ti(t) 1 b ∑ j∈Sτ ∥∥∥∥∇θ(Lj ◦ hj)(θ(τ)) ∥∥∥∥ 2,2
≤ t−1∑
τ=ti(t)
A(∇L)O (m d ) ≤ A(∇L)O (m d ) (t− ti(t))
(102)
As a result, we obtain∥∥∥p̂(L̂)− p̂(L̆)∥∥∥ 2 ≤ ηC(L)A(∇L)O ( βLm3/2 nρd3/2 )√√√√ n∑ i=1 (t− ti(t))2 (103)
For all i and for any τ the probability that the sample i is not in batch Sτ is lesser than (1− α)b.
Therefore, for any k ≥ 1 and for any t,
P (t− ti(t) ≥ k) ≤ (1− α)kb (104)
For k ≥ 1bΩ ( log2(m)
log( 11−α )
) , we have (1− α)kb ≤ exp ( −Ω ( log2(m) )) , and thus with probability at
least 1− exp ( −Ω ( log2(m) )) ,
∀t, t− ti(t) ≤ O  log2(m) b log ( 1
1−α
)  (105)
As a result, we finally obtain that with probability at least 1− exp ( −Ω ( log2(m) )) ,
∥∥∥p̂(L̂)− p̂(L̆)∥∥∥ 2 ≤ ηC(L)A(∇L)O ( βLm3/2 nρd3/2 )√ nO  log2(m) b log ( 1
1−α
) 
≤ ηαO  βLm3/2 log2(m) αn1/2ρd3/2b log ( 1
1−α
) 
≤ ηαK ′
(106)

heading: D.8.2 PROOF OF TECHNICAL LEMMA 3
text: Let us first denote
A = ∣∣∣∣∣〈∇θ(R ◦ L ◦h)(θ(t))− n∑ i=1 p̄i(L̂)∇θ(L i ◦hi)(θ(t))), n∑ i=1 p̄i(L̂)∇θ(L i ◦hi)(θ(t)))〉 ∣∣∣∣∣ =
∣∣∣∣∣〈 n∑ i=1 ( p̄i(L̆)− p̄i(L̂) ) ∇θ(L i ◦hi)(θ(t))), n∑ i=1 p̄i(L̂)∇θ(L i ◦hi)(θ(t)))〉 ∣∣∣∣∣ (107)
Using Cauchy-Schwarz inequality
A = ∣∣∣∣∣∣ n∑ i=1 ( p̄i(L̆)− p̄i(L̂) ) 〈∇θ(L i ◦hi)(θ(t))), n∑ j=1 p̄j(L̂)∇θ(L j ◦hj)(θ(t)))〉 ∣∣∣∣∣∣ ≤ ∥∥∥p̂(L̂)− p̂(L̆)∥∥∥
2 √√√√√ n∑ i=1 〈∇θ(L i ◦hi)(θ(t))), n∑ j=1 p̄j(L̂)∇θ(L j ◦hj)(θ(t)))〉 2 (108)
Let
B = 〈∇θ(L i ◦hi)(θ(t))), n∑ j=1 p̄j(L̂)∇θ(L j ◦hj)(θ(t)))〉 (109)
Using again Cauchy-Schwarz inequality
B ≤ ∥∥∥∇θ(L
i ◦hi)(θ(t))) ∥∥∥ 2,2 ∥∥∥∥∥∥ n∑ j=1 p̄j(L̂)∇θ(L j ◦hj)(θ(t))) ∥∥∥∥∥∥ 2,2
(110)
As a result, A becomes
A ≤ ∥∥∥p̂(L̂)− p̂(L̆)∥∥∥
2 ∥∥∥∥∥∥ n∑ j=1 p̄j(L̂)∇θ(L j ◦hj)(θ(t))) ∥∥∥∥∥∥ 2,2 √√√√ n∑ i=1 ∥∥∥∇θ(L i ◦hi)(θ(t))) ∥∥∥2 2,2
≤ ∥∥∥p̂(L̂)− p̂(L̆)∥∥∥
2 ∥∥∥∥∥∥ n∑ j=1 p̄j(L̂)∇θ(L j ◦hj)(θ(t))) ∥∥∥∥∥∥ 2,2 √√√√ n∑ i=1 1 α2 ∥∥∥p̄j(L̂)∇θ(L i ◦hi)(θ(t))) ∥∥∥2 2,2
≤ 1 α ∥∥∥p̂(L̂)− p̂(L̆)∥∥∥ 2 ∥∥∥∥∥∥ n∑ j=1 p̄j(L̂)∇θ(L j ◦hj)(θ(t))) ∥∥∥∥∥∥ 2
2,2
(111)
Using the triangular inequality, Theorem D.2, and Lemma D.8.1, we finally obtain
A ≤ m αd ∥∥∥p̂(L̂)− p̂(L̆)∥∥∥ 2 n∑ j=1 ∥∥∥∥p̄j(L̂)∇hj Lj (hj(θ(t))) ∥∥∥∥2 2,2
≤ ηm d K ′ n∑ j=1 ∥∥∥∥p̄j(L̂)∇hj Lj (hj(θ(t))) ∥∥∥∥2 2,2
(112)

heading: D.8.3 PROOF OF TECHNICAL LEMMA 4
text: We have∥∥∥∇h(R ◦ L)(h(θ(t)))∥∥∥ 1,2 = n∑ j=1 p̄j(L̆) ∥∥∥∥∇hj Lj (hj(θ(t))) ∥∥∥∥ 2,2
= n∑ j=1 p̄j(L̂) ∥∥∥∥∇hj Lj (hj(θ(t))) ∥∥∥∥ 2,2
+ n∑ j=1
( p̄j(L̆)− p̄j(L̂)
p̄j(L̂)
) p̄j(L̂) ∥∥∥∥∇hj Lj (hj(θ(t))) ∥∥∥∥
2,2
(113)
Using Cauchy-Schwarz inequality ∥∥∥∇h(R ◦ L)(h(θ(t)))∥∥∥ 1,2 ≤ √n+ √√√√ n∑
j=1
( p̄j(L̆)− p̄j(L̂)
p̄j(L̂)
)2√√√√ n∑ j=1 ∥∥∥∥p̄j(L̂)∇hj Lj (hj(θ(t))) ∥∥∥∥2 2,2
(114)
Using Lemma D.8.1 n∑ j=1
( p̄j(L̆)− p̄j(L̂)
p̄j(L̂)
)2 ≤ 1 α ∥∥∥p̂(L̂)− p̂(L̆)∥∥∥ 2
≤ ηK ′ (115)
Therefore, we finally obtain
∥∥∥∇h(R ◦ L)(h(θ(t)))∥∥∥ 1,2 ≤ (√ n+ ηK ′ )√√√√ n∑ j=1 ∥∥∥∥p̄j(L̂)∇hj Lj (hj(θ(t))) ∥∥∥∥2 2,2
(116)

Review summaries:
This paper proposes a method for Distributionally Robust Optimization (DRO). DRO has recently been proposed (Namkoong and Duchi 2016 and others) as a robust learning framework compared to Empirical Risk Minimization (ERM).  My analysis of this work in short is that the problem that this paper addresses is interesting and not yet solved. This paper proposes a simple and efficient method for DRO. However, convergence guarantees are weak which is reflected in their weak experimental results. And for a 10-page paper, the writing has to improve quite a lot.

This paper studies the Distributionally Robust Optimization (DRO), in the sense that the weights assigned to the training data can change, but the training data itself remains unchanged. They demonstrate that SGD with hardness weighted sampling is a principled and efficient optimization method for DRO in machine learning and is particularly suited in the context of deep learning. On the theoretical side, they prove the convergence of our DRO algorithm for over-parameterized