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: Many theoretical works have studied SGD, but they commonly rely on restrictive and unrealistic assumptions about the noise. In this work, the authors construct example optimization problems illustrating that, if these assumptions are relaxed, SGD can exhibit many strange behaviors, including (1) SGD can converge to local maxima, (2) SGD may escape saddle points arbitrarily slowly, (3) SGD can prefer sharp minima over flat ones, and (4) AMSGrad can converge to local maxima. Therefore, the authors conclude that in the most general nonconvex case, many counter-intuitive phenomena of SGD may arise and contrast to the commonly held presumptions.