SGD with Hardness Weighted Sampling for Distributionally Robust Deep Learning

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 { ( 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 can not 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 . 2 RELATED WORK IN DRO WITH A WASSERSTEIN DISTANCE . 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 ) .

Summary: 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