STORM: Fully Adaptive SGD with Momentum for Nonconvex Optimization
Year: 2022

Abstract:
In this work we investigate stochastic non-convex optimization problems where the objective is an expectation over smooth loss functions, and the goal is to find an approximate stationary point. The most popular approach to handling such problems is variance reduction techniques, which are also known to obtain tight convergence rates, matching the lower bounds in this case. Nevertheless, these techniques require a careful maintenance of anchor points in conjunction with appropriately selected “mega-batchsizes". This leads to a challenging hyperparameter tuning problem, that weakens their practicality. Recently, [Cutkosky and Orabona, 2019] have shown that one can employ recursive momentum in order to avoid the use of anchor points and large batchsizes, and still obtain the optimal rate for this setting. Yet, their method called STORM crucially relies on the knowledge of the smoothness, as well a bound on the gradient norms. In this work we propose STORM, a new method that is completely parameter-free, does not require large batch-sizes, and obtains the optimal O(1/T ) rate for finding an approximate stationary point. Our work builds on the STORM algorithm, in conjunction with a novel approach to adaptively set the learning rate and momentum parameters.

heading: 1 Introduction
text: Over the past decade non-convex models have become principal tools in ML (Machine Learning), and in data-science. This predominantly includes deep models, as well as Phase Retrieval [Candes et al., 2015], non-negative matrix factorization [Hoyer, 2004], and matrix completion problems [Ge et al., 2016] among others.
The main workhorse for training ML models is SGD (stochastic gradient descent) and its numerous variants. One parameter that significantly affects the SGD performance is the learning rate, which often requires a careful and costly hyper-parameter tuning. Adaptive approaches to setting the learning rate like AdaGrad [Duchi et al., 2011] and Adam [Kingma and Ba, 2014] as well as non-adaptive heuristics [Loshchilov and Hutter, 2017, He et al., 2019] are very popular in modern ML applications, yet these methods also require some tuning of hyper-parameters like momentum and the scale of the learning rate schedule.
A popular SGD heuristic that has proven to be crucial in many applications is the use of momentum, i.e., the use of a weighted average of past gradients instead of the current gradient [Sutskever et al., 2013, Kingma and Ba, 2014]. Although adaptive approaches to setting the momentum have been investigated in the past [Srinivasan et al., 2018, Hameed et al., 2016], principled and theoreticallygrounded approaches to doing so are less investigated. Another aspect that has not been extensively studied, which we take into account in this work, is the interplay between learning rate and momentum.
∗A Viterbi fellow. Corresponding author
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
In this work we explore momentum-based adaptive and parameter-free methods for stochastic non-convex optimization problems. Concretely, we focus on the setting where the objective is an expectation over smooth losses (see Eq. (4)), and the goal is to find an approximate stationary point.
In the general case of smooth non-convex objectives it is known that one can approach a stationary point at a rate of O(1/T 1/4), where T is the total number of samples [Ghadimi and Lan, 2013]. While this rate is optimal in the general case, it is known that one can obtain an improved rate of O(1/T 1/3) if the objective is an expectation over smooth losses [Fang et al., 2018, Zhou et al., 2018, Cutkosky and Orabona, 2019, Tran-Dinh et al., 2019]. Besides, this rate was recently shown to be tight [Arjevani et al., 2019].
Nevertheless, most of the methods developed for this setting rely on variance reduction techniques [Johnson and Zhang, 2013, Zhang et al., 2013, Mahdavi et al., 2013, Wang et al., 2013], which require careful maintenance of anchor points in conjunction with appropriately selected large batchsizes. This leads to a challenging hyper-parameter tuning problem, weakening their practicality. One exception is the recent STORM algorithm of Cutkosky and Orabona [2019].
STORM does not require large batches nor anchor points; instead, it uses a corrected momentum based gradient update that leads to implicit variance reduction, which in turn facilitates fast convergence. Unfortunately, none of the aforementioned methods (including STORM ) is parameter-free. Indeed, the knowledge of smoothness parameter together with either the noise variance or a bound on the norm of the gradients are crucial to establish their guarantees.
In this work, we essentially develop a parameter-free variant of STORM algorithm. We summarize our contributions as follows,
• We present STORM+ , a parameter-free momentum based method that ensures the optimal O(1/T 1/3) rate for the expectation over smooth losses setting. Similarly to STORM , our method does not require large-batches nor anchor points.
• STORM+ implicitly adapts to the variance of the gradients. Concretely, it obtains convergence rate of O(1/ √ T + σ1/3/T 1/3), which recovers the optimal O(1/ √ T ) rate in the
noiseless case. We also improve over STORM by shaving off a (log T )3/4 factor from the 1/ √ T term.
• In STORM+ we demonstrate a novel way to set the learning rate by introducing an adaptive interplay between learning rate and momentum parameters.

