Title: SGD WITH HARDNESS WEIGHTED SAMPLING FOR DISTRIBUTIONALLY ROBUST DEEP LEARNING

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

1. INTRODUCTION:

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 )} n i=1 , where x i are the inputs and y i 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(x i ; θ), y i ) (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 q i L (h(x i ; θ), y i ) -

1 β n i=1 1 n φ (nq i ) (2)

where φ is a convex function that defines a φ-divergence (Csiszár et al., 2004 ), q = (q i ) 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.

2. RELATED WORK IN DRO WITH A WASSERSTEIN DISTANCE:

In this work, we focus on DRO with a φ-divergence (Csiszá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 {(x i , y i )} n i=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 → y, and omit explicitly mentioning any parameter regularisation term.

Let ∆ n ⊂ R n be the set of empirical weighted training data distributions defined according to a given training dataset

∆ n = p = (p i ) n i=1 ∈ [0, 1] n | n i=1 p i = 1

(3) and let pdata be the uniform empirical training data distribution, i.e. for all i, pdata,i = 1 n . Let θ be the set of parameters of the predictor h( . ; θ) : x → y we want to train, and h : θ → (h(x i ; θ)) 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

EQUATION

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

EQUATION

However, pdata is typically biased compared to the true data distribution. 

EQUATION

Example 3.1. For φ : z → z log(z), D φ is the Kullback-Leibler (KL) divergence:

D φ (q p) = D KL (q p) = n i=1 q i log q i p i (7)

And, for φ : z → (z -1) 2 , D φ is the Pearson χ 2 divergence:

D φ (q p) = χ 2 (q p) = n i=1 (q i -p i ) 2 p i (8) Definition 3.4 (Distributionally Robust Loss). R(L(h(θ))) = max q∈∆n E q [L (h(x; θ), y)] - 1 β D φ (q pdata ) = max q∈∆n n i=1 q i L (h(x i ; θ), y i ) - 1 nβ n i=1 φ (nq i ) (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).

EQUATION

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 E q [L(h (x; θ) , y)] - 1 β D φ (q pdata ) = E pdata [L(h (x; θ) , y)] + β 2φ (1) V pdata [L(h (x; θ) , y)] + o(β) max q∈∆n E q [L(h (x; θ) , y)] - 1 β D φ (q pdata ) -----→ β-→+∞ max i L(h (x i ; θ) , y i ) (11)

where V pdata 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

EQUATION

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

4. SGD WITH HARDNESS WEIGHTED SAMPLING:



4.1. DISTRIBUTIONALLY ROBUST OPTIMIZATION WITH SGD AND ADAPTIVE SAMPLING:

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 : R m → R ∪ {+∞} be a proper function. The Fenchel conjugate of f is defined as

EQUATION

where δ ∆n is the characteristic function of the closed convex set ∆ n , i.e.

∀p ∈ R n , δ ∆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:

EQUATION

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

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.1 for efficient 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 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 → z log(z) -z + 1), we have (see D.2 for a proof)

EQUATION

And for the Pearson χ 2 divergence (i.e. φ : z → (z -1)

2 ), we have:

EQUATION

In both cases, we can verify consistency with (11) Li ← L(h(x i ; θ), y i )

as ∀i ∈ {1, . . . , n}, pi (L(h(θ))) ---→ β-→0 1 n (22)

23:

// Run backward pass and update the parameters of the model 24:

θ ← θ -η 1 b i∈I ∇ θ L(h(x i ; θ), y i ) 25: Output: θ 4.2 EFFICIENT ALGORITHM FOR DISTRIBUTIONALLY ROBUST DEEP LEARNING

The second equality in (19) implies that ∇ θ L(h i (θ), y i ) 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(x i ; θ (ti) ), y i ) . Where for all i, t i 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.

4.3. CONVERGENCE OF SGD WITH HARDNESS WEIGHTED SAMPLING FOR OVER-PARAMETERIZED DEEP NEURAL NETWORKS WITH ReLU:

