Distributionally Robust Optimization Leads to Better Generalization: on SGD and Beyond

1 INTRODUCTION . In recent years , supervised learning based on deep neural networks ( DNNs ) has achieved state-ofthe-art performances in various tasks of different domains , such as computer vision and natural language processing ( Goodfellow et al. , 2016 ) . However , theoretical explanations for the generalization ability of deep learning remain challenging . Roughly speaking , supervised learning aims to learn a model from the training data . Let Θ denote the parameter space of the model , and Z denote the sample space . Let ` ( z , θ ) : Z ×Θ→ R+ denote the loss function related to the parameter θ and the instance z . In the general setting of supervised learning , there is a fixed but unknown distribution over Z ( denoted P ) . We attempt to find a θ ∈ Θ , which gives a small generalization error . That is , we address the following optimization problem : min θ∈Θ R ( θ ) , EP ` ( z , θ ) . ( 1 ) However , the real distribution P is unknown , so we consider an empirical distribution . Given the training samples S = ( z1 , z2 , · · · , zm ) , the empirical distribution is defined as P̂m ( z ) = 1 m ∑m i=1 I [ z=zi ] , where I is the indicator function . The loss related to this distribution is called the empirical loss , which is R̂ ( θ ) , EP̂m ` ( z , θ ) = 1 m m∑ i=1 ` ( zi , θ ) , ( 2 ) and we minimize the empirical loss instead . This procedure is referred to as Empirical Risk Minimization ( ERM ) . The overwhelming capacity and the paucity of data make deep learning architectures susceptible to over-fitting . This urges us to propose a better approach to the minimization problem ( 1 ) . Suppose there is a distribution familyM = { Pλ : λ ∈ Λ } parametrized by λ , and the underlying distribution P = Pλ0 ∈ M , where λ0 is fixed but unknown . It is a common case in deep learning tasks that the structure of the data is complicated while the amount of data is small . Hence , when we use an estimator λ̂ to estimate the parameter λ0 , the variance of λ̂ would be quite large . Thus , rather than directly minimizing the true expected loss R ( θ ) = EPλ0 ` ( z , θ ) , we minimize a distributionally robust loss alternatively , which can ensure the true expected loss to be small with high probability . The distributionally robust loss is defined as : R ( θ , K ) , sup λ∈K EPλ ` ( z , θ ) , where K ⊂ Λ is a neighbourhood of the estimator λ̂ . This procedure is called Distributionally Robust Optimization ( DRO ) . Because we only have access to the given training dataset , we construct an unbiased1 estimation of the expected loss EPλ ` ( z , θ ) : 1 m m∑ i=1 Pλ ( zi ) Pλ0 ( zi ) ` ( zi , θ ) . This can be viewed as the expected loss with respect to a perturbed empirical generalized distribution2 P̂λ = 1m ∑m i=1 Pλ ( z ) Pλ0 ( z ) I [ z=zi ] , so the DRO problem becomes minimizing the empirical distributionally robust loss : R̂ ( θ , K ) , sup λ∈K EP̂λ ` ( z , θ ) = sup λ∈K m∑ i=1 Pλ ( zi ) mPλ0 ( zi ) ` ( zi , θ ) . Keeping the motivations above in mind , we focus our attention on the robust loss , defined as the maximal possible weighted loss over a pre-specified weight set P for simplicity : R̂S , P ( θ ) , sup p∈P m∑ i=1 pi ` ( zi , θ ) . In deep learning , the high capacity of the model and the non-convexity of the objective function lead to an extremely complicated landscape and induce a lot of minima , and different minima have different generalization performances . In this case , an optimization algorithm is both an optimizer and a minimum selector , and the minimum preference of the optimization algorithm is crucial in deep learning . Stochastic Gradient Descent ( SGD ) and its variants are widely used to minimize the empirical loss ( 2 ) . The simplest mini-batch SGD consists of iterative sampling of a batch of data and then updating the parameter : θt+1 = θt − ηt ( 1 |Bt| ∑ zi∈Bt ∇θ ` ( zi , θt ) ) , where Bt ⊂ S is the sampled mini-batch , ηt is the learning rate , and |Bt| denotes the batch size . We interpret SGD from another perspective : at each iteration , SGD randomly chooses an empirical distribution over the data and performs a gradient descent update w.r.t . the corresponding loss . This interpretation shows a profound connection between SGD and DRO . Motivated by the connection , we develop a brand new version of SGD that can achieve a better generalization performance . The main contribution of this paper is summarized as follows : • We clarify the connection between mini-batch SGD and DRO . We show that SGD helps escape from bad stationary points and the robust loss of local minima where SGD converges is not large . • We analyze the generalization property of robust loss and show the trade off : the robust loss generalizes well when P is relatively small , but a larger P is beneficial to reducing the localized Rademacher complexity . • We propose a framework of improved SGD algorithm , which enables the algorithm to explore more empirical distributions over the data . We give a practical improved SGD algorithm that can directly optimize the robust loss . 1To make the fraction reasonable , we assume Pλ0 ( z ) > 0 whenever Pλ ( z ) > 0 . 2Strictly speaking , this is not a distribution because P̂λ ( Z ) may not equals to 1 , but its expectation is 1 ( the expectation is taken over the sample S ) . Nonetheless , we can view P̂λ as a measure over Z . 2 RELATED WORK . SGD is widely used in large-scale machine learning problems . The convergence theory of SGD in convex cases is solid ( Ghadimi & Lan , 2013 ; Shamir & Zhang , 2013 ) . However , due to the nonconvexity of deep learning models , the behaviour of SGD is more complicated . Dauphin et al . ( 2014 ) argued that it is the prevalence of saddle points rather than local minima that causes the difficulty in high dimensional non-convex optimization . Ge et al . ( 2015 ) analyzed the property of perturbed GD ( i.e. , gradient descent with an isotropic noise injected ) . They showed that perturbed GD has the ability to escape from saddle points in a polynomial number of iterations , and then converges to a second-order stationary point . It is also well known that adding perturbation and noise at different levels of the training processes helps generalization , such as dropout ( Srivastava et al. , 2014 ) , and random noise injection to the input data ( Chapelle et al. , 2001 ) . The good performance of SGD might benefit from the noise . Keskar et al . ( 2016 ) discovered that large-batch training will incur a degradation in the generalization performance of the model . They explained this phenomenon with the concept of sharp minima , and carried out numerical experiments to support the idea that SGD with large batch size tends to converge to sharp minima , which leads to a bad generalization ability . However , Dinh et al . ( 2017 ) provided theoretical results indicating that most existing definitions of sharpness and flatness of minima can not be applied to some commonly used non-negative homogeneous activation functions such as ReLU . DRO was originally applied to decision making under uncertainty . Namkoong & Duchi ( 2017 ) introduced DRO to machine learning , and viewed it as a variance-based regularization that allows us to automatically trade-off between bias and variance . They proved an O ( 1m ) approximation to the generalization errorR ( θ ) , in contrast to the generalO ( 1√ m ) approximation of the empirical risk . Recently , DRO is used to achieve fairness without demographic information in fair machine learning ( Hashimoto et al. , 2018 ) . In learning theory , Rademacher complexity is widely used to measure the complexity of a hypothesis set ( Bartlett & Mendelson , 2002 ) . Different from the general Rademacher complexity which considers the entire hypothesis set , localized Rademacher complexity only cares about the near optimal hypothesis ( Bartlett et al. , 2002 ) . In fact , the hypothesis selected by a learning algorithm usually has better performance than the worst hypothesis . Thus , localized Rademacher complexity is usually much smaller than the general Rademacher complexity . 3 NOTATION AND PRELIMINARIES . In our work , we assume the loss function ` ( z , θ ) is bounded between 0 and 1 , and is continuously differentiable . Let Σ denote the hyperplane { p : ∑m i=1 pi = 1 } and Σ+ denote the simplex { p : ∑m i=1 pi = 1 , p ≥ 0 } . Denote p0 = ( 1 m , . . . , 1 m ) > . For p = ( p1 , p2 , . . . , pm ) > ∈ Σ , let R̂S , p ( θ ) denote ∑m i=1 pi ` ( zi , θ ) . Then the ERM loss can be written as R̂S , p0 ( θ ) , and the robust loss w.r.t . the weight set P can be viewed as : R̂S , P ( θ ) = sup p∈P R̂S , p ( θ ) . Let conv ( P ) denote the convex hull of P . Given a v ∈ Σ , let σ ( v ) be the set of all the vectors whose entries are permutations of v ’ s . Let P ( v ) denote the convex hull of σ ( v ) . We define P ( k ) as the set of all the vectors that have k 1k ’ s and ( m− k ) 0 ’ s : P ( k ) = { p = 1 k v : v ∈ { 0 , 1 } m , m∑ i=1 vi = k } . Clearly , P ( k ) has a natural connection to SGD with batch size k : for any batch sampled from the dataset , there is a corresponding p ∈ P ( k ) , where pi = 1k if the i-th instance is in the batch . In the remainder of the paper , we assume that P is symmetric , i.e. , if p ∈ P , then σ ( p ) ⊂ P . This follows that p0 ∈ conv ( P ) . And we assume that P is a compact set . We assume the entries of p ∈ P sums to 1 , but we relax the restrict that each entry of p ∈ P is non-negative , i.e. , we assume P ⊂ Σ . To describe the properties of the set P , we use the notation RADp ( P ) to denote the minimum radius of the Lp ball centered at p0 that contains P , and we use radp ( P ) to denote the maximum radius of the Lp ball centered at p0 whose intersection with Σ is contained within P . We denote ‖P‖p = supp∈P ‖p‖p , and ‖P‖0 as the maximum number of non-zero entries of p ∈ P . Definition 1 . A function f : Θ→ R is L-Lipschitz , if : |f ( θ1 ) − f ( θ2 ) | ≤ L‖θ1 − θ2‖ , ∀θ1 , θ2 ∈ Θ . Definition 2 . A function f : Θ→ R is µ-strongly convex if f ( θ ) − µ2 ‖θ‖ 2 is convex . Some basic properties of the robust loss R̂S , P ( θ ) are summarized in Appendix A .

Summary: This paper consider the connections between SGD and distributionally robust optimization. There has long been observed a connection between robust optimization and generalization. Recently, this has been explored through the lens of distributionally robust optimization. e.g., in the papers of Namkoong and Duchi, but also many others, e.g., Farnia and Tse, etc. Primarily, this paper appears to build off the work of Namkoong.