heading: 2 Related Work
text: In the context of stochastic non-convex optimization with general smooth losses, it was shown in Ghadimi and Lan [2013] that SGD with an appropriately selected learning rate can obtain a rate of O(1/T 1/4) for finding an approximate stationary point, which is known to match the respective lower bound [Arjevani et al., 2019]. While the method of Ghadimi and Lan [2013] requires knowledge of the smoothness and variance parameters, recent works have shown that adaptive methods like AdaGrad are able to obtain this bound in a parameter free manner, as well as to adapt to the variance of the problem [Li and Orabona, 2019, Ward et al., 2019, Reddi et al., 2018]. These results, in a sense, explain the success of adaptive2 methods like AdaGrad [Duchi et al., 2011], Adam [Kingma and Ba, 2014], and RMSProp [Tieleman and Hinton, 2012] in handling non-convex problems.
The idea of using variance reduction techniques for non-convex problems was first suggested in the context of finite sum problems by Allen-Zhu and Hazan [2016], Reddi et al. [2016], showing a rate of O(1/T 1/4). This was later improved by Lei et al. [2017] to a rate of O(1/T 3/10). The first works that have obtained the optimal O(1/T 1/3) for this setting were Fang et al. [2018], Zhou et al. [2018]. Additionally, Fang et al. [2018] shows that the same convergence behavior applies to the more general expectation over smooth losses setting (see Eq. (4)) – a setting that captures finite-sum problems as a private case.
2An adaptive method is a method that updates its learning rate according to the (noisy) gradient feedback that it receives throughout the training process.
The STORM algorithm suggested in Cutkosky and Orabona [2019] is the first algorithm to obtain the optimal O(1/T 1/3) for this setting without the need to maintain anchor points or large batches. Instead, it relies on a clever correction of the momentum by making only one extra call to the oracle, which leads to an implicit variance reduction effect. Moreover, STORM adapts to the variance of the problem by obtaining a rate of O((log T )3/4/ √ T + σ1/3/T 1/3) without any prior knowledge of variance. However, it needs to know the smoothness parameter and a bound on the gradient norms to set the step size and momentum parameters. Simultaneously to the work of Cutkosky and Orabona [2019], another paper [Tran-Dinh et al., 2019] have obtained the same optimal bound by proposing a similar update rule. Note that Tran-Dinh et al. [2019] does calculate a single anchor point, and it still requires the knowledge of the smoothness and variance parameters.

heading: 3 Setting and Preliminaries
text: We discuss stochastic non-convex optimization problems where the objective f : Rd 7→ R is of the following form, f(x) := Eξ∼D[f(x; ξ)] , and D is an unknown distribution from which we may draw i.i.d. samples. Our goal is to find an approximate stationary point of f , i.e. after T draws from D we should output a point x̄ ∈ Rd such that E‖∇f(x)‖ ≤ Poly(1/T ). We focus on first order methods, i.e., methods that may access the gradients of f(·, ξ), and make the following assumptions regarding the noisy gradients and function values.
Bounded values: There exists B > 0 such that,
max x,y∈Rd
|f(x)− f(y)| ≤ B. (1)
Bounded gradients: There exists G > 0 such that,
‖∇f(x; ξ)‖2 ≤ G2 ; ∀x ∈ Rd, ξ ∈ support{D}. (2) Bounded variance: There exists σ > 0 such that,
E‖∇f(x; ξ)−∇f(x)‖2 ≤ σ2 ; ∀x ∈ Rd. (3) Expectation over smooth losses: There exists L > 0 such that,
‖∇f(x; ξ)−∇f(y; ξ)‖ ≤ L‖x− y‖ ; ∀x, y ∈ Rd, ξ ∈ support{D} . (4) The last assumption also implies that the expected loss f(·) is L smooth. A property of smooth functions that we will exploit throughout the paper is the following,
f(y) ≤ f(x) +∇f(x)>(y − x) + (L/2)‖y − x‖2 ; ∀x, y ∈ Rd (5)
In the rest of this manuscript,∇f(x; ξ) relates to gradients with respect to x, i.e.,∇ := ∇x. We use ‖ · ‖ to denote the Euclidean norm, and x∗ denotes a global minima of f(·), i.e., x∗ = minx∈Rd f(x).

