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: This paper studies stability of SGD which is a popular optimization algorithm. The authors aim to build a tight stability analysis. In particular, the authors show by constructing specific problems that the existing stability bounds for SGD applied to convex problems are tight within a constant factor. Then, the authors refine the stability analysis in non-convex case by presenting new bounds, and show the new bounds are tight. The authors also provide new conditions weaker than strongly convex assumption in both convex and non-convex case. The paper is clearly written and is easy to follow.