Quasi-potential theory for escape problem: Quantitative sharpness effect on SGD's escape from local minima

1 INTRODUCTION . In recent years , the successes of deep learning have been a major driving force of machine learning development ( LeCun , 2019 ) . Owing to its strong generalization capability , deep learning has diverged into a wide range of domains , such as computer vision ( Krizhevsky et al. , 2012 ) , speech recognition ( Mikolov et al. , 2011 ) , and natural language processing ( Collobert et al. , 2011 ) . The high performance of deep learning is underpinned by gradient-based learning algorithms , including stochastic gradient descent ( SGD ) and its variations ( Kingma & Ba , 2014 ; Schmidt et al. , 2021 ) . However , at the same time , those unprecedented successes raise a question : Why does SGD learn parameters of neural networks with high generalization performance ? Although the optimization problems of neural networks were thought to be difficult to solve ( Blum & Rivest , 1992 ) , SGD can find nearly optimal solutions empirically , and further , the obtained solutions generalize well ( Keskar et al. , 2016 ; Brutzkus et al. , 2017 ) . Analyzing SGD ’ s role on deep learning is an area of research that is currently attracting strong interest ( Masters & Luschi , 2018 ; Jastrzebski et al. , 2021 ) . One of the promising directions for this question is to study the geometric properties of a training loss landscape . Many empirical studies have found that minima obtained by SGD have distinctive geometric properties . Keskar et al . ( 2016 ) have shown that the shape of the minima obtained by SGD tends to be flat . He et al . ( 2019b ) have deepened the investigation by picturing that SGD settles on the flatter side of asymmetric loss surface , which they named “ asymmetric valley . '' Draxler et al . ( 2018 ) and Garipov et al . ( 2018 ) have shown that separate minima obtained by independent training processes are internally connected through pathways . Li et al . ( 2017 ) have proposed a dimension reduction technique to visualize the geometry of loss surfaces , visually confirming flat minima . Most significantly , Jiang et al . ( 2019 ) conducted large-scale experiments and verified that minima in flat and wide regions have the strongest correlations with generalization capabilities . To attain a theoretical understanding of SGD , it is key to quantitatively analyze the connection between SGD and the geometric properties of loss surfaces . An escape problem is a scheme to analyze the dynamic of SGD escaping from local minima ( Zhu et al. , 2019 ; Jastrzębski et al. , 2017 ; Hu et al. , 2019 ; Nguyen et al. , 2019 ; Xie et al. , 2020 ) . This scheme allows us to investigate why SGD avoids ( potentially ) bad local minima and settles on good minima . Zhu et al . ( 2019 ) first investigated the SGD ’ s escape phenomenon and showed that SGD ’ s escape is enhanced by its unique noise structure , called the “ anisotropic noise structure . '' Invoked by their analysis , many studies have been attempting to theoretically quantify this phenomenon . Hu et al . ( 2019 ) rigorously identified the role of learning rate in escaping . Nguyen et al . ( 2019 ) used the Levy process to provide the precise description of SGD as well as its escaping phenomena . Jastrzębski et al . ( 2017 ) developed a theory of stochastic differential equation and quantified how the anisotropic noise affects its fast escape from sharp minima . Xie et al . ( 2020 ) refined the mathematical aspect and showed that the SGD ’ s noise structure exponentially enhances escaping under a setup of diffusion theory . In this paper , we introduce a quasi-potential theory to the escape problem , and investigate a mean exit time , which formally quantifies escaping . The notion of quasi-potential is defined in a fundamental theory of stochastic dynamical systems , named a Large Deviation Theory ( Freidlin & Wentzell , 2012 ; Dembo & Zeitouni , 2010 ) . Quasi-potential can formulate the distribution of trajectories that a stochastic dynamical system takes . To illustrate quasi-potential for SGD ’ s escaping problem , we introduce an intuitive notion , steepness of a trajectory ( Fig . 1 and Definition 3 ) , and show that it is an effective tool to analyze the SGD ’ s escaping . To the best of our knowledge , this is the first work that applies the quasi-potential to formalize the relationship between SGD ’ s escape and sharpness . Our main findings and contributions are as follows : • We develop a novel quasi-potential theory that rigorously describes the escape of SGD with no auxiliary variables . Our theory can concisely incorporate the effect of essential factors , i.e. , a batch size ( B ) , a learning rate ( η ) , and geometric parameters of loss surfaces ( r and λ ) • Our theory can be flexibly applied to several practical setups : SGD with discrete update and state-dependent gradient noise , while those were sometimes omitted in previous works for mathematical convenience . • We find that a loss surface with sharp minima slows down the escape of SGD , which seems to contradict the common knowledge , i.e. , SGD escapes efficiently from sharp minima . We show that our result does not contradict the common knowledge but is a generalization of existing results , such as ( Jastrzębski et al. , 2017 ) and ( Xie et al. , 2020 ) . 1.1 COMPARISON WITH EXISTING STUDIES ON ESCAPE PROBLEM . In Table 1 , we compare the escape time derived with the results of other studies that analyze the escape problem . There are two main points of focus . The first is the time to escape that we derive . Our theory provides a unified analysis of exit time incorporating all the essential parameters , batch size , learning rate , a radius of the region around a minimum , and sharpness of a minimum . As a consequence , we show that the eigenvalues of the Hesse matrix increase the time to escape , which has not been found in other studies . Such effect was less apparent in previous studies such as ( Jastrzębski et al. , 2017 ) and ( Xie et al. , 2020 ) because part of the sharpness dependency is hidden in ∆L . Provided that ∆L = r2λ under Assumption 1 , one can see that our results are consistent with previous studies . The second is our theory ’ s flexibility . Different from Jastrzębski et al . ( 2017 ) and Xie et al . ( 2020 ) , our theory does not require that SGD reaches the stationary distribution before escaping , which is known to take exponentially many iterations ( Xu et al. , 2017 ; Raginsky et al. , 2017 ) . Additionally , our theory can evaluate the correspondence with the practical SGD , which has a discrete update rule and state-dependent noise . Notations : For a k × k matrix A , λj ( A ) is the j-th largest eigenvalue of A , and λmax ( A ) = λ1 ( A ) and λmin ( A ) = λk ( A ) denote the largest and the smallest eigenvalue of a square matrix . O ( · ) denotes Landau ’ s Big-O notation . ‖·‖ denotes the Euclidean norm . Given a time-dependent function θt , θ̇t denotes the differentiation of θt with respect to t. N ( µ , Σ ) denotes the multivariate Gaussian distribution with the mean µ , and the covariance Σ . 2 SETTING AND PROBLEM . 2.1 STOCHASTIC GRADIENT DESCENT AND DYNAMICAL SYSTEM . Consider a learning model parameterized by θ ∈ Rd . Given training examples { xi } Ni=1 and a loss function ℓ ( θ , xi ) , we consider a training loss L ( θ ) : = 1N ∑N i=1 ℓ ( θ , xi ) and a mini-batch loss LB ( θ ) : = 1B ∑ xi∈B ℓ ( θ , xi ) , where B is a randomly sampled subset of the training data such that |B| = B . We assume that L ( θ ) is differentiable and its derivative ∇L ( θ ) is Lipschitz continuous . We mainly consider two types of SGD : a discrete SGD and a continuous SGD . Although a discrete SGD is used in practice , we study continuous SGD as a starting point of our analysis because of its mathematical convenience . This is a widely used approach in general SGD analyses ( Ali et al. , 2019 ; Advani et al. , 2020 ) as well as in the escaping analyses ( Jastrzębski et al. , 2017 ; Xie et al. , 2020 ) . Discrete SGD : First , we give the usual discrete formulation of SGD . Given an initial parameter θ0 ∈ Rd , SGD generates a sequence of parameters { θk } k∈N by the following update rule : θk+1 = θk − η∇LB ( θk ) , ( 1 ) for k ∈ N , where η > 0 is a learning rate . In particular , we focus on SGD whose noise on gradients has a Gaussian distribution . We decompose −∇LB ( θk ) in ( 1 ) into a gradient term −∇L ( θk ) and a noise term ∇L ( θk ) −∇LB ( θk ) , and consider a case that the noise is Gaussian . With this setting , the update rule in ( 1 ) is rewritten as θk+1 = θk − η∇L ( θk ) + √ η B Wk , ( 2 ) where Wk ∼ N ( 0 , ηC ( θk ) ) is a parameter-dependent Gaussian noise with its covariance C ( θ ) : = Ei∼Uni ( { 1 , ... , N } ) [ ( ∇L ( θ ) −∇ℓ ( θ , xi ) ) ⊤ ( ∇L ( θ ) −∇ℓ ( θ , xi ) ) ] .We assume that C ( θ ) is Lipschitz continuous . The Gaussianity of the noise on gradients is justified by the following reasons : ( i ) if the batch sizeB is sufficiently large , the central limit theorem ensures the noise term becomes Gaussian noise , and ( ii ) several empirical studies show that the noise term becomes Gaussian noise ( Mandt et al. , 2016 ; Jastrzębski et al. , 2017 ; He et al. , 2019a ) , although different findings have been obtained in other settings ( Simsekli et al. , 2019 ) . Continuous SGD : We also give a continuous SGD , which is exactly discretized to ( 2 ) by a classic Euler scheme ( Definition 5.1.1 of Gobet ( 2016 ) ) . With a time index t ≥ 0 and the given initial parameter θ0 ∈ Rd , the continuous dynamic of SGD is written as follows : θ̇t = −∇L ( θt ) + √ η B C ( θt ) 1/2 wt ( 3 ) where wt is a d-dimensional Wiener process , i.e . an Rd-valued stochastic process with t such that w0 = 0 and wt+u−wt ∼ N ( 0 , uI ) for any t , u > 0 . We note this system can be seen as a Gaussian perturbed dynamical system with a noise magnitude √ η B because η and B do not evolve by time .

Summary: This paper uses the quasi potential theory to formalize the escape behavior of SGD happening while training a deep neural network. The quasi potential is defined based on the steepness of a trajectory and is the smallest steepness to go from the minimizer to another point. Utilizing this they analyze both continuous and discrete SGD and show that the average exit time for a trajectory of SGD to move out of the neighborhood of minima depends exponentially on the mini-batch size, the sharpness of the minima, the radius of the neighborhood, and the inverse of learning rate.