Convergence results for over-parameterized deep learning has recently been proposed in (Allen-Zhu 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.

4.4. DRO AS PRINCIPLED ONLINE HARD EXAMPLE MINING:

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 p i of sampling example x i is an non-decreasing function of the (potentially stale) loss value associated with x i . 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 p i of sampling example x i is an non-increasing function of the loss value associated with x j for all j = i. 

See:



5. EXPERIMENTS:

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

5.1. ROBUSTNESS TO DOMAIN GAP:

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. This suggests that our SGD with hardness weighted sampling is more robust than SGD to a domain gap between the training and the testing dataset.

5.2. STABILITY OF DRO:

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.

6. CONCLUSION AND DISCUSSION:

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. 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.2 COMMENT ON EARLY-STOPPING 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 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.

B CONVERGENCE OF SGD WITH HARDNESS WEIGHTED SAMPLING FOR 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 h i : θ → 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 = j, x i -x j ≥ δ.

• 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+1 l=0 are initialized at random such that:

-[θ 0 ] i,j ∼ N 0, 2 m for every (i, j) ∈ {1, . . . , m} × {1, . . . , d} -[θ l ] i,j ∼ N 0, 2

m for every (i, j) ∈ {1, . . . , m} 2 and l ∈ {1, . . . , L} -[θ L+1 ] i,j ∼ N 0, 1 d for every (i, j) ∈ {1, . . . , d} × {1, . . . , m} Assumption B.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 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 C 1 , C 2 , C 3 > 0 such that the number of hidden units satisfies m ≥ C 1 (d -1 × poly(n, L, δ -1 )), δ ≥ C2 L , and the learning rate be η exact = C 3 min 1,

αn 2 ρ βC(L) 2 +2nρC(∇ L) × bδd poly(n,L)m log 2 (m)

. There exists constants C 4 , C 5 > 0 such that with probability at least 1 -exp -C 4 (log 2 (m)) over the randomness of the initialization and the mini-batches, Robust SGD with exact loss vector finds

∇ θ (R • L •h)(θ) ≤ after T = C 5 Ln 3 ηexactδ 2 iterations.

where α = min θ min i pi (L(θ)) is lower bound on the sampling probabilities. For the Kullback-Leibler φ-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 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, ρ = 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 (Allen-Zhu 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. 

αρd 3/2 δb log( 1 1-α ) βC(L)A(∇ L)Lm 3/2 n 3/2 log 2 (m) × η exact for a constant C 6 > 0.

With probability at least 1 -exp -C 4 (log 2 (m)) over the randomness of the initialization and the mini-batches, Robust SGD with exact loss vector finds

∇ θ (R • L •h)(θ) ≤ after T = C 5 Ln 3 η 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,

αρd 3/2 δb log( 1 1-α ) βC(L)A(∇ L)Lm 3/2 n 3/2 log 2 (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.

C SUMMARY OF THE NOTATIONS USED IN THE PROOFS:

For the ease of following the proofs we first summarize our notations.

C.1 PROBABILITY THEORY NOTATIONS:

• ∆ n = {p = (p i ) n i=1 ∈ [0, 1] n , i p i = 1} • Let q = (q i ) ∈ ∆ n , and f a function, we denote E q [f (x)] := n i=1 q i f (x i ). • Let q ∈ ∆ n , and f a function, we denote V q [f (x)] := n i=1 q i f (x i ) -E q [f (x)]

2 .

• pdata is the uniform training data distribution, i.e. pdata =

1 n n i=1 ∈ ∆ n C.2 MACHINE LEARNING 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 )} n i=1

, where for all i ∈ {1, . . . , n}, x i ∈ R d and y i ∈ R d • h : x → y is the predictor • θ is the set of parameters of the predictor • For all i, h i : θ → 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 → 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 ) n i=1 → (L i (h i )) n i=1 • b ∈ {1, . . . , n} is the batch size • η > 0 is the learning rate • ERM is short for Empirical Risk Minimization C.3 DISTRIBUTIONALLY ROBUST OPTIMISATION NOTATIONS • Forall θ, R(L(h(θ))) = max q∈∆n E q [L (h(x; θ), y)] -1 β D φ (q pdata

) 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 C.4 MISCELLANEOUS 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:

EQUATION

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 (Csiszár et al., 2004) .

D.2 PROOF OF EXAMPLE 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 → 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 q i L i • h i (θ) - 1 β n i=1 q i log (nq i ) = max q∈∆n n i=1 q i L i • h i (θ) - 1 β q i log (nq i )