heading: 4 Method
text: In this section we present STORM+ (STochastic Recursive Momentum +): a parameter-free stochastic optimization method that finds approximate stationary points at an optimal rate. We describe our method in Alg. 1 and Eq. (8), and state its guarantees in Theorem 1.
The original STORM algorithm: The original STORM template of Cutkosky and Orabona [2019] relies on an SGD-style update with a corrected momentum. Concretely, the idea is to maintain a gradient estimate dt which is a corrected weighted average of past stochastic gradients, and then update the iterates similarly to SGD,
xt+1 = xt − ηtdt . (6) Standard momentum is a weighted average of past gradients,
dt = at∇f(xt, ξt) + (1− at)dt−1 ; where at ∈ [0, 1] .
Under this construction, dt is generally a biased estimate of ∇f(xt). In STORM it is suggested to add a correction term ,(1 − at)(∇f(xt, ξt) − ∇f(xt−1, ξt)), which leads to the following update rule (again, at ∈ [0, 1]),
dt = ∇f(xt, ξt) + (1− at)(dt−1 −∇f(xt−1, ξt)) , (Corrected Momentum)
The correction term plays a crucial role here: it exploits the smoothness of f(·, ξ) in a way that leads to a variance reduction effect. To see this effect one can inspect the error of the momentum dt compared to the exact gradient at xt,
t := dt −∇f(xt) .
The STORM update rule induces the following error dynamics,
t = (1− at) t−1 + at(∇f(xt, ξt)−∇f(xt)) + (1− at)Zt
where Zt := (∇f(xt, ξt) − ∇f(xt−1, ξt)) − (∇f(xt) − ∇f(xt−1)). Due to the smoothness of the objective we have ‖Zt‖ ≤ O(‖xt − xt−1‖) = O(ηt−1‖dt−1‖). Intuitively, as we approach a stationary point (and use a small enough learning rate) then ηt−1‖dt−1‖ decreases which in turn reduces the magnitude of Zt’s. Moreover, the second term in the above dynamics, at(∇f(xt, ξt)− ∇f(xt)), can be controlled by choosing a small enough momentum at. Thus, carefully controlling the learning rate and momentum parameters leads to a variance reduction effect which facilitates fast convergence.
The original STORM paper [Cutkosky and Orabona, 2019] makes the following choices,
ηt = θ/
( w +
t∑ i=1
‖gi‖2 )1/3
& at = cL2η2t−1 , (7)
where we denote gt := ∇f(xt, ξt). The above choice of learning rate is inspired by AdaGrad Duchi et al. [2011], which also sets the learning rate inversely proportional to the cumulative square norms of past gradients. Note that θ and w are constants that depend on the smoothness of the objective L, as well as on the bound on the gradients G, and c is an absolute constant independent of the problem’s characteristics. These choices of the constants and especially the choice of at ∝ L2η2t−1 is crucial for the analysis of the original STORM . In fact, the convergence proof for STORM breaks down unless we encode this prior knowledge into ηt and at. Next, we describe our parameter-free version.
Our STORM+ algorithm: STORM+ relies on the original STORM template described in Equations (6) and (Corrected Momentum), with the following parameter-free choices of learning rate and momentum parameter,
ηt = 1/ ( t∑ i=1 ‖di‖2/ai+1 )1/3 & at = 1/ ( 1 + t−1∑ i=1 ‖gi‖2 )2/3 , (8)
where again we denote gt := ∇f(xt, ξt). Note that in contrast to the original STORM our adaptive learning rate builds on history of estimates {d1, . . . , dt} as well as on the momentum parameters {a1, . . . , at+1}. Our momentum term is similar to the adaptive choice of STORM , yet it does not require a bound on the gradients nor on the smoothness parameter, which was crucial for the original analysis. Finally, note that the above choice ensures at ∈ [0, 1].
For completeness we present our method in Alg. 1, where it can be seen that STORM+ is a combination of the original STORM template (Equations (6) and (Corrected Momentum)) together with the specific choices of ηt and at appearing in Eq. (8). Note that the solution that STORM+ outputs is a point chosen uniformly at random among all iterates, which is quite standard in (stochastic) non-convex optimization.
Notation: In Alg. 1 and throughout the rest of the paper we will employ the following notation,
gt := ∇f(xt, ξt) ; g̃t := ∇f(xt, ξt+1) ; ḡt := ∇f(xt) .
Now, we are at a position to present our main theorem regarding STORM+ (Alg. 1):
Algorithm 1 STORM+
Input: #iterations T , x1 ∈ Rd 1: Sample ξ1 and set d1 = g1 = ∇f(x1, ξ1) 2: for t = 1, ..., T do
3: at+1 ← 1/ ( 1 + ∑t i=1 ‖gi‖2 )2/3 & ηt ← 1/ (∑t i=1 ‖di‖2/ai+1 )1/3 4: xt+1 ← xt − ηtdt 5: Sample ξt+1 and set gt+1 := ∇f(xt+1; ξt+1), and g̃t := ∇f(xt; ξt+1) 6: dt+1 ← gt+1 + (1− at+1)(dt − g̃t) 7: end for 8: Choose x̄T uniformly at random from {x1, . . . , xT } 9: return x̄T
Theorem 1. Under the assumption in Eq. (1), (2), (3) and (4) in Section 3, STORM+ ensures,
E‖∇f(x̄T )‖ ≤ O ( M√ T + κσ1/3 T 1/3 ) ,
where κ = O(B3/4 + L3/2); M = O(1 + L9/4 +B9/8 +G5 + (LG4)3/2), and the expectation is with respect to the randomization of the samples as well as the algorithm’s.
Theorem 1 demonstrates that in the stochastic case STORM+ achieves the optimal O(1/T 1/3) rate for our setting. Moreover, it can be seen that STORM+ implicitly adapts to the variance of the noise; in the noiseless case where σ = 0, STORM+ recovers the optimal O(1/ √ T ) rate. We note that scaling the learning rate by some (absolute) constant factor may enable us to obtain better dependence on L and B.

