Revisiting the Stability of Stochastic Gradient Descent: A Tightness Analysis

1 INTRODUCTION . Stochastic gradient descent ( SGD ) has gained great popularity in solving machine learning optimization problems ( Kingma & Ba , 2014 ; Johnson & Zhang , 2013 ) . SGD leverages the finite-sum structure of the objective function , avoids the expensive computation of exact gradients , and thus provides a feasible and efficient optimization solution in large-scale settings ( Bottou , 2012 ) . The convergence and the optimality of SGD have been thoroughly studied ( Ge et al. , 2015 ; Rakhlin et al. , 2012 ; Reddi et al. , 2018 ; Zhou & Gu , 2019 ; Carmon et al. , 2019a ; b ; Shamir & Zhang , 2013 ) . In recent years , new research questions have been raised regarding SGD ’ s impact on a model ’ s generalization power . The seminal work ( Hardt et al. , 2016 ) tackled the problem using the algorithmic stability of SGD , i.e. , the progressive sensitivity of the trained model w.r.t . the replacement of a single ( test ) datum in the training set . The stability-based analysis of the generalization gap allows one to bypass classical model capacity theorems ( Vapnik , 1998 ; Koltchinskii & Panchenko , 2000 ) or weight-based complexity theorems ( Neyshabur et al. , 2017 ; Bartlett et al. , 2017 ; Arora et al. , 2018 ) . This framework also provides theoretical insights into many phenomena observed in practice , e.g. , the “ train faster , generalize better ” phenomenon , the power of regularization techniques such as weight decay ( Krogh & Hertz , 1992 ) , Dropout ( Srivastava et al. , 2014 ) , and gradient clipping . Other works have applied the stability analysis to more sophisticated settings such as Stochastic Gradient Langevin Dynamics and momentum SGD ( Mou et al. , 2018 ; Chaudhari et al. , 2019 ; Chen et al. , 2018 ) . Despite the promises of this stability-based analysis , it remains open whether this framework can explain the strong generalization performance of deep neural networks in practice . Existing theoretical upper bounds of the stability ( and thus , generalization ) ( Hardt et al. , 2016 ) are ideal for strongly convex loss functions : the upper bound remains constant even as the number of training iterations increases . However , the same bound deteriorates significantly when we relax to more general and realistic settings . In particular , for convex ( but not strongly convex ) and non-convex loss functions , if SGD has constant step size , then the upper bound grows linearly and exponentially with the number of training iterations . This bound fails to match the superior generalization performance of deep neural networks , and leads to the following question : Question 1 : Can we find a better stability upper bound for convex or non-convex loss functions ? In this paper , we first address the question above and investigate the tightness of the algorithmic stability analysis for stochastic gradient methods ( SGM ) proposed by ( Hardt et al. , 2016 ) . R1 . We show in Theorem 1 that the analysis in ( Hardt et al. , 2016 ) is tight for convex and smooth objective functions ; in other words , there is a convex loss function whose stability grows linearly with the number of training iterations , with constant step size ( αt = α ) in SGD . R2 . We show that in Theorem 2 that for linear models , the analysis in the convex case can be tightened to show that stab does not increase with t. R3 . In Theorem 3 we show that the analysis in ( Hardt et al. , 2016 ) for decreasing step size ( αt = O ( 1/t ) ) is loose for non-convex objective functions by providing a tighter upper bound on the stability ( and hence generalization error ) . R4 . The bound on the stability of SGD by ( Hardt et al. , 2016 ) is achieved by bounding the divergence at time t , defined as δt : = E||wt − w ′ t|| , where wt is the model trained on data set S and w ′ t is the model trained on a data set S′ that differs from S in exactly one sample . In Theorem 4 we provide evidence that our new upper bound in the non-convex case is tight , by showing a non-convex loss function whose divergence matches the upper bound for our divergence . R5 . Although it is not derived formally , the techniques in ( Hardt et al. , 2016 ) can be employed to show an exponential upper bound for non-convex loss functions minimized using SGD with constant-size step . In Theorem 5 , we give evidence that this abysmal upper bound is likely tight for non-convex loss functions , by exhibiting a non-convex loss function for which the divergence δt increases exponentially . Thus the only functions whose stability provably does not increase with the number of iterations when a constant step-size during SGD is employed , are strongly convex functions . However , a ) it has been empirically observed that for deep neural network loss , near the local minima , the Hessians are usually low rank ( Chaudhari et al. , 2017 ; Yao et al. , 2019 ) , and b ) neural networks trained with constant step-size SGD do generalize well in practice ( Lin & Jegelka , 2018 ; Huang et al. , 2017 ; Smith et al. , 2017 ) . Combined with our lower bounds on convex and non-convex functions , we seem to hit an obstacle on the way to explaining generalization using the stability framework . Question 2 : What is it that makes constant-step SGD on deep learning loss function generalize well ? Realizing the limitation of the current state of stability analysis , we investigate whether a strongerthan-convex , but weaker-than-strongly-convex assumption of the loss function can be made , at least near local minima . If we can show algorithmic stability near local minima , we can still show the stability using similar argument as ( Du et al. , 2019 ; Allen-Zhu et al. , 2019 ) . Aiming towards a characterization of loss functions exhibiting good stability , we propose a new condition for loss near local minima . This condition , called Hessian contractive , is slightly stronger than a general convex condition , but considerably weaker than strongly convex . Formally , the Hessian contractive condition stipulates that near any local minima , ( 1 ) the function is convex ; and ( 2 ) a data dependent Hessian is positive definite in the gradient direction . Theoretically , we show that such a condition is sufficient to guarantee a constant stability bound for SGD ( constant step size ) near the local minima , while allowing the Hessian to be low rank . We also provide examples showing Hessian Contractive is a reasonable condition for several loss functions . Empirically , we verify the Hessian Contractive condition near a local minima of the loss while training deep neural networks . We sample points from a neighborhood of current iterates by adding Gaussian noise and verify the HC condition locally by Hessian product approximation . Summarizing our second set of contributions : R6 . In Observation 1 we show that the family of widely used ( convex ) linear model loss functions will satisfy the Hessian Contractive condition . One typical example of such linear model loss is the regression loss function . These observation suggests that Hessian Contractive is a condition satisfied by ( potentially many ) machine learning loss functions . R7 . In Theorem 6 we show that the Hessian Contractive condition will localize SGD iterates in a neighborhood of minima , which implies a constant stability bound for SGD near the local minima . 1.1 RELATED WORKS . Stability and generalization . The stability framework suggests that a stable machine learning algorithm results in models with good generalization performance ( Kearns & Ron , 1999 ; Bousquet & Elisseeff , 2002 ; Elisseeff et al. , 2005 ; Shalev-Shwartz et al. , 2010 ; Devroye & Wagner , 1979a ; b ; Rogers & Wagner , 1978 ) . It serves as a mechanism for provable learnability when uniform convergence fails ( Shalev-Shwartz et al. , 2010 ; Nagarajan & Kolter , 2019 ) . The concept of uniform stability was introduced in order to derive high probability bounds on the generalization error ( Bousquet & Elisseeff , 2002 ) . Uniform stability describes the worst case change in the loss of a model trained on an algorithm when a single data point in the dataset is replaced . In ( Hardt et al. , 2016 ) , a uniform stability analysis for iterative algorithms is proposed to analyze SGD , generalizing the one-shot version in ( Bousquet & Elisseeff , 2002 ) . Algorithmic uniform stability is widely used in analyzing the generalization performance of SGD ( Mou et al. , 2018 ; Feldman & Vondrak , 2019 ; Chen et al. , 2018 ) . The worst case leave-one-out type bounds also closely connect uniform stability with differential private learning ( Feldman et al. , 2018 ; 2020 ; Dwork et al. , 2006 ; Wu et al. , 2017b ) , where the uniform stability can lead to provable privacy guarantee . While the upper bounds of algorithmic stability of SGD have been extensively studied , the tightness of those bounds remains open . In addition to uniform stability , an average stability of the SGD is studied in Kuzborskij & Lampert ( 2018 ) where the authors provide data-dependent upper bounds on stability1 . In this work , we report for the first time lower bounds on the uniform stability of SGD . Our tightness analysis suggests necessity of additional assumptions for analyzing the generalization of SGD on deep learning . Geometry of local minima . The geometry of local minima plays an important role in the generalization performance of deep neural networks ( Hochreiter & Schmidhuber , 1995 ; Wu et al. , 2017a ) . The flat minima , i.e. , minima whose Hessians have a large portion of zero-valued eigenvalues , are believed to attain better generalization ( Keskar et al. , 2016 ; Li et al. , 2018 ) . In ( Chaudhari et al. , 2019 ) , the authors construct a local entropy-based objective function which converges to a solution with good generalization in a flat region , where “ flatness ” means that the Hessian matrix has a large portion of nearly-zero eigenvalues . However , these observations have not been supported theoretically . In this paper , we propose the Hessian contractive condition that is slight stronger than flat minima . Such condition suggests that the minima is sharp only in the gradient direction while remains flat in other directions , which unifies the geometrical interpretation of flat minima and uniform stability analysis . 2 PRELIMINARIES . In this section we introduce the notion of uniform stability and establish our notations . We first introduce the quantities Empirical and Population Risk and Generalization Gap . Given an unknown distribution D on labeled sample space Z = X × { −1 , +1 } , let S = { z1 , ... , zn } denote a set of n samples zi = ( xi , yi ) drawn i.i.d . from D. Let w ∈ Rd be the parameter ( s ) of a model that tries to predict y given x , and let f be a loss function where f ( w ; z ) denotes the loss of the model with parameter w on sample z . Let f ( w ; S ) denote the empirical risk f ( w ; S ) = Ez∼S [ f ( w ; z ) ] = 1While it is an interesting open problem to get data-dependent lower bounds by lower bounding the average stability , we construct lower bounds on the worst-case stability . Thus our lower bounds are general and not data-dependent . 1 n ∑n i=1 f ( w ; zi ) with corresponding population risk Ez∼D [ f ( w ; z ) ] . The generalization error of the model with parameter w is defined as the difference between the empirical and population risks : |Ez∼D [ f ( w ; z ) ] − Ez∼S [ f ( w ; z ) ] | . Next we introduce the Stochastic Gradient Descent ( SGD ) method . We follow the setting of ( Hardt et al. , 2016 ) , and starting with some initialization w0 ∈ Rd , consider the following SGD update step : wt+1 = wt − αt∇wf ( w ; zit ) , where it is drawn from [ n ] : = { 1 , 2 , · · · , n } uniformly and independently in each round . The analysis of SGD requires the following crucial properties of the loss function f ( . , z ) at any fixed point z , viewed solely as a function of the parameter w : Definition 1 ( L-Lipschitz ) . A function f ( w ) is L-Lipschitz if ∀u , v ∈ Rd : |f ( u ) −f ( v ) | ≤ L‖u−v‖ . Definition 2 ( β-smooth ) . A function f ( w ) is β-smooth if ∀u , v ∈ Rd : |∇f ( u ) −∇f ( v ) | ≤ β‖u−v‖ . Definition 3 ( γ-strongly convex ) . A function f ( w ) is γ strongly convex if ∀u , v ∈ Rd : f ( u ) > f ( v ) +∇f ( v ) > [ u− v ] + γ 2 ‖u− v‖2 . Algorithmic Stability Next we define the key concept of algorithmic stability , which was introduced by ( Bousquet & Elisseeff , 2002 ) and adopted by ( Hardt et al. , 2016 ) . Informally , an algorithm is stable if its output only varies slightly when we change a single sample in the input dataset . When this stability is uniform over all datasets differing at a single point , this leads to an upper bound on the generalization gap . More formally : Definition 4 . Two sets of samples S , S′ are twin datasets if they differ at a single entry , i.e. , S = { z1 , ... zi , ... , zn } and S′ = { z1 , ... , z′i , ... , zn } . Consider a possibly randomized algorithm A that given a sample S of size n outputs a parameter A ( S ) . Define the algorithmic stability parameter of A by : εstab ( A , n ) : = inf { ε : sup z∈Z , S , S′ EA|f ( A ( S ) ; z ) − f ( A ( S′ ) ; z ) | ≤ ε } . Here EA denote expectation over the random coins of A . Also , for such an algorithm , one can define its expected generalization error as : GE ( A , n ) : = ES , A [ Ez∼D [ f ( A ( S ) ; z ) ] − Ez∼S [ f ( A ( S ) ; z ) ] ] . Stability and generalization : It was proved in ( Hardt et al. , 2016 ) that GE ( A , n ) ≤ εstab ( A , n ) . Furthermore , the authors observed that an L-Lipschitz condition on the loss function f enforces a uniform upper bound : supz |f ( w ; z ) − f ( w′ ; z ) | ≤ L‖w − w′‖ . This implies that for Lipschitz loss , the algorithmic stability εstab ( A , n ) ( and hence the generalization error GE ( A , n ) ) can be bounded by obtaining bounds on ‖w − w′‖ . Let wT and w ′ T be the parameters obtained by running SGD starting on twin datasets S and S ′ , respectively . Throughout this paper we will focus on the divergence quantity δT : = EA||wT − w ′ T || . While ( Hardt et al. , 2016 ) reports upper bounds on δT with different types of loss functions , e.g. , convex and non-convex loss functions , we investigate the tightness of those bounds .

Summary: The paper studies the stability of stochastic gradient descent (SGD), which is one of the framework used for explaining generalization. More specifically, the paper investigates the tightness of the algorithmic stability bounds for SGD given by Hardt et al. (2016). Furthermore, the authors propose the Hessian contractive condition, which characterizes deep learning loss functions with good generalization properties, when being trained with SGD.