Faster Perturbed Stochastic Gradient Methods for Finding Local Minima

Abstract:
Escaping from saddle points and finding local minimum is a central problem in 1 nonconvex optimization. Perturbed gradient methods are perhaps the simplest 2 approach for this problem. However, to find ( , √ )-approximate local minima, the 3 existing best stochastic gradient complexity for this type of algorithms is Õ( −3.5), 4 which is not optimal. In this paper, we propose Pullback, a faster perturbed 5 stochastic gradient framework for finding local minima. We show that Pullback 6 with stochastic gradient estimators such as SARAH/SPIDER and STORM can 7 find ( , H)-approximate local minima within Õ( −3 + −6 H ) stochastic gradient 8 evaluations (or Õ( −3) when H = √ ). The core idea of our framework is a 9 step-size “pullback” scheme to control the average movement of the iterates, which 10 leads to faster convergence to the local minima. 11

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 cannot 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

Review summaries:
This papers looks into stochastic gradient methods with the goal of avoiding saddle points in nonconvex optimization. The main contribution of this work is a *perturbed* stochastic gradient framework (called "Pullback") which, by using a particular normalization scheme, finds an $(\epsilon, \epsilon_H)$-approximate local minimum (that is, an $x$ s.t. $||\nabla F(x)|| \leq \epsilon$ and $\lambda_{\textrm{min}}(\nabla^2 F(x)) \geq -\epsilon_H$) in $\tilde{O}(\epsilon^{-3} + \epsilon_H^{-6})$ stochastic gradient evaluations. When $\epsilon_H = \sqrt{\epsilon}$ (referred to as the "classic setting" in the paper), their method results in a total of $\tilde{O}(\epsilon^{-3})$ stochastic gradient evaluations, which improves upon all previous perturbed methods. This rate also matches (in the classic setting) the rate previously attained by SPIDER-SFO+ (+Neon2) (though this previous method requires an additional negative curvature search routine).

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.