heading: 5 Analysis
text: In this section we provide the convergence analysis of the STORM+ algorithm. We begin with the analysis in the offline case where σ = 0, and establish a convergence rate of O(1/ √ T ) in Section 5.1 for completeness. In Section 5.2, we introduce a simplified version of STORM+ , with a nonadaptive momentum parameter of the form at+1 := 1/t2/3. Due to simplicity and space limitations, it is inconvenient to share the full proof of STORM+ , and this simplified version enables us to illustrate the main steps of the original proof. We show that this version achieves a convergence rate of O(1/T 1/3) in the stochastic case (though it does not adapt to the variance). Finally, in Section 5.1 we provide a proof sketch for STORM+ in Alg. 1 that establishes the result in Theorem 1.

heading: 5.1 Offline Case
text: Here we analyze STORM+ in the case where σ = 0, and demonstrate a rate of O(1/ √ T ) for
finding an approximate stationary point.
Theorem 2. Let f satisfy Eq. (1), (4) and x̄T be generated after running Alg. 1 for T iterations under deterministic oracle. Then it holds that,
E‖∇f(x̄T )‖ ≤ O( √ 1 + L3 +B9/4/ √ T ) .
where we take expectation due to randomness governing the generation of x̄T (see line 8 in Alg. 1).
Proof. In the case where σ = 0 one can directly show by induction that dt = ḡt := ∇f(xt). So the update rule becomes xt+1 = xt − ηtḡt. Now, using the smoothness of the objective implies,
∆t+1 −∆t = f(xt+1)− f(xt) ≤ −ηt‖ḡt‖2 + Lη2t ‖ḡt‖2/2 ,
here we denoted ∆t := f(xt)− f(x∗), where x∗ ∈ arg min f(x). Dividing by ηt, re-arranging and summing gives,
T∑ t=1 ‖ḡt‖2 ≤ ∆1 η1 − ∆T+1 ηT + T∑ t=2 ( 1 ηt − 1 ηt−1 ) ∆t + L 2 T∑ t=1 ηt‖ḡt‖2
≤ B η1 +B T∑ t=2 ( 1 ηt − 1 ηt−1 ) + L 2 T∑ t=1 ‖ḡt‖2(∑t i=1 ‖ḡi‖2
)1/3 ≤ B ηT + L ( t∑ i=1 ‖ḡi‖2 )2/3 ≤ B ( T∑ t=1 ‖ḡt‖2/at+1 )1/3 + L ( T∑ t=1 ‖ḡt‖2 )2/3
≤ B ( 1 +
T∑ t=1
‖ḡt‖2 )2/9( T∑
t=1
‖ḡt‖2 )1/3 + L ( T∑ t=1 ‖ḡt‖2 )2/3
(9)
where the second inequality uses ηt = (∑t i=1 ‖ḡi‖2/ai+1 )−1/3 ≤ (∑t i=1 ‖ḡi‖2 )−1/3 which holds since dt = ḡt and at ≤ 1. We also use that ∆t ∈ [0, B] together with η−1t − η−1t−1 ≥ 0. The third inequality uses Lemma 3 below; and the last inequality uses 1/at+1 ≤ (1/aT+1) =(
1 + ∑T t=1 ‖ḡt‖2 )2/3 , which holds since at is monotonically non-increasing.
By treating the inequality in Eq. (9) as a polynomial of x = ∑T t=1 ‖ḡt‖2, one can derive the following
bound, ∑T t=1 ‖ḡt‖2 ≤ O(1 + L3 +B9/4) . Using the definition of x̄T as well as Jensen’s inequality implies,
E‖∇f(x̄T )‖ := E‖ḡ(x̄T )‖ ≤ √ E‖ḡ(x̄T )‖2 = √√√√ T∑ t=1 ‖ḡt‖2/T ≤ O( √ 1 + L3 +B9/4/ √ T ) .
which establishes the bound. In the proof we have used the technical lemma below,
Lemma 3. Let b1 > 0, b2, ..., bn ≥ 0 be a sequence of real numbers, p ∈ (0, 1) be a real number. n∑ i=1 bi(∑i j=1 bj )p ≤ 11− p ( n∑ i=1 bi )1−p
5.2 Stochastic Case Analysis of Simplified STORM+
Here we analyze a simplified version of STORM+ in the stochatic setting. While this version does not adapt to the noise variance, it exhibits the optimal rate of O(1/T 1/3) in the stochastic case, and its analysis illustrates some of the main ideas that we employ in the proof of the fully adaptive STORM+ (which is more involved).
The version that we analyze here differs from STORM+ in the choice of the momentum parameters. Here we choose a1 = 1 and at+1 = 1/t2/3 ;∀t ≥ 1, in contrast to the adaptive choice that we make
in Alg. 1. Note that we keep the same expression for the step size, ηt = 1/ (∑t i=1 ‖di‖2/ai+1 )1/3 .
Theorem 4. Under Assumptions in Eq. (1), (2), (3) and (4), simplified STORM+ ensures, E‖∇f(x̄T )‖ = O( √ L3 + σ2 +B3/2/T 1/3) ,
Proof. The proof is composed of two parts. In the first we bound the cumulative expectation of errors E ∑T t=1 ‖ t‖2, where t is the difference between the corrected momentum dt and the exact gradient ḡt, i.e. t = dt − ḡt. Thus, in the first part we relate the above sum to the sum of exact
gradients E ∑T t=1 ‖ḡt‖2. Then, in the second part we divide into two sub-cases the first where
E ∑T t=1 ‖ t‖2 ≤ (1/2)E ∑T t=1 ‖ḡt‖2 and its complement. In one of these sub-cases we also use the smoothness of the objective together with the update rule, similarly to what we do in Eq. (9). First Part: Bounding E ∑T t=1 ‖ t‖2. The update rule for dt induces the following error dynamics, t = (1− at) t−1 + at(gt − ḡt) + (1− at)Zt (10)
where Zt := (gt − g̃t−1) − (ḡt − ḡt−1). Letting Ht be the history to time t, i.e., Ht := {x1, ξ1, ξ2, ξ3 . . . , ξt} and recalling that both at and xt depend on history up to t − 1, i.e., Ht−1, immediately implies that E[at(gt − ḡt)|Ht−1] = E[(1 − at)Zt|Ht−1] = 0, as well as E[(1− at) t−1|Ht−1] = (1− at) t−1. Thus, taking the square of the above equation and then taking the expectation gives,
E ‖ t‖2 ≤ (1− at)2E ‖ t−1‖2 + ‖(1− at)Zt + at(gt − ḡt)‖2
≤ (1− at)2E ‖ t−1‖2 + 2(1− at)2 ‖Zt‖2 + 2a2tE‖gt − ḡt‖2
≤ (1− at)E ‖ t−1‖2 + 8L2Eη2t−1 ‖dt−1‖ 2 + 2a2tσ 2 , (11)
where the second line uses ‖b+ c‖2 ≤ 2‖b‖2 + 2‖c‖2, and the last line uses E‖gt − ḡt‖2 ≤ σ2 and (1− at) ∈ [0, 1], as well as the smoothness assumption that implies ‖Zt‖ ≤ ‖gt − g̃t−1‖+ ‖ḡt − ḡt−1‖ ≤ 2L‖xt − xt−1‖ = 2Lηt−1‖dt−1‖. Dividing Eq. (11) by at and re-arranging implies,
E ‖ t−1‖2 ≤ 1
at (E ‖ t−1‖2 − E ‖ t‖2) + 8L2E[η2t−1 ‖dt−1‖ 2 /at] + 2atσ 2 .
Summing the above, and using 0 := 0 gives, E T∑ t=1 ‖ t−1‖2 ≤ − E ‖ T ‖2
aT︸ ︷︷ ︸ (A) +
T−1∑ t=1 ( 1 at+1 − 1 at
)E ‖ t‖2︸ ︷︷ ︸ (B)
+8L2 E[ T∑ t=1
η2t−1 ‖dt−1‖ 2 /at]︸ ︷︷ ︸
(C)
+2σ2 T∑ t=1
at︸ ︷︷ ︸ (D)
(12)
Next we bound all the term on the RHS of the above equation:
Bounding (A): Since aT ≤ 1 we can bound −E ‖ T ‖2/aT ≤ −E ‖ T ‖2
Bounding (B): Note that G(z) = z2/3 is a concave function in R+. Thus applying the gradient inequality implies that ∀z1, z2 ≥ 0 we have (z1 + z2)2/3 − z2/31 ≤ 23z −1/3 1 z2. Hence, for all t ≥ 2,
1/at+1 − 1/at = t2/3 − (t− 1)2/3 ≤ 2(t− 1)−1/3/3 ≤ 2/3 . Moreover, 1/a2 − 1/a1 = 0. These imply that (B) ≤ (2/3)E ∑T t=1 ‖ t‖ 2.
Bounding (C): By definition of ηt we have, (C) = E T∑ t=1 ‖dt−1‖2 /at(∑t−1 i=1 ‖di‖2/ai+1 )2/3 ≤ 3E ( T−1∑ t=1 ‖dt‖2/at+1 )1/3 ≤ 3T 2/9 ( E T∑ t=1 ‖dt‖2 )1/3 .
where the first inequality uses Lemma 3, and the second inequality uses 1/at ≤ 1/aT+1 ≤ T 2/3 as well as Jensen’s inequality with respect to the concave function U(z) = z1/3, defined over R+.
Bounding (D): Lemma 3 immediately implies that (D) = 1 + ∑T−1 t=1 1/t
2/3 ≤ 1 + 3T 1/3 ≤ 4T 1/3. Plugging these bounds into Eq. (12) and re-arranging yields,
E T∑ t=1
‖ t‖2 ≤ 72L2T 2/9 (
E T∑ t=1
‖dt‖2 )1/3 + 24σ2T 1/3 . (13)
Second Part: Bounding E ∑T t=1 ‖ḡt‖2. Here we use the bound of Eq. (13) in order to bound the sum of square gradients. Let us divide into two sub-cases. Case 1 : Assume that E ∑T t=1 ‖ t‖ 2 ≥ 12E ∑T t=1 ‖ḡt‖2. Combining the condition of Case 1 with
‖dt‖2 ≤ 2‖ḡt‖2 + 2‖ t‖2 (due to dt = ḡt + t), implies that E ∑T t=1 ‖dt‖2 ≤ 6E ∑T t=1 ‖ t‖
2. Plugging this inside Eq. (13) yields,
E T∑ t=1
‖ t‖2 ≤ 72L2T 2/9 (
6E T∑ t=1
‖ t‖2 )1/3 + 24σ2T 1/3 .
The above immediately implies that E ∑T t=1 ‖ t‖
2 ≤ O((L3 + σ2)T 1/3), and due to the condition of Case 1 we therefore have, E ∑T t=1 ‖ḡt‖ 2 ≤ O((L3 + σ2)T 1/3). This concludes the first case.
Case 2 : Assume that E ∑T t=1 ‖ t‖ 2 ≤ 12E ∑T t=1 ‖ḡt‖2. Combining the condition of Case 2 with
‖dt‖2 ≤ 2‖ḡt‖2 + 2‖ t‖2 (due to dt = ḡt + t), implies that E ∑T t=1 ‖dt‖2 ≤ 3E ∑T t=1 ‖ḡt‖ 2.
Now using the update rule xt+1 = xt − ηtdt together with smoothness of f(·), one can show in a similar manner to our derivation of Eq. (9) the following bound (we defer this to the appendix),
T∑ t=1 ‖ḡt‖2 ≤ T∑ t=1
‖ t‖2 + 2BT 2/9 (
T∑ t=1
‖dt‖2 )1/3 + 3
2 L ( T∑ t=1 ‖dt‖2 )2/3
(14)
Taking the expectation of the above equation and plugging in E ∑T t=1 ‖dt‖2 ≤ 3E ∑T t=1 ‖ḡt‖
2 as well as E ∑T t=1 ‖ t‖ 2 ≤ 12E ∑T t=1 ‖ḡt‖2 gives,
E T∑ t=1 ‖ḡt‖2 ≤ 1 2 E T∑ t=1
‖ḡt‖2 + 2BT 2/9 (
3E T∑ t=1
‖ḡt‖2 )1/3 + 3
2 L
( 3E
T∑ t=1
‖ḡt‖2 )2/3
(15)
where we also used Jensen’s inequality with respect to he concave functions z1/3 and z2/3 defined over R+. The above immediately implies, E ∑T t=1 ‖ḡt‖
2 ≤ O(L3 +B3/2T 1/3). This concludes the second case. Summary. We have shown that E ∑T t=1 ‖ḡt‖
2 ≤ O((L3 + σ2 +B3/2)T 1/3), combining this with the definition of x̄T and using Jensen’s inequality similarly to what we did in the offline analysis provides,
E‖∇f(x̄T )‖ = O( √ L3 + σ2 +B3/2/T 1/3) ,
which concludes the proof.
5.3 Stochastic Case Analysis of STORM+
Finally, we provide a sketch of the proof for the STORM+ algorithm appearing in Alg. 1. In a high level, the analysis follows similar lines to that of of simplified STORM+ ’s appearing in Section 5.2.
There are two extra challenges compared to the analysis of simplified STORM+ :
1. Now at is a random variable that depends on the noisy samples. 2. The differences 1/at+1 − 1/at are not necessarily smaller than 1.
Recall that in the analysis appearing in Section 5.2 we used 1/at+1− 1/at ≤ 2/3, which was crucial to bounding term (B).
Among the tools that we use to address the first challenge is a version of Young’s inequality, that we mention in the appendix. To cope with the second challenge, when we bound the expectation of∑T
t=1 ‖ t‖2, it is split into two, T∑ t=1 ‖ t‖2 = τ∗∑ t=1 ‖ t‖2 + T∑ t=τ∗+1 ‖ t‖2
where τ∗ is a time-step after which we can ensure that 1/at+1 − 1/at ≤ 2/3. Next we proceed with the proof sketch.
Proof Sketch of Theorem 1. The proof is composed of three parts: (a) In the first part we bound the cumulative expectation of errors E ∑τ∗ t=1 ‖ t‖2, where t := dt − ḡt, and τ∗ is a stopping time after which we can ensure that 1/at+1 − 1/at ≤ 2/3. (b) In the second part we use our bound on E ∑τ∗ t=1 ‖ t‖2 in order to bound the total sum of square errors, E ∑T t=1 ‖ t‖2. (c) Then, in the
last part we divide into two sub-cases the first where E ∑T t=1 ‖ t‖2 ≤ (1/2)E ∑T t=1 ‖ḡt‖2 and its complement. In one of these sub-cases we also use the smoothness of the objective together with the update rule, similarly to what we do in Eq. (9). First Part: Bounding E ∑τ∗ t=1 ‖ t‖2. Recall the error dynamics of STORM
+ appearing in Eq. (10). Taking the square and summing up to some τ∗ ∈ [T ] enables to bound,
τ∗∑ t=1 ‖ t‖2 ≤ τ∗∑ t=1 (1− at)‖ t−1‖2 + 2 τ∗∑ t=1 ‖Zt‖2 + 2 τ∗∑ t=1 a2t‖gt − ḡt‖2 + τ∗∑ t=1 Mt ,
where Mt = 2〈(1− at) t−1, at(gt− ḡt) + (1− at)Zt〉 is a martingale difference sequence such that E[Mt|Ht−1] = 0, where Ht is the history to time t, i.e., Ht := {x1, ξ1, ξ2, ξ3 . . . , ξt}. Also, recall that Zt := (gt − g̃t−1)− (ḡt − ḡt−1). Now let us define β := min{1, 1/G4}, and τ∗ = max{t ∈ [T ] : at ≥ β}. Recalling that at+1 is measurable with respect toHt implies that τ∗ ∈ [T ] is a stopping time. Re-arranging the above and using the definition of τ∗ implies,
β τ∗∑ t=1 ‖ t‖2 ≤ ‖ τ∗‖2 + τ∗−1∑ t=1 at+1‖ t‖2 ≤ 2 T∑ t=1
‖Zt‖2︸ ︷︷ ︸ (i) +2
T∑ t=1
a2t‖gt − ḡt‖2︸ ︷︷ ︸ (ii) +
τ∗∑ t=1
Mt︸ ︷︷ ︸ (iii)
where we used τ∗ ≤ T , as well as β ≤ 1. Next we bound the expected value of the above terms.
Bounding (i). As in the previous section, the smoothness property implies that ‖Zt‖2 ≤ 4L2η2t−1‖dt−1‖2. Using the expression for ηt−1 together with Lemma 3 enables to show,
(i) ≤ 4L2 T∑ t=1
‖dt−1‖2 ( ∑t−1 i=1 ‖di‖2)2/3 ≤ 12L2( T∑ t=1 ‖dt‖2)1/3 .
Bounding (ii). One can directly show that E[a2t‖gt − ḡt‖2] ≤ E[a2t‖gt‖2]. Using this together with the expression for at, it is possible to show that,
E(ii) ≤ E T∑ t=1
‖gt‖2 (1 + ∑t−1 i=1 ‖gi‖2)4/3 ≤ C1 .
where C1 is a constant, and the second inequality is due to a lemma that we describe in the appendix.
Bounding (iii). Since τ∗ ∈ [T ] is a bounded stopping time, and Mt is a martingale difference sequence, then Doob’s optional stopping theorem Levin and Peres [2017] implies E(iii) = E ∑τ∗ t=1Mt = 0.
Conclusion. The above together with Jensen’s inequality for U(z) = z1/3 defined over R+, yields,
E τ∗∑ t=1 ‖ t‖2 ≤ 2C1/β + 24(L2/β)(E T∑ t=1 ‖dt‖2)1/3 . (16)
Second Part: Bounding E ∑T t=1 ‖ t‖2. Recall the error dynamics of STORM
+ appearing in Eq. (10). Dividing by √ at, taking the square and summing up to some T enables to bound,
1 at ‖ t‖2 ≤ ( 1 at − 1)‖ t−1‖2 + 2
‖Zt‖2
at + 2at‖gt − ḡt‖2 + Yt
where Yt = 2〈 1−at√at t−1, √ at(gt − ḡt) + 1−at√ at Zt〉 is a martingale difference sequence such E[Yt|Ht−1] = 0. Re-arranging the above and summing one can show, T∑ t=1 ‖ t−1‖2 ≤ − 1 aT ‖ T ‖2︸ ︷︷ ︸
(A)
+ T∑ t=1 ( 1 at+1 − 1 at
)‖ t‖2︸ ︷︷ ︸ (B) +2
T∑ t=1 ‖Zt‖2
at︸ ︷︷ ︸ (C)
+2 T∑ t=1
at‖gt − ḡt‖2︸ ︷︷ ︸ (D) +
T∑ t=1
Yt︸ ︷︷ ︸ (E)
Now, due to the martingale property E(E) = 0. Next, we focus on bounding term (B), Bounding (B). Using the definition of τ∗ one can show that 1/at+1 ≤ 1/β̃ ;∀t ≤ τ∗, where 1/β̃ := (1/β3/2 +G2)2/3. Moreover, we can show,
1/at+1 − 1/at ≤ 2/3 ; ∀t ≥ τ∗ + 1
This enables to decompose and bound (B) according to τ∗,
T∑ t=1 ( 1 at+1 − 1 at )‖ t‖2 = τ∗∑ t=1 ( 1 at+1 − 1 at )‖ t‖2 + T∑ t=τ∗+1 ( 1 at+1 − 1 at )‖ t‖2
≤ 1 β̃ τ∗∑ t=1 ‖ t‖2 + 2 3 T∑ t=τ∗+1 ‖ t‖2 ≤ 1 β̃ τ∗∑ t=1 ‖ t‖2 + 2 3 T∑ t=1 ‖ t‖2 . (17)
This enables to use Eq. (16) to bound the expected value of term (B). From here the analysis of the other terms and bounding E ∑T t=1 ‖ t−1‖2 is done similarly to our analysis of simplified STORM+ . Third Part: Bounding E ∑T t=1 ‖ḡt‖2. In this part we divide into two sub-cases depending whether
E ∑T t=1 ‖ t‖ 2 ≥ (1/2)E ∑T t=1 ‖ḡt‖2 or not. And continue similarly to our analysis of simplified STORM+ . The rest of the details appear in the appendix.