To simplify the notations, let us denote 

v = (v i ) n i=1 = L • h(θ) = (L i • h i (θ))

EQUATION

First, let us remark that the function q → n i=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 (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∈R n + n i=1 q i v i - 1 β q i log (nq i ) s.t. n i=1 q i = 1 (25)

There exists a Lagrange multiplier λ ∈ R, such that the solution p of ( 25) is characterized by

EQUATION

Which we can rewrite as

EQUATION

The last equality gives exp (βλ -1) = n For the ease of reading, let us first recall that given a φ-Divergence that satisfies assumptions 3.1, we have defined in ( 9)

EQUATION

And in ( 13)

EQUATION

where δ ∆n is the characteristic function of the closed convex set ∆ n , i.e.

EQUATION

)

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

EQUATION

φ 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 = (p i ) n i=1 , q = (q i ) 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λp i + n(1 -λ)q i ) + δ ∆n (λp + (1 -λ)q) ≤ λG(p) + (1 -λ)G(q) - 1 βn n i=1 ρλ(1 -λ) 2 |nq i -np i | 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 & Lemaréchal, 2013).

We obtain (32) using Definition 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).

arg max

EQUATION

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 & Lemaré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

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échal, 2013, Theorem 6.1.2 p.280). Lemma D.2. A necessary and sufficient 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 s 2 -s 1 , x 2 -x 1 ≥ c x 2 -x 1 2 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(p 2 ), v 2 -v 1 , p 2 -p 1 ≥ nρ β p 2 -p 1 2 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

v 2 -v 1 p 2 -p 1 ≥ v 2 -v 1 , p 2 -p 1 We conclude that nρ β ∇ v R(v 2 ) -∇ v R(v 1 ) ≤ v 2 -v 1

Which implies that R is β nρ -gradient Lipchitz continuous.

D.4 PROOF OF LEMMA 4.2: FORMULA OF THE DISTRIBUTIONALLY ROBUST LOSS GRADIENT:

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

EQUATION

where ∇ v R 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), φ is strongly convex, and ∆ n is closed and convex, so the following optimization problem has one and only one solution:

∇ θ (R • L • h)(θ) = n i=1 ∂R ∂v i (L • h(θ)))∇ θ (L i • h i )(θ) = n i=1 pi (L(h(θ)))∇ θ (L i • h i )(θ) = E p(L(h(θ))) [∇ θ L (h(x; θ),

max p=(pi) n i=1 ∈∆ n v, p - 1 βn n i=1 φ(np i ) (40)

Making the constraints associated with p ∈ ∆ n explicit, this can be rewritten as

max p=(pi) n i=1 ∈R n v, p - 1 βn n i=1 φ(np i ) s.t. ∀i ∈ {1, . . . , n}, p i ≥ 0 s.t. n i=1 p i = 1 (41)

There exists KKT multipliers λ ∈ R and ∀i, µ i ≥ 0 such that the solution p = (p i )

n i=1 satisfies:                    ∀i ∈ {1, . . . , n}, v i - 1 β φ (np i ) + λ -µ i = 0 ∀i ∈ {1, . . . , n}, µ i p i = 0 ∀i ∈ {1, . . . , n}, p i ≥ 0 n i=1 pi = 1 (42)

Since φ is continuously differentiable and strongly convex, we have (φ ) -1 = (φ * ) , where φ * is the Fenchel conjugate of φ (see Hiriart-Urruty & Lemaréchal, 2013, Proposition 6.1.2). As a result, (42) can be rewritten has:

                 ∀i ∈ {1, . . . , n}, pi = 1 n (φ * ) (β(v i + λ -µ i ))

∀i ∈ {1, . . . , n}, µ i p i = 0 ∀i ∈ {1, . . . , n}, p i ≥ 0

1 n n i=1 (φ * ) (β(v i + λ -µ i )) = 1 (43)

We now show that the KKT multipliers are uniquely defined.

The µ i 's are uniquely defined 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 (43

) it comes (φ * ) (β(v i + λ -µ 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:

(φ * ) (β(v i + λ -µ i )) = 0

And (43) can be rewritten as:

         ∀i ∈ {1, . . . , n}, pi = ReLU 1 n (φ * ) (β(v i + λ)) = 1 n ReLU (φ * ) (β(v i + λ)) 1 n n i=1 ReLU (φ * ) (β(v i + λ)) = 1 (44)

λ is uniquely defined by v and a continuous function of v: Let λ ∈ R that satisfies (44).

We have

1 n n i=1 ReLU (φ * ) (β(v i + λ)) = 1.

So there exists at least one index i 0 such that

ReLU (φ * ) (β(v i0 + λ)) = (φ * ) (β(v i0 + λ)) ≥ 1

Since (φ * ) -1 is continuous and striclty increasing, λ → ReLU (φ * ) (β(v i0 + λ )) is continuous and strictly increasing on a neighborhood of λ.

In addition ReLU is continuous and increasing, so for all i ∈ {1, . . . , n}, λ → ReLU (φ * ) (β(v i + λ )) is a continuous and increasing function.

As a result, λ → 1

n n i=1 ReLU (φ * ) (β(v i + λ )

) 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 → λ(v) is continuous.

Hard Example Mining Sampling::

For any pseudo loss vector v = (v i ) n i=1 ∈ R n , there exists a unique λ and a unique p that satisfies (44), so we can define the mapping:

EQUATION

where for all v, λ(v) is the unique λ ∈ R satisfying (44).

We will now demonstrate that each pi0 (v) for i 0 ∈ {1, . . . , n} is an increasing function of v i and a decreasing function of the v i for i = i 0 . Without loss of generality we assume i 0 = 1.

Let v = (v i ) n i=1 ∈ R n , and > 0. Let us define v = (v i ) n i=1 ∈ R n , such that v 1 = v 1 + and ∀i ∈ {2, . . . , n}, v i = v i .

Similarly as in the proof of the uniqueness of λ above, we can show that there exists η > 0 such that the function

F : λ → 1 n n i=1 ReLU (φ * ) (β(v i + λ ))

is continuous and strictly increasing on [λ(v) -η, λ(v) + η], and

F (λ(v)) = 1. v → λ(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 (φ * ) (β(v i + λ(v ))) ≥ F (λ(v )) > F (λ(v)) > 1

which is a contradiction. As a result

EQUATION

Using ( 46), ( 44), and the fact that ReLU • (φ * ) is an increasing function, we obtain for all i ∈ {2, . . . , n}

EQUATION

In addition

EQUATION

This holds for any i 0 and any v, which concludes the proof.

D.6 PROOF THAT R • L IS ONE-SIDED GRADIENT LIPCHITZ

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, L i is C(L)-Lipschitz continuous and C(∇ L)-gradient Lipschitz continuous.

Let z = (z i ) n i=1 , z = (z i ) n i=1 ∈ R dn .

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 ), z i -z i = n i=1 pi (L(z))∇ zi L i (z i ) -pi (L(z ))∇ zi L i (z i ), z i -z i = n i=1 pi (L(z)) ∇ zi L i (z i ) -∇ zi L i (z i ), z i -z i + n i=1 (p i (L(z)) -pi (L(z ))) ∇ zi L i (z i ), z i -z i (49)

Where for all i ∈ {1, . . . , n} we have used the chain rule

∇ zi (R • L)(z) = n j=1 ∂R * ∂v j (L(z))∇ zi L j (z j ) = pi (L(z))∇ zi L i (z i ) Let A = n i=1 pi (L(z)) ∇ zi L i (z i ) -∇ zi L i (z i ), z i -z i

For all i, L i is C(∇ L)-gradient Lipchitz continuous, so using Cauchy-Schwarz inequality

A ≤ n i=1 C(∇ L) z i -z i 2 = C(∇ L) z -z 2 (50) Let B = n i=1 (p i (L(z)) -pi (L(z ))) ∇ zi L i (z i ), z i -z i

Using the triangular inequality:

EQUATION

Combining ( 49), ( 50) and ( 51) we finally obtain:

EQUATION

From there, we can obtain the following inequality that will be used for the proof of the semismoothness property in Theorem D.1:

EQUATION

D.7 PROOF OF THE CONVERGENCE OF ROBUST SGD:

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:

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

D.7.1 SEMI-SMOOTHNESS PROPERTY FOR THE DISTRIBUTIONALLY ROBUST LOSS:

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 ω ∈ Ω

d 3/2 m 3/2 L 3/2 log 3/2 (m) , O 1 L 4.5 log 3 (m)

, and the θ (0) being initialized randomly as described in assumption B.1. With probability as least 1 -exp (-Ω(mω 3/2 L)) over the initialization, we have for all θ, θ ∈ (R m×m )

L with θθ (0) 2 ≤ ω, and θ -

EQUATION

where for all layer l ∈ {1, . . . , L}, θ l is the vector of parameters for layer l, and

θ -θ 2,∞ = max l θ l -θ l 2 θ -θ 2 2,∞ = 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 pi (L(h(θ)))∇ hi L i (h i (θ)) 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) Let θ, θ ∈ (R m×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

EQUATION

Then using the chain rule (55)

EQUATION

For all i ∈ {1, . . . , n}, let us denote

loss i := ∇ 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 loss i , we can conclude the proof of our Theorem D.1 (as in Allen-Zhu et al., 2019, Appendix A).

D.7.2 GRADIENT BOUNDS FOR THE DISTRIBUTIONALLY ROBUST LOSS:

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 n 9/2 L 6 log 3 (m) , and θ (0) being initialized randomly as described in assumption B.1. With probability as least 1 -exp (-Ω(mω 3/2 L)) over the initialization, we have for all θ ∈ (R m×m )

L with θ -

EQUATION

It is worth noting that the loss vector L used for computing the robust probabilities p( L) = pi ( 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 θ ∈ (R m×m ) L satisfying the conditions of Theorem D.2, and L ∈ R n .

Let us denote:

v := pi ( L)∇ hi L i (h i (θ)) 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},

pi ( L)∇ θ l (L i •h i )(θ) 2 2 ≤ O m d pi ( L)∇ hi L i (h i (θ)) 2 2 ∀l ∈ {1, . . . , L}, ∀ L ∈ R n n i=1 pi ( L)∇ θ l (L i •h i )(θ) 2 2 ≤ O mn d n i=1 pi ( L)∇ hi L i (h i (θ)) 2 2 n i=1 pi ( L)∇ θ L (L i •h i )(θ) 2 2 ≥ Ω mδ dn max i pi ( L)∇ hi L i (h i (θ)) 2 2

In addition 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.

max i pi ( L)∇ hi L i (h i (θ)) 2 2 ≥ 1 n n i=1 pi ( L)∇ hi L i (h i (θ))

Let θ ∈ (R m×m )

L satisfying the conditions of Theorem B.1, and L be the exact loss history at θ, i.e. 

L = L:



∼ p( L):

(61)

Let θ ∈ (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 θ satisfies the conditions of Theorem D.1, we obtain

EQUATION

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

EQUATION

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 E S be the expectation with respect to S. Applying E S to (63), we obtain

E S [R(L(h(θ ))] ≤R(L(h(θ)) -η ∇ θ (R • L •h)(θ) 2 2,2 + η ∇ h (R • L)(h(θ)) 2,2 O nL 2 ω 1/3 m log(m) √ d n i=1 max l pi ( L)∇ θ l (L i •h i )(θ) 2 + η 2 O βC(L) 2 nρ + 2C(∇ L) nL 2 m d 1 α n i=1 max l pi ( L)∇ θ l (L i •h i )(θ)

2:

(65) where we have used the following results:

• For any integer k ≥ 1, and all (a i ) n i=1 ∈ R k n , we have (see the proof in D.7.4)

EQUATION

• Using (66) for (a i )

n i=1 = (∇ θ (L i •h i )(θ)) n i=1

, and the chain rule ( 55)

EQUATION

• Using the triangular inequality

EQUATION

And using (66) for (a i )

n i=1 = ∇ θ (L i •h i )(θ) 2,∞ n i=1 , E S   1 b i∈S ∇ θ (L i •h i )(θ) 2,∞   ≤ n i=1 pi ( L) ∇ θ (L i •h i )(θ) 2,∞ ≤ n i=1 max l ∇ θ l (p i ( L) L i •h i )(θ) 2 ≤ √ n n i=1 max l ∇ θ l (p i ( L) L i •h i )(θ) 2 2 (69)

where we have used Cauchy-Schwarz inequality for the last inequality. • Using (68) and the convexity of the function

EQUATION

And using ( 66) for (a i )

n i=1 = ∇ θ (L i •h i )(θ) 2 2,∞ n i=1 , E S   1 b i∈S ∇ θ (L i •h i )(θ) 2 2,∞   ≤ n i=1 pi ( L) ∇ θ (L i •h i )(θ) 2 2,∞ ≤ n i=1 1 pi ( L) max l ∇ θ l (p i ( L) L i •h i )(θ) 2 2 ≤ 1 α n i=1 max l ∇ θ l (p i ( L) L i •h i )(θ) 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 α = 1 n . 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)(θ) 2 2,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)

E S [R(L(h(θ ))] ≤ R(L(h(θ)) -Ω ηmδ dn 2 n i=1 pi ( L)∇ hi L i (h i (θ)) 2 2 + ηO nL 2 ω 1/3 m log(m) √ d O m d n i=1 pi ( L)∇ hi L i (h i (θ)) 2 2 + η 2 O βC(L) 2 nρ + 2C(∇ L) nL 2 m d O m dα n i=1 pi ( L)∇ hi L i (h i (θ)) 2 2 ≤ R(L(h(θ)) -Ω ηmδ dn 2 n i=1 pi ( L)∇ hi L i (h i (θ)) 2 2 + O ηnL 2 mω 1/3 log(m) d + K η 2 (n/α)L 2 m 2 d 2 n i=1 pi ( L)∇ hi L i (h i (θ)) 2 2

(72) 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 D.7.3 for more details), and an additional multiplicative term K. So in total, this term differs by a multiplicative factor αn K from the analogous term in the proof of (Allen-Zhu et al., 2019 ).

• we have (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 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 L 2 loss.

n i=1 pi ( L)∇ hi L i (h i (θ)) 2 2 instead of F (W

Our choice of learning rate in Theorem B.2 can be rewritten as

EQUATION

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δ dn 2 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

E St R(L(h(θ (t+1) )) ≤ R(L(h(θ (t) )) -Ω ηmδ dn 2 n i=1 pi ( L)∇ hi L i (h i (θ (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

EQUATION

where according to (54)

EQUATION

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)

EQUATION

where c 1 > 0 is the constant hidden in O 

∇ θ (R • L •h)(θ (T ) ) 2,2 ≤ c 1 Lmn d ∇ h (R • L)(h(θ (T ) )) 2,2 ≤ c 1 Lmn d ≤ (80) If T = O dn 2 ηδm 2 = O dn 2 ηδm Lmn d 2 = O Ln 3 ηδ 2 (81)

which concludes the proof.

D.7.4 PROOF OF TECHNICAL LEMMA 1:

For any integer k ≥ 1, and all (a i )

EQUATION

D.8 CONVERGENCE OF ROBUST SGD WITH STALE LOSS HISTORY:

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) ∈ (R m×m ) L the parameters of the deep neural network at iteration t.

We define 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 defined.

We also define 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 B.2 remains satisfied during the warm-up if it was satisfied by the initial parameters.

Similarly to (62) we define

EQUATION

and using Theorem D.1, similarly to (63), we obtain

EQUATION

We can still define α as in ( 64)

EQUATION

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

EQUATION

where the differences with respect to (72) comes from the fact that L is not the exact loss history here, i.e. L = L, which leads to

EQUATION

And

EQUATION

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 -Ω log 2 (m) ,

• according to lemma D. • according to lemma D.8.2

∇ θ (R • L •h)(θ (t) ) - n i=1 pi ( L)∇ θ (L i •h i )(θ (t) )), n i=1 pi ( L)∇ θ (L i •h i )(θ (t) )) ≤ η m d K n i=1 pi ( L)∇ θ (L i •h i )(θ (t) )) 2 2

(93)

• according to lemma D.8.3 In addition, since η stale ≤ η exact , Ω ηmδ dn 2 still dominates also the ther terms as in the proof of Theorem B.1.

∇ h (R • L)(h(θ (t) )) 1,2 ≤ √ n + ηK n i=1 pi ( L)∇ θ (L i •h i )(θ (t) ))

As a consequence, we obtain as in (76) that for any iteration t > 0 (after the end of the warm-up) 

E St R(L(h(θ (t+1) )) ≤ R(L(h(θ (t) )) -Ω ηmδ dn 2 n i=1 pi ( L)∇ hi L i (h i (θ (t) ))

EQUATION

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 And using the triangular inequality

EQUATION

As a result, we obtain p( L) -p( L)

EQUATION

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 -t i (t) ≥ k) ≤ (1 -α) 

EQUATION

Using Cauchy-Schwarz inequality

EQUATION

Using again Cauchy-Schwarz inequality 

EQUATION

:



Figures and Tables:

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

Figure 1: Comparison of learning curves for ERM with SGD (blue) and DRO with our SGD with hardness weighted sampling (red). The models are trained on an imbalanced MNIST dataset (only 1% of the digits 3 kept in the training datatset) and evaluated on the original MNIST testing dataset.This suggests that our SGD with hardness weighted sampling is more robust than SGD to a domain gap between the training and the testing dataset.

.0, lr=0.01) DRO (beta=0.1, lr=0.01) DRO (beta=1.0, lr=0.01) DRO (beta=10, lr=0.01)DRO (beta=100, lr=0.01)

Figure 2: Comparison of learning curves at testing (top panels) and at training (bottom panels) for ERM with SGD (blue) and DRO with our SGD with hardness weighted sampling for different values of β on MNIST (β = 0.1, β = 1, β = 10, β = 100). The models are trained on an imbalanced MNIST dataset (only 1% of the digits 3 kept in the training datatset) and evaluated on the original MNIST testing dataset.

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 = C 6 min 1,

) 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

i=1 exp (βv i ) And by replacing in the formula of the pi ∀i ∈ {1, . . . , n}, pi = 1 n exp (βv i ) exp (βλ -1) = exp (βv i ) n j=1 exp (βv j ) Which corresponds exactly to p = softmax (βv) D.3 PROOF OF LEMMA 4.1: REGULARITY PROPERTIES OF R

y)] D.5 PROOF OF THEOREM 4.2: DISTRIBUTIONALLY ROBUST OPTIMIZATION AS PRINCIPLED HARD EXAMPLE MINING Let D φ an φ-divergence satisfying Assumption 3.1, and v = (v i ) n i=1 ∈ R n . v will play the role of a generic loss vector.

, 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).

This allows us to conclude the proof of our Theorem D.2.D.7.3 CONVERGENCE OF ROBUST SGD WITH EXACT LOSS HISTORYWe 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 = Θ αn 2 ρ βC(L) 2 +2nρC(∇ L) × bδd poly(n,L)m log 2 (m) , with probability at least 1 -exp -Ω(log 2 (m)) over the randomness of the initialization and the mini-batches, Robust SGD with exact loss vector finds ∇ θ (R • f • h)(θ) ≤ after T = O Ln 3 ηδ 2 iterations.

size b ∈ {1, . . . , n}, let S = {i j } b j=1 a batch of indices drawn from p( L) without replacement, i.e. ∀j ∈ {1, . . . b}, i j i.i.d.

inequalities with (88) we obtainE S R(L(h(θ (t+1) )) -R(L(h(θ (t) that compared to (72), there is only the additional term 1 + m d K . Using our choice of η, η = η stale ≤ O δ , η 2 1 + m d K is dominated by the term Ω ηmδ dn 2

concludes the proof using the same arguments as in the end of the proof of Theorem B.1 starting from (76).D.8.1 PROOF OF TECHNICAL LEMMA 2Using Lemma 4.2 and Lemma 4.1 we obtainp( L) -p( L) 2 = ∇ v R( L) -∇ v R( L)

Let A(∇ L) a bound on the gradient of the per-example loss function. Using Theorem D.2 and the chain rule∀j, ∀τ ∇ θ (L j • h j )(θ (τ ) )

(1 -α) kb ≤ exp -Ω log 2 (m) , and thus with probability at least 1 -exp -Ω log 2 (m) , ∀t, t -t i (t) , we finally obtain that with probability at least 1 -exp -Ω log 2 (m) ,p( L) -p( L) 2 ≤ ηC(L)A(∇ L)O βLm 3-pi ( L) ∇ θ (L

finally obtain∇ h (R • L)(h(θ (t) )) )∇ hj L j (h j (θ (t) ))

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 ∆

Table content: <table/>

Jean-Jacques Moreau. Proximité et dualité dans un espace hilbertien. Bulletin de la Société mathématiquede France, 93:273-299, 1965. Hongseok Namkoong and John C Duchi. Stochastic gradient methods for distributionally robust optimization with f-divergences. In Advances in Neural Information Processing Systems, pp. 2208-2216, 2016. Abhinav Gupta, and Ross Girshick. Training region-based object detectors with online hard example mining. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 761-769, 2016. Aman Sinha, Hongseok Namkoong, and John Duchi. Certifying some distributional robustness with principled adversarial training. arXiv preprint arXiv:1710.10571, 2017. Matthew Staib and Stefanie Jegelka. Distributionally robust deep learning as a generalization of adversarial training. In NIPS workshop on Machine Learning and Computer Security, 2017.

Table content: <table><tr><td>A MORE ON OUR EXPERIMENTS</td><td/><td/></tr><tr><td>A.1 EXPERIMENTS ON MNIST</td><td/><td/></tr><tr><td colspan="3">Hassan Rafique, Mingrui Liu, Qihang Lin, and Tianbao Yang. Non-convex min-max optimization:</td></tr><tr><td colspan="3">Provable algorithms and applications in machine learning. arXiv preprint arXiv:1810.02060,</td></tr><tr><td>2018.</td><td/><td/></tr><tr><td colspan="3">Shai Shalev-Shwartz and Yonatan Wexler. Minimizing the maximal loss: How and why. In ICML,</td></tr><tr><td>pp. 793-801, 2016.</td><td/><td/></tr><tr><td>Abhinav Shrivastava, Sergey Zagoruyko and Nikos Komodakis. 0 100 200 300 400 500</td><td>Wide residual networks.</td><td>arXiv preprint</td></tr><tr><td>arXiv:1605.07146, 2016.</td><td/><td/></tr><tr><td colspan="3">Difan Zou and Quanquan Gu. An improved analysis of training over-parameterized deep neural</td></tr><tr><td colspan="2">networks. arXiv preprint arXiv:1906.04688, 2019.</td><td/></tr></table>

Peer review (iclr2020):
Disclaimer: I was able to read the other reviews and the author’s responses before finalizing this review. The paper proposes an easy-to-implement algorithm for DRO on example weights, with the same computational cost as SGD. The algorithm is based on hardness weighted sampling and links are shown to hard example mining. Convergence results are shown for (finitely) wide networks. Additionally, the paper claims to demonstrate the usefulness of the algorithm in deep learning. I am unable to assess how the convergence results place in the literature, but I believe they motivate the algorithm. The claims of practical usefulness in deep learning do not seem supported by the provided empirical evidence. For the CIFAR experiments, the ERM baseline does not seem to train - probably due to the choice of learning rate. This makes it unclear how the proposed algorithm compares. They claim the algorithm is more robust to the learning rate, so is it possible to train with the original learning rate used for the ERM baseline? Why was a different learning rate chosen? If the experimental results showed an improvement over a properly trained ERM baseline on CIFAR I would lean toward weak accept. Many state-of-the-art deep learning pipelines do not use plain SGD - for example, the WideResNet you used on CIFAR. How is Algorithm 1 used on these? Do we only make changes to the update on line 24? Using momentum with nested optimization can introduce instabilities. Perhaps combining momentum with nested optimization of loss weights and parameters is why the baseline does not train? Maybe you could try an architecture that trained using vanilla SGD, so you can better leverage your theoretical results? It would provide evidence of the usefulness of the algorithm if we could take various pipelines and just drop their optimizer updates into algorithm 1 - hopefully without having to spend time re-tuning their optimizer parameters. It would also be nice to see the training loss/accuracy - perhaps over all classes and just cats - in the appendix. Things to improve that did not impact the score: (Page 1) “in term” -> “in terms” (Page 1 & 2) “an history” -> “a history” (Page 2) “Optmization” -> “Optimization” (Page 5) “allows to link” -> “allows us to link” (Page 7) “we focus at” -> “we focus on” (Page 8) “allows to guarantee” -> “allows us to guarantee” (Page 19) “continous” -> “continuous