SGD Can Converge to Local Maxima

1 INTRODUCTION . SGD is the main optimization algorithm for training deep neural networks , and understanding SGD is widely regarded as a key step on the path to understanding deep learning ( Bottou , 2012 ; Zhang et al. , 2018 ; Xing et al. , 2018 ; Mori et al. , 2021 ; Du et al. , 2018 ; Allen-Zhu et al. , 2018 ; Wojtowytsch , 2021a ; b ; Gower et al. , 2021 ; Ziyin et al. , 2021 ; Gurbuzbalaban et al. , 2021 ; Zou et al. , 2021 ; Li et al. , 2021 ; Feng and Tu , 2021 ) . Despite its algorithmic simplicity ( describable with only two lines of equations ) , SGD is hard to understand . The difficulty is threefold : ( 1 ) SGD is discrete-time in nature , and discrete-time dynamics are typically much more complicated to solve than their continuous-time counterparts ( May , 1976 ) ; ( 2 ) SGD noise is state-dependent ( Ziyin et al. , 2021 ; Hodgkinson and Mahoney , 2020 ) ; and ( 3 ) the loss landscape can be nonlinear and non-convex , involving many local minima , saddle points , and degeneracies ( the Hessian is not full-rank ) in the landscape ( Xing et al. , 2018 ; Wu et al. , 2017 ) . Each of these difficulties is so challenging that very few works attempt to deal with all of them simultaneously . Most previous works on SGD are limited to the cases when the loss landscape is strongly convex ( Ziyin et al. , 2021 ; Liu et al. , 2021 ; Hoffer et al. , 2017 ; Mandt et al. , 2017 ) , or when the noise is assumed to be Gaussian and time-independent ( Zhu et al. , 2019 ; Jastrzebski et al. , 2017 ; Xie et al. , 2021 ) ; for the works that tackle SGD in a non-convex setting , often strong conditions are assumed . The reliance on strong assumptions regarding each of the challenges means that our present understanding of SGD for deep learning could be speculative . This work aims to examine some commonly held presuppositions about SGD and show that when all the three challenging factors are taken together , many counter-intuitive phenomena may arise . Since most of these phenomena are potentially undesired , this work can also be seen as a worse case anaysis of SGD . In this work , we study the behavior of SGD in toy landscapes with non-convex loss and multi-minima . We approach the problem of SGD convergence from a different angle from many of the related works : instead of studying when SGD will converge , our result helps answer the question of when SGD might fail . In particular , the problem setting considers discrete-time SGD close to saddle points , where the noise is due to minibatch sampling , and the learning rate is held constant throughout training . In the next section , we define the SGD algorithm and the necessary notations . In Sec . 3 , we discuss the related works . A warmup example is provided in Sec . 4 . Sec . 5 introduces the main results . Sec . 6 presents the numerical simulations , including a minimal example involving a neural network . We also encourage the readers to examine the appendix . Sec . A presents additional experiments . Sec . C presents a continuous-time analysis of the problems in the main text and is relevant for discussing the unique discrete-time features of SGD . Sec . D presents all the proofs . 2 BACKGROUND . Notation and Terminology . We use λ to denote the learning rate . w is the model parameter . L̂ ( w ; x ) , with the caret symbol ⋅̂ , denotes the sampled loss function , which is a function of the data point , and L ∶= Ex [ L̂ ] is the true loss landscape averaged over the data distribution ; for notational conciseness , we often hide the dependence of L̂ on x when the context is clear . The lowercase t denotes the time step of optimization . In this work , the terms learning rate and step size are used interchangeably . The phrases loss , objective , and potential are also used interchangeably . Now , we introduce the minibatch SGD algorithm . The objective of SGD can be defined as a pair ( L̂ , p ( x ) ) such that L̂ is a differentiable function and p ( x ) is a probability density . We aim at finding the minimizer of the following differentiable objective : L ( w ) = Ex∼p ( x ) [ L̂ ( w ; x ) ] , ( 1 ) where x is a data point drawn from distribution p ( x ) and w ∈ RD is the model parameter . The gradient descent ( GD ) algorithm for L is defined as the update rule wt = wt−1 − λ∇wL ( w , x ) for some randomly initialized w0 , where λ is the learning rate . The following definition defines SGD . Definition 1 . The minibatch SGD algorithm computes the update to the parameter w with the following equation : wt = wt−1 − λĝt , where ĝt = ∇L̂ ( wt−1 ; xt ) , and xt is drawn from p ( x ) such that xi is independent of xj for i ≠ j . 3 RELATED WORKS . Escaping Saddle Points . Because neural networks are highly nonlinear , the landscape of neural networks is believed to have many local minima and saddle points ( Du et al. , 2017 ; Kleinberg et al. , 2018 ; Reddi et al. , 2018b ) . It is thus crucial for any optimization algorithm to be able to overcome saddle points efficiently and not become stuck in a suboptimal minimum . However , the majority of works on escaping saddle points assume no stochasticity or only artificially injected noise ( Du et al. , 2017 ; Jin et al. , 2017 ; Ge et al. , 2015 ; Reddi et al. , 2018b ; Lee et al. , 2016 ) . Some physics-inspired approaches take the stochasticity into account but often rely on continuous-time approximation and nontrivial assumptions regarding the stationary distribution of the parameter ( Mori et al. , 2021 ; Xie et al. , 2021 ; Liu et al. , 2021 ) . In the context of stochastic non-convex optimization , one setting that is relatively well-understood is when the learning rate decreases to zero in the training ( Pemantle , 1990 ; Vlaski and Sayed , 2019 ; Mertikopoulos et al. , 2020 ) . For example , Pemantle ( 1990 ) shows that stochastic approximation algorithms can avoid unstable stationary points ; Vlaski and Sayed ( 2019 ) shows that SGD avoids saddle points in expectation , Mertikopoulos et al . ( 2020 ) shows that SGD avoids saddle points in probability ; the later two also provide precise conditions under which the saddles points are avoided . These work , therefore , differ from the present work in that , while they assume that the learning rate decrease to zero , we only consider the case when the learning rate is held fixed . It is important to note that our result does not contradict these previous results . To the opposite , our result may be interpreted as reinforcing the crucial messages in these previous works , namely , that the learning rate needs to be chosen and scheduled carefully to guarantee convergence . Role of Minibatch Noise . One indispensable aspect of SGD is its special stochasticity due to minibatch sampling ( Wu et al. , 2020 ; Ziyin et al. , 2021 ; Zhu et al. , 2019 ; Hodgkinson and Mahoney , 2020 ) , whose strength and structure are determined by the model architecture and the distribution of the data points . In fact , it has been shown that the models trained with SGD significantly outperform the models directly trained with GD ( Hoffer et al. , 2017 ) , which suggests that minibatch noise in SGD is beneficial for training . Understanding the role of minibatch noise is therefore of fundamental and practical value . However , in the previous works , the focus is on analyzing the behavior of SGD on the landscape specified by L ; the unique structures due to the noise in SGD are often treated in an oversimplified manner without referring to the actual models in deep learning or the actual data distributions . For example , Daneshmand et al . ( 2018 ) assumes that the noisy gradient is negatively correlated with the direction of the Hessian with the smallest eigenvalue , the validity of which is unclear when the Hessian is not full rank . Kleinberg et al . ( 2018 ) assumes that the noise in SGD makes the landscape one-point strongly convex , which implies that the loss landscape before the convolution is already close to a strongly convex landscape . Another commonly assumed condition is that the loss-landscape satisfies the Lojasiewicz condition , which is non-convex but implies the unrealistic condition that there is no local minimum except the global minimum ( Karimi et al. , 2020 ; Wojtowytsch , 2021a ; Vaswani et al. , 2019 ) , which is not sufficient to understand the complicated dynamics that often happen in a setting with many minima . In Sec . 6.3 , we show that these assumptions are violated for a minimal toy nonlinear neural network with two layers and one hidden neuron . While our example may or may not be directly relevant for realistic settings in the deep learning practice , our work highlights the importance of analyzing the actual noise structure imposed by a deep neural network in future works . Large Learning Rate Regime . Recently , it is realized that networks trained at a large learning rate have a dramatically better performance than networks trained with a vanishing learning rate ( lazy training ) ( Chizat and Bach , 2018 ) . Lewkowycz et al . ( 2020 ) shows that there is a qualitative difference between the lazy training regime and the large learning rate regime ; the performance features two plateaus in testing accuracy in the two regimes , with the large learning rate regime performing much better . However , almost no previous theory exists for understanding SGD at a large learning rate . Our work is also relevant for understanding what happens at a large learning rate . 4 A WARM-UP EXAMPLE . This section studies a special example of the main results of our work . While the setting of this example is restricted , it captures the essential features and mechanisms of SGD that we will utilize to prove the main results . Consider the following loss function : L̂ ( w ) = x 2 w2 , ( 2 ) where x , w ∈ R and x is drawn from an underlying distribution p ( x ) at every time-step . We let E [ x ] = a and assume that Var [ x ] = σ2 is finite . This gives rise to a true , sample-averaged loss of L ( w ) = a 2 w2 . The stationary point of this loss function is w = 0 . When a > 0 , the underlying ( deterministic ) landscape is a minimum ; the a > 0 case is now well understood in the discrete-time limit and when the underlying noise is state-dependent ( Liu et al. , 2021 ) . When a < 0 , the point w = 0 is a local maximum ; one hopes that the underlying dynamics escape to infinity , and our goal is to understand the behavior of SGD in this case . The following proposition and proof show that for the seemingly benign data distribution p ( x ) ( zero-mean and bounded ) above , SGD can not escape from the local maximum when the learning rate is set to be 1 . Proposition 1 . Let λ = 1 and a < 0 and p ( x ) be the distribution such that x = 1 and x = −1 + a with equal probability . Then , w converges to w = 0 with probability 1 with the SGD algorithm , independent of the initialization . Proof . By the definition of the SGD algorithm , we have wt+1 = { 0 with probability 0.5 ; ( 2 − a ) wt otherwise . ( 3 ) Therefore , after t time steps , wt has at most 2−t probability of being non-zero . For infinite t , wt = 0 with probability 1 . ◻ This simple toy example illustrates a few important points . First of all , it suggests that one can not use the expected value of w to study the escape problem of SGD . The expectation value of wt is Ex [ wt ] = ( 1 − a/2 ) Ex [ wt−1 ] = ( 1 − a/2 ) tw0 = w0eln ( 1−a/2 ) t , ( 4 ) which is an exponentially fast escape from any initialization w0 for a < 0 . In fact , ln ( 1− a ) is exactly the escape time scale of the GD algorithm . This is counter-intuitive : SGD will converge to the local maximum with probability 1 , even if its expected value escapes the local maximum exponentially fast . Alternatively , one might also hope to use the expected loss or the norm Ewt [ L ( wt ) ] ∼ ∣∣wt∣∣22 as the metric of escape , but it is not hard to see that they suffer the same problem . At time t , because wt = 0 with probability 1− 2−t and wt = ( 2− a ) tw0 with probability 2−t , we have E [ L ( wt ) ] ∝ w20e 2t ln 2−a√ 2 , which also predicts an exponentially fast escape . Again , this fails to reflect the fact that SGD has only a vanishingly small probability of escaping the constructed local maximum . Statistically speaking , the norms of wt fail to be good metrics to measure the escape rate because these metrics are not robust against outliers . In our example , there is only 2−t probability for SGD to escape the local minimum , yet the speed at which this outlier escapes the maximum outweighs the decay in its probability through time , thus contributing more to the expected norm than any other events . This example also shows the difficulty and subtlety of dealing with the escape problem . This example also emphasizes the importance of showing convergence in probability for future works in non-convex optimization because convergence in expectation is not sufficient to guarantee a good empirical performance . The results in the rest of this work can be seen as extensions of this special case to less restrictive settings .

Summary: ## Post-discussion reassessment  See relevant post below.  ## Summary of contributions  In this paper, the authors examine whether SGD with a large, constant step-size avoids local maximizers (or, more generally, undesirable saddle points of the underlying minimization problem). More precisely, they focus on the algorithm $$ w_{t+1} = w_t - \lambda \hat g_t $$ where $\lambda>0$ is the algorithm's step-size, and $\hat g_t$ is a (stochastic) gradient of the (stochastic) loss function $\hat L(w;x)$, with $x$ a random variable.  The paper's results can be summarized as follows:  1. If $\hat L(w;x) = (x/2) \cdot w^2$, the authors identify a range of values of $\lambda$ (which depends on the distribution of $x$) such that, in probability, $w_t$ converges to $0$ – which, under the specified distributional assumptions for $x$, is the global maximum of $L = \mathbb{E}[\hat L]$. This is made precise in Propositions 1 and 2, and Corollary 1.  2. They provide a quartic loss function under which SGD converges to the function's sharper minimizers (as measured by the trace of the Hessian at said points). [Proposition 4]  3. They provide a specific range of parameters under which the AMSGrad algorithm converges to the undesirable maximizer of item (1) above.  [In the supplement, the authors also provide an analysis of a gradient-like diffusion (Appendices B and C), which they discuss as a continuous-time model of (SGD). This part is not directly connected to the rest of the paper, so I am not including it in my evaluation below.]