heading: 6 Conclusion
text: We have presented a novel parameter-free and adaptive algorithm for non-convex optimization that obtains the optimal rate in the setting of expectation over smooth losses while adapting to variance in gradient estimates. Our approach suggests a new way to set the learning rate and momentum jointly and adaptively throughout the learning process, which might open up new avenues to both practical and theoretical developments in this direction.

heading: Acknowledgments
text: This work has received funding from the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme (grant agreement no 725594 - time-data); Hasler Foundation Program: Cyber Human Systems (project number 16066); the Department of the Navy, Office of Naval Research (ONR) under a grant number N62909-17-1-2111; the Swiss National Science Foundation (SNSF) under grant number 200021_178865 / 1; and the Army Research Office under Grant Number W911NF-19-1-0404. K.Y. Levy acknowledges support from the Israel Science Foundation (grant No. 447/20).

Review summaries:
This paper proposes an adaptive variant of the recently proposed STORM algorithm for minimizing expectations of smooth, nonconvex losses. By doing so, it removes the hyperparameters necessary for running STORM. It offers an analysis of the method and rates of convergence, but no empirical validation of the method.

In this paper, the authors provide a parameter-free version named STORM+ of the STORM algorithm. STORM+ uses the same single-loop recursive momentum (which leads to implicit variance reduction effects) and constant batchsizes as in STORM, but features a parameter-free design by constructing the learning rate \eta_t and momentum coefficient a_t using past parameters and momentums (as a comparison, STORM needs to tune these two factors manually). Theoretically, the authors show that STORM+ achieves the same near-optimal complexity as STORM but without needing to tune the parameters to ensure this guarantee. 

This paper is interesting and technically sound. It devises a thoroughly parameter-free algorithm named STORM+, on the basis of STORM. STORM+ obtains the optimal theoretical results and achieves comparable performance with respect to other adaptive methods. The merit of this paper is that it does not require the knowledge of any problem parameters and only initial step-size tuning suffices.