Faster Perturbed Stochastic Gradient Methods for Finding Local Minima

Escaping from saddle points and finding local minimum is a central problem in1 nonconvex optimization . Perturbed gradient methods are perhaps the simplest2 approach for this problem . However , to find ( , √ ) -approximate local minima , the3 existing best stochastic gradient complexity for this type of algorithms is Õ ( −3.5 ) ,4 which is not optimal . In this paper , we propose Pullback , a faster perturbed5 stochastic gradient framework for finding local minima . We show that Pullback6 with stochastic gradient estimators such as SARAH/SPIDER and STORM can7 find ( , H ) -approximate local minima within Õ ( −3 + −6H ) stochastic gradient8 evaluations ( or Õ ( −3 ) when H = √ ) . The core idea of our framework is a9 step-size “ pullback ” scheme to control the average movement of the iterates , which10 leads to faster convergence to the local minima.11 1 Introduction12 In this paper , we focus on the following optimization problem13 min x∈Rd F ( x ) : = Eξ [ f ( x ; ξ ) ] , ( 1.1 ) where f ( x ; ξ ) : Rd → R is a stochastic function indexed by some random vector ξ , and it is14 differentiable and possibly nonconvex . We consider the case where only the stochastic gradients15 ∇f ( x ; ξ ) are accessible . ( 1.1 ) can unify a variety of stochastic optimization problems , such as finite-16 sum optimization and online optimization . Since in general , finding global minima of a nonconvex17 function could be an NP-hard problem [ 12 ] , one often seeks to finding an ( , H ) -approximate local18 minimum x , i.e. , ‖∇F ( x ) ‖2 ≤ and λmin ( ∇2F ( x ) ) ≥ − H , where ∇F ( x ) is the gradient of F19 and λmin ( ∇2F ( x ) ) is the smallest eigenvalue of the Hessian of F at x . In many machine learning20 applications such as matrix sensing and completion [ 5 , 11 ] , it suffices to find local minima due to the21 fact that all local minima are global minima.22 For the case where f is a deterministic function , it has been shown that vanilla gradient descent fails23 to find local minima efficiently since the iterates will get stuck at saddle points for exponential time24 [ 8 ] . To address this issue , the simplest idea is to add random noises as a perturbation to the stuck25 iterates . Jin et al . [ 13 ] showed that the simple perturbation step is enough for gradient descent to26 escape saddle points and find ( , √ ) -approximate local minima within Õ ( 1/ 2 ) gradient evaluations,27 which matches the number of gradient evaluations for gradient descent to find -stationary points28 [ 19 ] . Such matching results suggest that perturbed gradient methods can find local minima efficiently,29 at least for deterministic optimization . When it comes to stochastic optimization , a natural question30 arises:31 Can perturbed stochastic gradient methods find local minima efficiently ? 32 To answer this question , we first look into existing results of perturbed stochastic gradient methods33 for finding local minima . Ge et al . [ 10 ] showed that perturbed Stochastic gradient descent can find34 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . ( , √ ) -approximate local minima within Õ ( poly ( −1 ) ) stochastic gradient evaluations . Daneshmand35 et al . [ 7 ] showed that under a specific CNC condition , stochastic gradient descent is able to find36 ( , √ ) -approximate local minima within Õ ( 1/ 5 ) stochastic gradient evaluations . Later on , Li [ 17 ] 37 showed that simple stochastic recursive gradient descent ( SSRGD ) can find ( , √ ) -approximate38 local minima within Õ ( −3.5 ) stochastic gradient evaluations , which is the state-of-the-arts to date.39 However , none of these results by perturbed stochastic gradient methods matches the optimal result40 Õ ( −3 ) for finding -stationary points , achieved by SPIDER [ 9 ] , SNVRG [ 30 ] and STORM [ 6 ] ( See41 also Arjevani et al . [ 4 ] for the lower bound results ) . Therefore , whether perturbed stochastic gradient42 methods can find local minima as efficiently as finding stationary points still remains unknown.43 In this work , we give an affirmative answer to the above question . We propose a general framework44 named Pullback , which works together with existing popular stochastic gradient estimators such45 as SARAH/SPIDER and STORM to find approximate local minima efficiently . We summarize our46 contributions as follows:47 • We prove that Pullback finds ( , H ) -approximate local minima within Õ ( −3 + −6H ) stochastic48 gradient evaluations . Specifically , in the classic setting where H = √ , our Pullback together49 with the SARAH/SPIDER estimator enjoys an Õ ( −3 ) stochastic gradient complexity , which50 outperforms previous best known complexity result Õ ( −3.5 ) achieved by Li [ 17 ] . Our result51 also matches the best possible complexity result Õ ( −3 ) achieved by negative curvature search52 based algorithms [ 9 , 31 ] , which suggests that simple methods such as perturbed stochastic gradient53 methods can find local minima as efficiently as the more complicated ones.54 • Besides , we also show that Pullback with a recent proposed STORM estimator is also able to find55 ( , H ) -approximate local minima within Õ ( −3 + −6H ) stochastic gradient evaluations.56 • At the core of our Pullback is a novel step-size `` pullback '' scheme to control the average movement57 of the iterates , which may be of independent interest to other related nonconvex optimization58 algorithm design.59 To compare with , we summarized related results of stochastic first-order methods for finding local60 minima in Table 1.61 Notations We use lower case letters to denote scalars , lower and upper case bold letters to denote62 vectors and matrices . We use ‖ · ‖ to indicate Euclidean norm . We use Bx ( r ) to denote a Euclidean63 ball center at x with radius r.We also use the standard O and Ω notations . We use λmin ( M ) to denote64 the minimum eigenvalue of matrix M. We say an = O ( bn ) if and only if ∃C > 0 , N > 0 , ∀n > 65 N , an ≤ Cbn ; an = Ω ( bn ) if an ≥ Cbn . The notation Õ is used to hide logarithmic factors.66 2 Related Work67 In this section , we review some important related works.68 Variance reduction methods for finding stationary points . Our algorithm takes stochastic gradient69 estimators as its subroutine . In specific , Johnson and Zhang [ 14 ] , Xiao and Zhang [ 28 ] proposed70 Stochastic Variance Reduced Gradient ( SVRG ) for convex optimization in the finite-sum setting.71 Reddi et al . [ 25 ] , Allen-Zhu and Hazan [ 2 ] analyzed SVRG for nonconvex optimization . Lei et al . [ 16 ] 72 proposed a new variance reduction algorithm , dubbed stochastically controlled stochastic gradient73 ( SCSG ) algorithm , which finds a -stationary point within O ( −10/3 ) stochastic gradient evaluations.74 Nguyen et al . [ 21 ] proposed a SARAH algorithm which uses a recursive gradient estimator for convex75 optimization , and it was extended to nonconvex optimization in [ 22 ] . Fang et al . [ 9 ] proposed a76 SPIDER algorithm with a recursive gradient estimator and proved an O ( −3 ) stochastic gradient77 evaluations to find a -stationary point , which matches a corresponding lower bound . Concurrently,78 Zhou et al . [ 30 ] proposed an SNVRG algorithm with a nested gradient estimator and proved an Õ ( −3 ) 79 stochastic gradient evaluations to find a -stationary point . Wang et al . [ 27 ] proposed a Spiderboost80 algorithm with a constant step size , achieves the same O ( −3 ) gradient complexity . Pham et al.81 [ 23 ] extended SARAH [ 22 ] to proximal optimization and proved O ( −3 ) gradient complexity for82 finding stationary points . Recently , Cutkosky and Orabona [ 6 ] proposed a recursive momentum-based83 algorithm called STORM and proved an Õ ( −3 ) gradient complexity to find -stationary points.84 Tran-Dinh et al . [ 26 ] proposed a SARAH-SGD algorithm which hybrids both SGD and SARAH85 algorithm with an Õ ( −3 ) gradient complexity when is small . Li et al . [ 18 ] proposed a PAGE86 algorithm with probabilistic gradient estimator which also attains an Õ ( −3 ) gradient complexity.87 In our work , we employ SARAH/SPIDER and STORM as the gradient estimator in our Pullback88 framework since they are most representative and simple to use.89 Utilizing negative curvature descent to escape from saddle points . To escape saddle points , a90 widely used approach is to first compute the direction of the negative curvature of the saddle point and91 move away along that direction . In stochastic optimization , to find ( , H ) -approximate local minima,92 [ 1 ] proposed a Natasha algorithm using Hessian-vector product to compute the negative curvature93 direction with the total computation cost of Õ ( −3.25 + −3 −1H + −5 H ) . Later , Xu et al . [ 29 ] proposed94 a Neon method which computes the negative curvature direction with perturbed stochastic gradients,95 whose total computational cost is Õ ( −10/3 + −2 −3H + −6 H ) . [ 3 ] proposed a Neon2 negative96 curvature computation subroutine with Õ ( −10/3 + −2 −3H + −5 H ) stochastic gradient evaluations.97 Fang et al . [ 9 ] then showed that SPIDER equipped with Neon2 can find ( , H ) -approximate local98 minima within Õ ( −3 + −2 −2H + −5 H ) stochastic gradient evaluations , while independently Zhou99 et al . [ 31 ] proved that SNVRG equipped with Neon2 can find ( , H ) -approximate local minima100 within Õ ( −3 + −2 −3H + −5 H ) stochastic gradient evaluations . In contrast to this line of works , our101 algorithm is simpler since it does not need to use the negative curvature search routine.102 3 Preliminaries103 In this section , we present assumptions and definitions that will be used throughout our analysis.104 We first introduce the standard smoothness and Hessian Lipschitz assumptions.105 Assumption 3.1 . For all ξ , f ( · ; ξ ) is L-smooth and its Hessian is ρ-Lipschitz continuous w.r.t . x,106 i.e. , for any x1 , x2 , we have that107 ‖∇f ( x1 ; ξ ) −∇f ( x2 ; ξ ) ‖2 ≤ L‖x1 − x2‖2 , ‖∇2f ( x1 ; ξ ) −∇2f ( x2 ; ξ ) ‖2 ≤ ρ‖x1 − x2‖2 This assumption directly implies that the expected objective function F ( x ) is also L-smooth and its108 Hessian is ρ-Lipschitz continuous . This assumption is standard for finding approximate local minima109 in all the results presented in Table 1.110 Assumption 3.2 . The squared difference between the stochastic gradient and full gradient is bounded111 by σ2 < ∞ , i.e. , for any x , ξ ∈ Rd , ‖∇f ( x ; ξ ) −∇F ( x ) ‖22 ≤ σ2.112 Assumption 3.2 is standard in online/stochastic optimization for finding second-order stationary113 points [ 9 , 17 ] , and immediately implies that the variance of the stochastic gradient is bounded by σ2.114 It can be relaxed to be ‖∇f ( x ; ξ ) −∇F ( x ) ‖2 has a σ-Sub-Gaussian tail.115 Let x0 ∈ Rd be the starting point of the algorithm . We assume the gap between the initial function116 value and the optimal value is bounded.117 Assumption 3.3 . We have ∆ = F ( x0 ) − infx F ( x ) < +∞.118 Next , we give the formal definition of approximate local minima ( a.k.a. , second-order stationary119 points ) .120 Definition 3.4 . We call x ∈ Rd an ( , H ) -approximate local minimum , if121 ‖∇F ( x ) ‖2 ≤ , λmin ( ∇2F ( x ) ) ≥ − H . The definition of ( , H ) -approximate local minima is a generalization of the classical ( , √ ) -122 approximate local minima studied by Nesterov and Polyak [ 20 ] , Jin et al . [ 13 ] .123 4 The Pullback Framework124 In this section , we present our main algorithm Pullback . We begin with reviewing the mechanism125 of perturbed gradient descent in deterministic optimization , and then we discuss the main difficulty of126 extending it to the stochastic optimization case . Finally , we show how we overcome such a difficulty127 by presenting our Pullback framework.128 How does perturbed gradient descent escape from saddle points ? We review the perturbed129 gradient descent [ 13 ] ( PGD for short ) with its proof roadmap , which shows how PGD finds ( , √ ) -130 approximate local minima efficiently . In general , the whole process of perturbed gradient descent131 can be decomposed into several epochs , and each epoch consists of two non-overlapping phases : the132 gradient descent phase ( GD phase for short ) and the Escape from saddle point phase ( Escape phase133 for short ) . In each epoch , PGD starts with the GD phase by default . In the GD phase , PGD performs134 vanilla gradient descent to update its iterate , until at some iterate x̃ , the norm of the gradient135 ‖∇F ( x̃ ) ‖2 is less than the target accuracy Õ ( ) . Then PGD switches to the Escape phase . In the136 Escape phase , PGD first adds a uniform random noise ( or Gaussian noise ) to the current iterate x̃,137 then it runs ` thres = Õ ( −1/2 ) steps of vanilla gradient descent . PGD then compares the function138 value gap between the current iterate and the beginning iterate of Escape phase x̃ . If the gap is less139 than a threshold F = Õ ( 1.5 ) , then PGD outputs x̃ as the targeted local minimum . Otherwise , PGD140 starts a new epoch and performs gradient descent again.141 To see why PGD can find ( , √ ) -approximate local minima within Õ ( −2 ) gradient evaluations , we142 do the following calculation . First , when PGD is in the GD phase , the function value decreases Õ ( 2 ) 143 per step ( following the standard gradient descent analysis ) . When PGD is in the Escape phase , the144 function value decreases F/ ` thres = Õ ( 2 ) per step on average . Therefore , the total number of steps145 will be bounded by Õ ( −2 ) , which is in the same order as GD for finding -stationary points.146 Limitation of existing methods . However , extending the two-phase PGD algorithm from determinis-147 tic optimization to stochastic optimization with a competative gradient complexity is very challenging.148 We take the SSRGD algorithm proposed by Li [ 17 ] as an example , which uses SARAH/SPIDER [ 9 ] as149 its gradient estimator . Unlike deterministic optimization where we can access the exact function value150 F ( x ) and gradient ∇F ( x ) defined in ( 1.1 ) , in the stochastic optimization case we can only access151 the stochastic function f ( x ; ξ ) and the stochastic gradient∇f ( x ; ξ ) . Therefore , in order to estimate152 the gradient norm ‖∇F ( x ) ‖2 ( which is required at the beginning of Escape phase ) , a naive approach153 ( adapted by Li [ 17 ] ) is to sample a big batch of stochastic gradient∇f ( x ; ξ1 ) , . . . , ∇f ( x ; ξB ) and154 uses their mean to approximate ∇F ( x ) . Standard concentration analysis suggests that in order to155 achieve an -accuracy , the batch size B should be in the order Õ ( −2 ) . Thus , each Escape phase156 leads to a F = Õ ( 1.5 ) function value decrease with at least Õ ( −2 ) number of stochastic gradient157 evaluations , which contributes Õ ( 1/ 1.5 · −2 ) = Õ ( −3.5 ) gradient complexity in the end . This is158 already worse than the O ( −3 ) gradient complexity of SPIDER for finding -stationary points.159 Our approach . Here we propose our Pullback framework in Algorithm 1 , which overcomes the160 aforementioned limitation . In detail , Pullback inherits the two-phase structure of PGD and SSRGD,161 and it takes either SARAH/SPIDER or STORM [ 6 ] as its gradient estimator . The two gradient162 estimators are summarized as subroutines GradEst-SPIDER and GradEst-STORM in Algorithms 2163 and 3 respectively , and we use dt to denote their estimated gradient at iterate xt . The key improvement164 of Pullback is that , it directly takes the output of the gradient estimator GradEst to estimate the165 true gradient ∇F ( x ) , which avoids sampling a big batch of stochastic gradients as in Li [ 17 ] and166 thus saves the total gradient complexity . A similar strategy has also been adapted in [ 9 ] , but for the167 negative curvature search subroutine . However , such a strategy leads to a new problem to be solved.168 Since we use dt to directly estimate ∇F ( xt ) , in order to make such an estimation valid , we need to169 guarantee that the error between dt and ∇F ( xt ) is small enough , e.g. , up to O ( ) accuracy . Notice170 Algorithm 1 Pullback Input : Initial point x1 , step size η and ηH , perturbation radius r , threshold parameter ` thres , average movement D. 1 : d1 ← GradEst ( 0,0,0 , x1 ) , s← 0 , t← 1 , FIND←false 2 : while FIND = false do 3 : s← s+ 1 , ts ← t , FIND←true 4 : while ‖dt‖2 > do 5 : ηt ← η/‖dt‖2 , { `` PullBack '' } 6 : xt+1 ← xt − ηtdt , dt+1 ← GradEst ( t , dt , xt , xt+1 ) , t← t+ 1 7 : end while 8 : ms ← t , ξ ∼ Uniform B0 ( r ) , xt+1 ← xt + ξ , dt+1 ← GradEst ( t , dt , xt , xt+1 ) , t← t+ 1 9 : for k = 0 , . . . , ` thres − 1 do 10 : ηt ← ηH , D ← ∑t i=ms η2i ‖di‖22 11 : if D > ( t−ms + 1 ) D then 12 : Set ηt such that ∑t i=ms η2i ‖di‖22 = ( t−ms + 1 ) D { `` PullBack '' } 13 : xt+1 ← xt − ηtdt , dt+1 ← GradEst ( t , dt , xt , xt+1 ) , t← t+ 1 , FIND← false , break 14 : end if 15 : xt+1 ← xt − ηtdt , dt+1 ← GradEst ( t , dt , xt , xt+1 ) , t← t+ 1 16 : end for 17 : end while Output : xms Algorithm 2 GradEst-SPIDER ( t , dt , xt , xt+1 , b , q , B ) Input : Big batch size B , mini-batch size b , loop length q 1 : if t mod q = 0 then 2 : Generate ξ1t+1 , . . . , ξ B t+1 . Set dt+1 ← ∑B i=1∇f ( xt+1 ; ξit+1 ) /B 3 : else 4 : Generate ξ1t+1 , . . . , ξ b t+1 . Set dt+1 ← dt + ∑b i=1 [ ∇f ( xt+1 ; ξit+1 ) −∇f ( xt ; ξit+1 ) ] /b 5 : end if Output : dt+1 that the recursive structure of SARAH/SPIDER and STORM suggests the following error bound:171 ∀t , ‖dt −∇F ( xt ) ‖22 = Õ ( t−1∑ i=st ‖xi+1 − xi‖22 ) , ( 4.1 ) where st is some reference index only related to t. Therefore , in order to make the error172 ‖dt −∇F ( xt ) ‖2 small , it suffices to make the movement of the iterates ‖xi+1 − xi‖2 small either173 individually or on average . We achieve this goal by our proposed step-size “ Pullback ” scheme . In de-174 tail , in the GD phase , when the norm of the estimated gradient ‖dt‖2 is large , we pull the step-size ηt175 back to a smaller value via normalization , which forces the movement ‖xi+1 − xi‖2 = ηt‖dt‖2 = η176 to be small . Such an approach is also adapted by Fang et al . [ 9 ] as a normalized gradient de-177 scent for finding first-order stationary points . In the Escape phase , which starts at ms-th step , we178 record the accumulative squared movement starting from xms+1 ( after the perturbation step ) as179 D : = ∑t i=ms+1 ‖xi+1−xi‖22 = ∑t i=ms+1 η2i ‖di‖22 . When the average movement D/ ( t−ms + 1 ) 180 is large , we pull the last step size ηt back to a smaller value , which forces the average movement181 D/ ( t−ms + 1 ) to be small . Fortunately , such a simple step-size calibration scheme allows us to182 well-control the error between dt and∇F ( xt ) , and to reduce the gradient complexity.183 5 Main Results184 In this section , we present the main theoretical results . We first present the convergence guarantee of185 Pullback-SPIDER , which uses GradEst-SPIDER to estimate the gradient dt in Algorithm 1.186 Theorem 5.1 . Under Assumptions 3.1 , 3.2 and 3.3 , choose batch size B = Õ ( σ2 −2 +187 σ2ρ2 −4H ) , b = q = √ B , set step size η = σ/ ( 2 √ BL ) , ηH = Õ ( L−1 ) , perturbation ra-188 dius r ≤ min { σ/ ( 2 √ BL ) , log ( 4/δ ) ηHσ 2/ ( 2B ) , √ 2 log ( 4/δ ) ηHσ2/ ( BL ) } , threshold ` thres =189 Algorithm 3 GradEst-STORM ( t , dt , xt , xt+1 , a , b , B ) Input : Initial batch size B , mini batch size b and weight parameter a . 1 : if t = 0 then 2 : Generate ξ1t+1 , . . . , ξ B t+1 . Set dt+1 ← ∑B i=1∇f ( xt+1 ; ξit+1 ) /B 3 : else 4 : Generate ξ1t+1 , . . . , ξ b t+1 5 : Set dt+1 ← ( 1− a ) [ dt − ∑b i=1∇f ( xt ; ξit+1 ) /b ] + ∑b i=1∇f ( xt+1 ; ξit+1 ) /b 6 : end if Output : dt+1 Õ ( 1/ ( ηH H ) ) and D = σ2/ ( 4BL2 ) . Then with high probability , Pullback-SPIDER can find190 ( , H ) -approximate local minima within Õ ( σL∆ −3 + σρ3L∆ −6H ) stochastic gradient evaluations.191 Remark 5.2 . In the classical setting = √ H , our result gives Õ ( −3 ) gradient complexity , which192 outperforms the best existing result Õ ( −3.5 ) for perturbed stochastic gradient methods achieved193 by SSRGD [ 17 ] . For sufficiently small , Arjevani et al . [ 4 ] proved the lower bound of gradient194 complexity Ω ( −3 + −5H ) for any first-order stochastic methods to find ( , H ) -approximate local195 minima . Our results matches the lower bound Õ ( −3 ) when H ≤ 3/5 . For the general case , there is196 still a gap in the dependency of H between our result and the lower bound , and we leave to close it197 as future work.198 Next , we present the convergence guarantee of Pullback-STORM , which uses GradEst-STORM to199 estimate the gradient dt in Algorithm 1.200 Theorem 5.3 . Under Assumptions 3.1 , 3.2 and 3.3 , choose the mini batch size b = Õ ( σ −1 +201 σρ −2H ) , and initial batch size B = b2 , set step size η = σ/ ( 2bL ) , ηH = Õ ( L−1 ) ,202 weight a = 562 log ( 4/δ ) /b , threshold ` thres = Õ ( 1/ ( ηH H ) ) , perturbation radius r ≤203 min { σ/ ( 2bL ) , log ( 4/δ ) 2ηHσ 2/ ( b2 ) , √ 2 log ( 4/δ ) 2ηHσ2/ ( b2L ) } , and D = σ2/ ( 4b2L2 ) . Then204 with high probability , Pullback-STORM can find ( , H ) -approximate local minima within205 Õ ( σL∆ −3 + σρ3L∆ −6H ) stochastic gradient evaluations.206 Remark 5.4 . Different from Pullback-SPIDER , the estimation error ‖dt − ∇F ( xt ) ‖2 of207 Pullback-STORM is controlled by the weight parameter a . This allows us to come up with a208 simpler single-loop algorithm instead of a double-loop algorithm.209 6 Proof Outline of the Main Results210 Due to the page limit , we only outline the proof of Theorem 5.1 and leave the proof of Theorem 5.3211 to the appendix.212 Let t denote the difference between true gradient∇F ( xt ) and the estimated gradient dt , which is213 t : = dt −∇F ( xt ) . The following lemma suggests that the estimation error ‖ t‖2 can be bounded.214 Lemma 6.1 . Under Assumptions 3.1 and 3.2 , set b = q = √ B , η ≤ σ/ ( 2 √ BL ) , r ≤ σ/ ( 2 √ BL ) 215 and D ≤ σ2/ ( 4BL2 ) , then with probability at least 1− δ , for all t we have216 ‖ t‖2 ≤ √ 8 log ( 4/δ ) σ/ √ B . Specifically , by the choice of B in Theorem 5.1 we have that ‖ t‖2 ≤ /2.217 Proof of Lemma 6.1 . By GradEst-SPIDER presented in Algorithm 2 we have218 t+1 = 1 B B∑ i=1 [ ∇f ( xt+1 ; ξit+1 ) −∇F ( xt+1 ) ] , t mod q = 0 , t+1 = t + 1 b b∑ i=1 [ ∇f ( xt+1 ; ξit+1 ) −∇f ( xt ; ξit+1 ) −∇F ( xt+1 ) +∇F ( xt ) ] , t mod q 6= 0 . By the L-smoothness in Assumption 3.1 we have219 ∥∥∇f ( xt+1 ; ξit+1 ) −∇f ( xt ; ξit+1 ) −∇F ( xt+1 ) +∇F ( xt ) ∥∥2 ≤ 2L‖xt+1 − xt‖2 . Then by Assumption 3.2 and Azuma–Hoeffding inequality ( See Lemma D.1 for details ) , with220 probability at least 1− δ , we have221 ∀t > 0 , ‖ t+1‖22 ≤ 4 log ( 4/δ ) ( σ2 B + 4L2 b t∑ i=bt/qcq ‖xi+1 − xi‖22 ) . ( 6.1 ) Notice that GradEst-SPIDER is parallel with Pullback . Thus we need to further bound ( 6.1 ) by222 considering iterates in three different cases : ( 1 ) for step i in the GD phase , we have ‖xi+1−xi‖22 ≤ η2223 due to the `` Pullback '' scheme ; ( 2 ) for i = ms for some s in the Escape phase , we have ‖xi+1−xi‖22 ≤224 r2 ; and ( 3 ) for the other steps in Escape phase , we have on average , ‖xi+1 − xi‖22 ≤ D. Therefore225 we have226 ‖ t+1‖22 ≤ 4 log ( 4/δ ) ( σ2 B + 4L2 b · q ·max { η2 , r2 , D } ) ≤ 8 log ( 4/δ ) σ 2 B . 227 Lemma 6.1 guarantees that with high probability ‖∇F ( xt ) ‖2 ≤ ‖dt‖2 + , which ensures228 ‖∇F ( xms ) ‖2 ≤ 2 when the algorithm terminates . Next lemma bounds the function value de-229 crease in the GD phase , which is also valid for Pullback-STORM.230 Lemma 6.2 . Suppose the event in Lemma 6.1 holds , η ≤ / ( 2L ) , then for any s , we have231 F ( xts ) − F ( xms ) ≥ ( ms − ts ) η 8 . The choice of η in Theorem 5.1 further implies that the loss decreases by at least σ / ( 16 √ BL ) per232 step on average.233 Proof of Lemma 6.2 . For any ts ≤ t < ms , we can show the following property ( See Lemma D.2 ) ,234 F ( xt+1 ) ≤ F ( xt ) − ηt 2 ‖dt‖22 + ηt 2 ‖ t‖22 + L 2 ‖xt+1 − xt‖22 . ( 6.2 ) Plugging the update rule xt+1 = xt − ηtdt into ( 6.2 ) gives,235 F ( xt+1 ) = F ( xt ) − ‖xt+1 − xt‖22 ( 1 2ηt − L 2 ) + ηt‖ t‖22 2 ≤ F ( xt ) − η2 ( 1 2ηt − L 2 ) + ηt 2 8 , ≤ F ( xt ) − η 8 where the first inequality holds due to the fact that ηt = η/‖dt‖2 and ‖ t‖2 ≤ /2 , and the second236 inequality is by ηt = η/‖dt‖2 ≤ η/ ≤ 1/ ( 2L ) .237 Following Lemma shows that if xms is a saddle point , then with high probability , the algorithm will238 break during the Escape phase and set FIND←false . Thus , whenever xms is not a local minima , the239 algorithm can not terminate.240 Lemma 6.3 . Under Assumptions 3.1 and 3.2 , set perturbation radius r ≤ LηH H/ ( Cρ ) , step241 size ηH ≤ min { 1/ ( 16L log ( ηH H √ dLC−1ρ−1δ−1r−1 ) ) , 1/ ( 8CL log ` thres ) } = Õ ( L−1 ) , ` thres =242 2 log ( ηH H √ dLC−1ρ−1δ−1r−1 ) / ( ηH H ) = Õ ( η −1 H −1 H ) , and D < C 2L2η2H 2 H/ ( ρ 2 ` 2thres ) , where243 C = O ( log ( d ` thres/δ ) = Õ ( 1 ) . We also set b = q = √ B ≥ 16 log ( 4/δ ) / ( η2H 2H ) . Then for244 any s , when λmin ( ∇2F ( xms ) ) ≤ − H , with probability at least 1 − 2δ algorithm breaks in the245 Escape phase.246 Proof of Lemma 6.3 . Let { xt } , { x′t } be two coupled sequences by running Pullback-SPIDER from247 xms+1 , x ′ ms+1 with xms+1 − x ′ ms+1 = r0e1 , where xms+1 , x ′ ms+1 ∈ Bxms ( r ) . Here r0 = δr/ √ d248 and e1 denotes the smallest eigenvector direction of Hessian∇2F ( xms ) .249 When λmin ( ∇2F ( xms ) ) ≤ − H , under the parameter choice in Lemma 6.3 , we can show that250 at least one of two sequence will escape the saddle point ( See Lemma D.3 ) . To be specific , with251 probability at least 1− δ,252 max ms < t < ms+ ` thres { ‖xt − xms+1‖2 , ‖x′t − x′ms+1‖2 } ≥ LηH H Cρ . ( 6.3 ) ( 6.3 ) suggests that for any two points xms+1 , x ′ ms+1 satisfying xms+1−x ′ ms+1 = r0e1 , at least one253 of them will generate a sequence of iterates which finally move more than LηH H/ ( Cρ ) . Thus , let254 S ⊆ Bms ( r ) be the set of xms+1 which will not generate a sequence of iterates moving more than255 LηH H Cρ , then in the direction e1 , the `` thickness '' of S is smaller than r0 . Simple integration shows that256 the ratio between the volume of S and Bms ( r ) is bounded by δ . Therefore , since xms+1 is generated257 from xms by adding a uniform random noise in ball Bms ( r ) , we conclude that the probability for258 xms+1 locating in S is less than δ . Applying union bound , we get with probability at least 1− 2δ,259 ∃ms < t < ms + ` thres , ‖xt − xms+1‖2 ≥ LηH H Cρ . ( 6.4 ) Denote E as the event that the algorithm does not break in the Escape phase . Then under E , for any260 ms < t < ms + ` thres , we have261 ‖xt − xms+1‖2 ≤ t−1∑ i=ms+1 ‖xi+1 − xi‖2 ≤ √√√√ ( t−ms ) t−1∑ i=ms ‖xi+1 − xi‖22 ≤ ( t−ms ) √ D , where the first inequality is due to the triangle inequality and the second inequality is due to Cauchy-262 Schwarz inequality . Thus , by the choice of ` thres and D , we have263 ‖xt − xms+1‖2 ≤ ( t−ms ) √ D ≤ ` thres √ D < C · LηH H ρ . Then by ( 6.4 ) , we know that P ( E ) ≤ 2δ . Therefore when λmin ( ∇2F ( xms ) ) ≤ − H , with probability264 at least 1− 2δ , Pullback breaks in the Escape phase.265 Next lemma bounds the decreasing value of the function during the Escape phase if the algorithm266 breaks in the Escape phase ( i.e . FIND is false ) .267 Lemma 6.4 ( localization ) . Suppose the result of Lemma 6.1 holds , and268 set the step size ηH ≤ 1/ ( L √ 128 log ( 4/δ ) ) , perturbation radius r ≤269 min { log ( 4/δ ) ηHσ 2/ ( 2B ) , √ 2 log ( 4/δ ) ηHσ2/ ( BL ) } , and D = σ2/ ( 4BL2 ) . Suppose the270 algorithm breaks in the Escape phase starting at xms , then we have271 F ( xms ) − F ( xts+1 ) ≥ ( ts+1 −ms ) log ( 4/δ ) ηHσ 2 B . Proof of Lemma 6.4 . For anyms < i < ts+1 , we can show the following property ( See Lemma D.2 ) ,272 F ( xi+1 ) ≤ F ( xi ) − ηi 2 ‖di‖22 + ηi 2 ‖ i‖22 + L 2 ‖xi+1 − xi‖22 . ( 6.5 ) Plugging the update rule xi+1 = xi − ηidi into ( 6.5 ) gives,273 F ( xi+1 ) ≤ F ( xi ) + ηi 2 ‖ i‖22 − ( 1 2ηi − L 2 ) ‖xi+1 − xi‖22 ≤ F ( xi ) + ηH 2 8 log ( 4/δ ) σ2 B − 1 4ηH ‖xi+1 − xi‖22 ( 6.6 ) where the the second inequality holds due to Lemma 6.1 and ηi ≤ ηH ≤ 1/ ( 2L ) for any ms < i < 274 ts+1 . Telescoping ( 6.6 ) from i = ms + 1 to ts+1 − 1 , we have275 F ( xts+1 ) ≤ F ( xms+1 ) + 4ηH log ( 4/δ ) ( ts+1 −ms − 1 ) σ2 B − 1 4ηH ts+1−1∑ i=ms+1 ‖xi+1 − xi‖22 . Finally , we have276 F ( xms+1 ) − F ( xts+1 ) ≥ ts+1−1∑ i=ms+1 ‖xi+1 − xi‖22 4ηH − 4 log ( 4/δ ) ( ts+1 −ms − 1 ) ηH σ2 B = ( ts+1 −ms − 1 ) ( D 4ηH − 4 log ( 4/δ ) ηHσ 2 B ) = ( ts+1 −ms − 1 ) ( σ2 16ηHBL2 − 4 log ( 4/δ ) ηHσ 2 B ) ≥ ( ts+1 −ms − 1 ) 4 log ( 4/δ ) ηHσ 2 B , ( 6.7 ) where the last inequality is by the choice of ηH ≤ 1/ ( L √ 128 log ( 4/δ ) ) . For i = ms , we have ( See277 Lemma D.2 ) 278 F ( xms+1 ) ≤ F ( xms ) + ( ‖dms‖2 + ‖ ms‖2 + Lr/2 ) r. ( 6.8 ) Plugging ‖dms‖2 ≤ and ‖ ms‖2 ≤ /2 into ( 6.8 ) gives,279 F ( xms+1 ) ≤ F ( xms ) + ( 4 + Lr/2 ) r ≤ F ( xms ) + 2 log ( 4/δ ) ηHσ 2 B , ( 6.9 ) where the last inequality is by the choice r ≤ min { log ( 4/δ ) ηHσ 2/ ( 2B ) , √ 2 log ( 4/δ ) ηHσ2/ ( BL ) } .280 Combining ( 6.7 ) and ( 6.9 ) and applying ts+1 −ms ≥ 2 gives,281 F ( xms ) − F ( xts+1 ) ≥ [ 4 ( ts+1 −ms − 1 ) − 2 ] log ( 4/δ ) ηHσ 2 B ≥ ( ts+1 −ms ) log ( 4/δ ) ηHσ 2 B . 282 Now , we can provide the proof of Theorem 5.1 .283 Proof of Theorem 5.1 . The analysis can be divided into two phases , i.e. , GD phase and Escape phase.284 The function value will decrease at different rates in different phases.285 GD phase : In this phase , ‖dt‖2 ≥ and ‖ ‖2 ≤ /2 due to Lemma 6.1 . Thus the gradients of the286 function are large ‖∇F ( x ) ‖2 ≥ /2 . Lemma 6.2 further shows that the loss decreases by at least287 σ / ( 16 √ BL ) on average.288 Escape phase : In this phase , the starting point xms satisfies ‖∇F ( xms ) ‖2 ≤ ‖dms‖2 + ‖ t‖2 ≤ 2 .289 If xms is a saddle point with λmin ( ∇2F ( xms ) ) ≤ − H , then by Lemma 6.3 , with high probability290 Pullback-SPIDER will break Escape phase , set FIND←False and begin a new GD phase . Further291 by Lemma 6.4 , the loss will decrease by at least log ( 4/δ ) ηHσ2/B per step on average.292 Sample Complexity : Note that the total amount for function value can decrease is at most ∆ =293 F ( x0 ) − infx F ( x ) < +∞ . So the algorithm must end and find an ( , H ) -approximate local294 minimum within Õ ( √ BL∆σ−1 −1 + BL∆σ−2 ) iterations . Notice that on average we sample295 max { b , B/q } = √ B examples per iteration , so the total sample complexity is Õ ( BL∆σ−1 −1 +296 B3/2L∆σ−2 ) . Plugging in the choice of B = Õ ( σ2 −2 + σ2ρ2 −4H ) in Theorem 5.1 , we have the297 total gradient complexity298 Õ ( σL∆ 3 + σρ2L∆ 4H + σρ3L∆ 6H ) = Õ ( σL∆ 3 + σρ3L∆ 6H ) , where the equation is due to the Young ’ s inequality.299 300 7 Conclusions301 In this paper , we propose a perturbed stochastic gradient framework named Pullback for finding302 local minima . Pullback can find ( , H ) -approximate local minima within Õ ( −3 + −6H ) stochastic303 gradient evaluations , which matches the best possible complexity results in the classical H = √ 304 setting . Our results show that simple perturbed gradient methods can be as efficient as more305 sophisticated algorithms for finding local minima.306 References307 [ 1 ] Allen-Zhu , Z . ( 2018 ) . Natasha 2 : Faster non-convex optimization than sgd . In Advances in308 neural information processing systems , pages 2675–2686.309 [ 2 ] Allen-Zhu , Z. and Hazan , E. ( 2016 ) . Variance reduction for faster non-convex optimization . In310 International Conference on Machine Learning , pages 699–707.311 [ 3 ] Allen-Zhu , Z. and Li , Y . ( 2018 ) . Neon2 : Finding local minima via first-order oracles . In Advances312 in Neural Information Processing Systems , pages 3716–3726.313 [ 4 ] Arjevani , Y. , Carmon , Y. , Duchi , J. C. , Foster , D. J. , Sekhari , A. , and Sridharan , K. ( 2020 ) .314 Second-order information in non-convex stochastic optimization : Power and limitations . In315 Conference on Learning Theory , pages 242–299 . PMLR.316 [ 5 ] Bhojanapalli , S. , Neyshabur , B. , and Srebro , N. ( 2016 ) . Global optimality of local search for low317 rank matrix recovery . arXiv preprint arXiv:1605.07221.318 [ 6 ] Cutkosky , A. and Orabona , F. ( 2019 ) . Momentum-based variance reduction in non-convex sgd.319 In Advances in Neural Information Processing Systems , pages 15210–15219.320 [ 7 ] Daneshmand , H. , Kohler , J. , Lucchi , A. , and Hofmann , T. ( 2018 ) . Escaping saddles with321 stochastic gradients . arXiv preprint arXiv:1803.05999.322 [ 8 ] Du , S. S. , Jin , C. , Lee , J. D. , Jordan , M. I. , Poczos , B. , and Singh , A . ( 2017 ) . Gradient descent323 can take exponential time to escape saddle points . arXiv preprint arXiv:1705.10412.324 [ 9 ] Fang , C. , Li , C. J. , Lin , Z. , and Zhang , T. ( 2018 ) . Spider : Near-optimal non-convex optimization325 via stochastic path-integrated differential estimator . In Advances in Neural Information Processing326 Systems , pages 689–699.327 [ 10 ] Ge , R. , Huang , F. , Jin , C. , and Yuan , Y . ( 2015 ) . Escaping from saddle points—online stochastic328 gradient for tensor decomposition . In Conference on Learning Theory , pages 797–842.329 [ 11 ] Ge , R. , Jin , C. , and Zheng , Y . ( 2017 ) . No spurious local minima in nonconvex low rank330 problems : A unified geometric analysis . In International Conference on Machine Learning , pages331 1233–1242 . PMLR.332 [ 12 ] Hillar , C. J. and Lim , L.-H. ( 2013 ) . Most tensor problems are np-hard . Journal of the ACM333 ( JACM ) , 60 ( 6 ) :1–39.334 [ 13 ] Jin , C. , Ge , R. , Netrapalli , P. , Kakade , S. M. , and Jordan , M. I . ( 2017 ) . How to escape saddle335 points efficiently . In International Conference on Machine Learning , pages 1724–1732 . PMLR.336 [ 14 ] Johnson , R. and Zhang , T. ( 2013 ) . Accelerating stochastic gradient descent using predictive337 variance reduction . In Advances in Neural Information Processing Systems , pages 315–323.338 [ 15 ] Kallenberg , O. and Sztencel , R. ( 1991 ) . Some dimension-free features of vector-valued martin-339 gales . Probability Theory and Related Fields , 88 ( 2 ) :215–247.340 [ 16 ] Lei , L. , Ju , C. , Chen , J. , and Jordan , M. I . ( 2017 ) . Non-convex finite-sum optimization via scsg341 methods . In Advances in Neural Information Processing Systems , pages 2348–2358.342 [ 17 ] Li , Z . ( 2019 ) . Ssrgd : Simple stochastic recursive gradient descent for escaping saddle points . In343 Advances in Neural Information Processing Systems , pages 1521–1531.344 [ 18 ] Li , Z. , Bao , H. , Zhang , X. , and Richtárik , P. ( 2020 ) . Page : A simple and optimal probabilistic345 gradient estimator for nonconvex optimization . arXiv preprint arXiv:2008.10898.346 [ 19 ] Nesterov , Y . ( 1998 ) . Introductory lectures on convex programming volume i : Basic course.347 Lecture notes , 3 ( 4 ) :5.348 [ 20 ] Nesterov , Y. and Polyak , B. T. ( 2006 ) . Cubic regularization of newton method and its global349 performance . Mathematical Programming , 108 ( 1 ) :177–205.350 [ 21 ] Nguyen , L. M. , Liu , J. , Scheinberg , K. , and Takáč , M. ( 2017a ) . Sarah : A novel method for351 machine learning problems using stochastic recursive gradient . In International Conference on352 Machine Learning , pages 2613–2621 . PMLR.353 [ 22 ] Nguyen , L. M. , Liu , J. , Scheinberg , K. , and Takáč , M. ( 2017b ) . Stochastic recursive gradient354 algorithm for nonconvex optimization . arXiv preprint arXiv:1705.07261.355 [ 23 ] Pham , N. H. , Nguyen , L. M. , Phan , D. T. , and Tran-Dinh , Q . ( 2020 ) . Proxsarah : An efficient356 algorithmic framework for stochastic composite nonconvex optimization . Journal of Machine357 Learning Research , 21 ( 110 ) :1–48.358 [ 24 ] Pinelis , I . ( 1994 ) . Optimum bounds for the distributions of martingales in banach spaces . The359 Annals of Probability , pages 1679–1706.360 [ 25 ] Reddi , S. J. , Hefny , A. , Sra , S. , Poczos , B. , and Smola , A . ( 2016 ) . Stochastic variance reduction361 for nonconvex optimization . In International Conference on Machine Learning , pages 314–323.362 [ 26 ] Tran-Dinh , Q. , Pham , N. H. , Phan , D. T. , and Nguyen , L. M. ( 2019 ) . Hybrid stochastic gradient363 descent algorithms for stochastic nonconvex optimization . arXiv preprint arXiv:1905.05920.364 [ 27 ] Wang , Z. , Ji , K. , Zhou , Y. , Liang , Y. , and Tarokh , V. ( 2019 ) . Spiderboost and momentum : Faster365 variance reduction algorithms . Advances in Neural Information Processing Systems , 32:2406–366 2416.367 [ 28 ] Xiao , L. and Zhang , T. ( 2014 ) . A proximal stochastic gradient method with progressive variance368 reduction . SIAM Journal on Optimization , 24 ( 4 ) :2057–2075.369 [ 29 ] Xu , Y. , Jin , R. , and Yang , T. ( 2017 ) . First-order stochastic algorithms for escaping from saddle370 points in almost linear time . arXiv preprint arXiv:1711.01944.371 [ 30 ] Zhou , D. , Xu , P. , and Gu , Q . ( 2018 ) . Stochastic nested variance reduction for nonconvex372 optimization . Advances in Neural Information Processing Systems , 31:3921–3932.373 [ 31 ] Zhou , D. , Xu , P. , and Gu , Q . ( 2020 ) . Stochastic nested variance reduction for nonconvex374 optimization . Journal of machine learning research.375 Checklist376 1 . For all authors ... 377 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s378 contributions and scope ? [ Yes ] 379 ( b ) Did you describe the limitations of your work ? [ Yes ] We note in Remark 5.2 that there380 is a gap in the dependency of H between our result and the lower bound . We will381 explore this in future work.382 ( c ) Did you discuss any potential negative societal impacts of your work ? [ N/A ] Our383 work studies the theoretical aspect of optimization algorithms , thus it does not have384 any negative social impact.385 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to386 them ? [ Yes ] 387 2 . If you are including theoretical results ... 388 ( a ) Did you state the full set of assumptions of all theoretical results ? [ Yes ] 389 ( b ) Did you include complete proofs of all theoretical results ? [ Yes ] 390 3 . If you ran experiments ... 391 ( a ) Did you include the code , data , and instructions needed to reproduce the main experi-392 mental results ( either in the supplemental material or as a URL ) ? [ N/A ] 393 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they394 were chosen ) ? [ N/A ] 395 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running experi-396 ments multiple times ) ? [ N/A ] 397 ( d ) Did you include the total amount of compute and the type of resources used ( e.g. , type398 of GPUs , internal cluster , or cloud provider ) ? [ N/A ] 399 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 400 ( a ) If your work uses existing assets , did you cite the creators ? [ N/A ] 401 ( b ) Did you mention the license of the assets ? [ N/A ] 402 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ N/A ] 403 404 ( d ) Did you discuss whether and how consent was obtained from people whose data you ’ re405 using/curating ? [ N/A ] 406 ( e ) Did you discuss whether the data you are using/curating contains personally identifiable407 information or offensive content ? [ N/A ] 408 5 . If you used crowdsourcing or conducted research with human subjects ... 409 ( a ) Did you include the full text of instructions given to participants and screenshots , if410 applicable ? [ N/A ] 411 ( b ) Did you describe any potential participant risks , with links to Institutional Review412 Board ( IRB ) approvals , if applicable ? [ N/A ] 413 ( c ) Did you include the estimated hourly wage paid to participants and the total amount414 spent on participant compensation ? [ N/A ] 415

Summary: The paper proposes a new approach called "Pullback" for escaping saddle points and finding local minimum that matches the same complexity as SGD. The approach perturbs a stochastic gradient estimators (SARAH/SPIDER) and STORM by an element uniformly on the ball of radius $r$. The proposed method is able to find an $(\varepsilon, \sqrt{\varepsilon})$-approximate local minima within $\mathcal{O}(\varepsilon^{-3.5})$ stochastic gradient evaluations. This is the fastest known algorithm achieving the optimal rate of $\mathcal{O}(\varepsilon^{-3})$. The proof techniques appear to be similar to previous works and the gradient estimator used to achieve this result is known.