High-probability Bounds for Non-Convex Stochastic Optimization with Heavy Tails
Authors: Ashok Cutkosky
Harsh Mehta
Venue: NeurIPS 2021 Oral (2021)
Keywords: stochastic optimization
heavy tails
non-convex optimization
optimization for deep learning

Abstract: We consider non-convex stochastic optimization using first-order algorithms for which the gradient estimates may have heavy tails. We show that a combination of gradient clipping, momentum, and normalized gradient descent yields convergence to critical points in high-probability with best-known rates for smooth losses when the gradients only have bounded $\mathfrak{p}$th moments for some $\mathfrak{p}\in(1,2]$. We then consider the case of second-order smooth losses, which to our knowledge have not been studied in this setting, and again obtain high-probability bounds for any $\mathfrak{p}$. Moreover, our results hold for arbitrary smooth norms, in contrast to the typical SGD analysis which requires a Hilbert space norm. Further, we show that after a suitable "burn-in" period, the objective value will monotonically decrease for every iteration until a critical point is identified, which provides intuition behind the popular practice of learning rate "warm-up'' and also yields a last-iterate guarantee.

High-probability bounds for Non-Convex Stochastic Optimization with Heavy Tails
Ashok Cutkosky Boston University ashokOcutkosky.o com
Harsh Mehta Google Research harshn@goggle.com com
Abstract
We consider non-convex stochastic optimization using first-order algorithms for which the gradient estimates may have heavy tails. We show that a combination of gradient clipping, momentum, and normalized gradient descent yields convergence to critical points in high-probability with best-known rates for smooth losses when the gradients only have bounded pth moments for some p E (1.2). We then consider the case of second-order smooth losses, which been studied in this setting, and again obtain high-probabee probability bounds for any p. Moreover, results hold for arbitrary smooth norms, in contrast to the typical SGD analysis which requires a Hilbert space norm. Further, we show that after a suitable "burn-in" period, the objective value will monotonically decrease whenever the current iterate is not a critical point, which provides intuition behind the popular practice of learning rate "warm-up" and also yields last-iterate guarantee.
knowledge have not
1 SGD and Heavy Tails
The standard algorithm for training large machine learning models today is stochastic gradient descent (SGD) and its myriad variants [13, 11, 19, 39, 28]. SGD enjoys many properties that help explain its empirical success: it can be implemented very easily, can efficiently process training data in a streaming manner, and even achieves the optimal convergence rate for finding critical points of smooth non-convex objectives [2]. In particular, SGD can be understood optimizing the following stochastic optimization problem:
min F i = [ff((i) =]]
where f fuū, 2) is some loss function taking model parameters and random example 2, where has some distribution P.. Then SGD consists of the update:
Th यh V f (w7, 21)
where 21...., 27 are i...d. random variables distributed according to Ps, and is some real number called the learning rate. Assuming that F is L-smooth and that f 2) satisfies TE[1 V f 2)] s G2 for some G for all wi SGD satisfies [13]:
LA T
G VT
E
V VF
5o
where = -in This implics that SGD takes O(E-4) iterations to find an e-critical point, which is point प such that [V [13]. When the stochastic gradients are further assumed to have light sub-uuusee -abssianniin tails, [21] shows that SGD with momentum can achieve this same result in high probability well.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

However, recent work [30, 31, 42, 41, 43] suggests that the assumptions of bounded variance or light tails may be too optimistic in practical deep learning problems. For example, Zhang et al. [43] provides empirical evidence that large NLP models based on attention and transformers [35, 10] have heavy-tailed gradients. In particular, the variance is extremely large (perhaps even essentially infinite for practical purposes), but the pth moment is bounded for some p (1, 2). This setting is
significantly more complicated, and SGD's convergence properties are difficult to characterize even when assuming slightly stronger tail bounds [29] various versions of convexity [36].
standard tool when working with heavy-tailed random variables is truncation, or clipping. Suc- cinctly, if X is a random variable such that x[]]X10 = oc, X itself may be poorly behaved due to the heavy tail, but a truncated value X = [XT min( will be much more benign so long as is chosen in an appropriate manner [6, 23]. Fortuitously, gradient clipping has long been a standard technique to improve training of deep neural networks [26], and recent analysis [43, 14] suggests that clipping allows SGD to deal with heavy-tailed stochastic gradients. Specifically, [43] shows that
3p-2 iterations.
SGD combined with clipping can find an e-critical point in expectation in O p
We propose to improve upon this backdrop in three important ways. First, the prior analysis of [43] for non-convex convergence of SGD under heavy tails only provides an in-expectation analysis. In-expectation results are unsatisfying because the quantity measured for convergence (i.e. gradient norm or function value) is itself a random variable, and so might also have a heavy tail. In this case individual training runs could be quite poor. Since training modern networks is quite expensive [27, 5), it is important to be confident that each and every training run will produce a high-quality result. We provide a variant of the normalized SGD with momentum algorithm of [9] that incorporates
gradient clipping and show that the method finds an e-critical point in Ô high probability.
iterations with
Second, the vast majority of first-order optimization methods for stochastic non non-convex optimization consider exclusively the standard Lg or Hilbert-space norm. Such analyses can take advantage of the fact that (V F()), 9) [>>))) when ÿ is an unbiased estimate of V Fz)). In contrast, our algorithms find critical points when the gradient is measured with any smooth norm (i.e. in a Banach space). In this case, the gradient estimate ÿ (which is a dual vector) must be converted into an appropriate primal vector. This operation is non-linear in general and so may add some bias. turns out the key to operating with different norms is the popular practice of augmenting SGD iterates with momentum. When using momentum-ban argument, and so the techniques can be extended to more general norms. This provides more concrete evidence of the theoretical advantages of momentum.
updates we do not rely on the same unbiasedness
Third, recent theoretical advances in non-convex optimization have produced a number of new algorithms that avoid the lower-bounds of [2] by assuming extra structure on the objective F, such as secor smoothness. In this case, [32, 1, 12, 9, 3] provide algorithms that achieve faster convergence rates. In particular, the algorithms of [12, 9] can find an c-rritical point in ((335) iterations using first-order stochastic oracle, which is the optimal rate [3]. To our knowledge, no analysis even in expectation exists for second-order smooth losses under heavy tails. We fill
this gap. providing an algorithm that identifies c-critical point in Ô 2p-2 iterations with
high probability. Further, in both smooth and second-order smooth cases we provide a convergence guarantee for the last iterate of the optimization algorithm: with high probability, the objective value for the final iterate will be smaller than the objective value at some approximate critical point.
Finally, we discuss intriguing connection the popular warm-up learning rate schedule [16, 15, 10J. in which the learning rate starts small and increases over the first iterations of training. We show that for normalized SGD with momentum, after a brief "burn-in" period, the objective value decreases monotonically until an approximate critical point is reached. This provides an alternative motivation for warm-up: it would be reasonable keep the learning rate small during the "burn-in" period. There is a long line of empirical research in understanding the need for warm-up. Warm-up has been shown to be connected the use of noryyarromaaiiitionn and that improved use of layer-norm can enable warm-up free training of transformer models [7, 37, 25]. Huang et al. [18] further build on these insights and proposes a weight initialization scheme which removes the need for both layer normalization and warm-up while training transformer models. Finally, Liu et al. [22] demonstrate that the variance of the learning rate in optimizers like Adam is abnormally high in the early stages of

---

the optimization process and further hypothesize that warm-up aids by reducing the variance in that stage. Our results provide a different view from an optimization perspective.
Paper Organization: in Section 1.1 we provide our concrete assumptions and setup. In Sections and 3, provide our algorithms and high-probability convergence bounds. In Section 4, we provide some empirical observations on the effect of the burn-in phase in our analysis.
1.1 Setup and assumptions
We assume that the objective F(ū) [i 2)] is a Banach space B to the reals R. Let I| be the norm in B and || ||+ be the dual norm in the dual space B* To preserve analogy and notation from the Hilbert space setting, we will write (v. w) to indicate the application of dual vector v E B* vector E B. We assume that || !| is Frechet differentiable and satisfies for some C for all x y E B:
function mapping a
| yll < - + (V]<>).,)) + C/
(1)
Such Banach space is called (v, CC Note that if || is the standard norm in R for p (1,2), then || || is the corresponding g norm, and the dual space (2, - )smooth. In the common case that B is Hilbert space (e.g. I with the euclidean norm), then B = B* ||-1I= and p and C
In this paper, we will exclusively consider the case p = 2 and C' 1. Finally,
given any E B*, we will write d(v) E B to indicate a unit vector satisfying both [[d(v)] || = 1 and (v, d(v)) ||=||.. which we assume exists for all v. For more background on Banach spaces, see
Appendix E.
We assume F is L-smooth, which means that IVF()) VF S Lli El| In Sec- tion 3, we will also assume that F is p-second-order smooth, which means that F(()=((w) p||u T|||||| for all ", B. We assume that Vf(w, satisfies both B f(w, z) - VF(T)) < G and [
F=F(())
< G for some G and some E (1,2). We will design algorithms such that after T stochastic gradient evaluations, can output a point w such that [[ is small with high probability.
2 Normalized SGD with Momentum
In this section, describe algorithms for finding critical points with high probability in non- convex objectives. Our analysis makes use of gradient clipping, momentum, and normalized updates. Clipping is standard technique for mitigating the effect of heavy tails, essentially by throwing out outlier data. Momentum useful as intuitively averages together many clipped gradients, resulting in gradient estimate that concentrates about its mean with high probability. Normalizing the updates allows for significantly simplified analyses because the bias terms introduced by the moving average in momentum are now very precisely controlled by the learning rate, originally analyzed for Hilbert spaces in [9]. In order to analyze this algorithm, we need the following critical Lemma, which characterizes the effect of taking a normalized SGD step (proved in appendix E).
Lemma 1. Suppose F + R be a Frechet-differentiad
L-Smooth function from a Banach space B to the reals. Let w E B. Let g" E B and let E B be a unit-vector satisfying {g" .9) I Define w' 18-119 ng. Define g" FF(w). Then:
Laf 2
F(w') S F(w) [F(()) + 2n/lcl/+ +
Theorem 2. Suppose E. V V =) S Go for all & for some G. Suppose F L-smooth. Set O, Tम¹ G//t/P, Then with probability least 1 -8:
b and 17
. for arbitrary constant b and satisfying S 1 and set
T
IVF(O)) So (log(T
T
T T
where the big-Ol Oh hides constant that depend on L, G b, C and F(Wi) F(0r+1). but not T 8.
3

---

Algorithm 1 Normalized SGD with Clipping and Momentum
Input: Initial Point W7, learning rate 77, momentum parameter clipping parameter т, time
horizon T:
Set mi = 0.
for 1 = 1.....T do
Sample 2, Pz.
Set 9t = [V [v min(T, V
vf(.).)))
Set m = 3m + 3)0r :lipp
Set uh++1
md(mm))
end for
This Theorem provides analog of the standard in-eppectation result for non-convex optimization, but now the result holds with probbbrbbability. Note that this matches the in-expectation convergence rates from Zhang et al. [43] up logarithmic factors. The next Theorem takes this theme one step further: it shows that after brief "burn-in" period, the decrease in objective value will in fact be monotonic with high probability:
Theorem 3. Under the assumptions of Theorem 2, define the constants:
K = 10C max(1, +0 / (og((1,)))(3110) +
Z = sL + GK
Then, with probability at least 1 8, for all t satisfying:
J p log(T) + log(G) - log(2)
t t2 T== =
b
3p 2
we have
2z
m, VF(ii))ll S
T
Further, so long as
6%
Ls
|m II+ 2
+
29-2
5p
2T
we have ((1++1) < F(W))) | Moreover, whenever (2) not satisfied for 2 T, we must have
14Z
Ls
VF <
T
T T3-2
Finully, if the iteration with smallest value of || such that 2 T, we have:
[FF()) < (log(T/8)
7
Theorem shows that after a small burn-in period T o(T), me very likely to be an extremely good estimate of the true gradient V VFZz). Moreover, we can empirically detect this event by monitoring the norm of the exponentially weighted moving average m. This burn-in period is especially interesting given the success of learning rate schedules that incorporate "warm-up" component, which the learning rate actually increases during the first iterations. Our result seems to hint mechanism for this schedule: by keeping the learning rate small during the burn-in period, we can limit any increase in the objective value until we reach the stage at which it begins to decrease steadily.
Further, Theorem 3 provides a very concrete algorithm to find critical points: after the burn-in period of Ö(Tsp=>)) iterates, we can simply return the point H, with the smallest value of [|l] Intuitively.
this works because the objective value can only increase by a total of O ()" = OT 3) "2 ((1) during the burn-in period so even without learning rate "warm-up", we should not expect the burn-in period to have too significant an effect on the objective value. After the burn-in period, m, is make steady progress until critical point is encountered.
always very close to V (.W,),

---

Proof of Theorem 2. Define & = n - VF(W)). and = clip VF(W)). Also, define S(a,b) b)
(()) (b... By smoothness, we have [| w1+1)ll S L W3 W1411l+
From the definition M441, we have the following recursive formulation for any 2 1:
m (1 (.FF)((((()))) + agg+1
clip
(+( (1-a))))(1(((( "6 + E)) + ac++1
4+1 (1 a S 1 m(+1) + (1 a) + ac++1
(3)
Further, by defining ma 0 and uāg w that én = -V (wT), have that (3) holds for t as well. Now, we unravel the recursion for all iterations:
4+1 (1-==11
a1)
Σ1- c01- + (1-a) E ai S( w2+--1)
n==0
p=0
Next, take the magnitude of both sides and use triangle inequality (using the fact that [1,, 11+1)ll= & Ln):
|441ll. $ (1 - a) I1ooll +a E - a) E+1-1
E a)
--=
M==
S (1 (1-a) (1-w) + (1 (1-a)mL
+ (11
where in the second line have used mo 0 and |VF G, which follows since F E=IlVV //P G.
Now, we invoke Lemma 9, a concentration inequality in Banach space (see appendix B). This implies that with probability at least 1 - 8/7:
(1 (1- (1-a)'
Σ 11 aa))ac)+)) c+t1r P'=
10CT max(1, log(3T/5)) +
Σ
TP-1
t'=0
1 < ak)
1/2
+4
2 -P-2-p
V /maxx(,log((T78))
F==
V //max(, 1, log (3T/8))
S 10CT max(1, log(3/) + aTp-1 +
1/2
Therefore, with probability at least - 5/T:
(1- a)n 10Cat log || s -1 a) (1
a
GP
+ +4(aCGP,-D))1) mmax(( +
7P-1
G Now, with -= .l/p and D3 = (og(3//8)) this becomes:
(1 (1- (1-a)nL annL 1-11 10CDs + 10CDs Dss
(1 + G
a anLL (1 + G+
Ga 1-1pKK
(4)
a
where we define the constant K = 10C'D + 12 DD + 1.
Now, we have from Lemma 1:
V(())) < F( - (WT))
LT'
T
+
|| +
"
2
tml
5

---

so with probability at least 1 - 8:
F (TT) F(UT)
LTT
2 (1 a) (1 (1-a)nL
+ Ga
1
2
F (1))
F(UT)
1.T
2G(1 a)
2(1 a)nTT
2GTa
2
a
a
Now set a
and 17 =
This yields:
T
2p
VF F
1- =0 (T
||
log Log
T
The proof of Theorem follows essentially the same idea as that of Theorem 2. However, this time use the explicit bound on ||&]]+ and observe that after the burn-in period, the contribution from ||ca |I+ has decreased exponentially to be insignificant, and so my will very high-quality estimate of F# every single iteration.
Proof of Theorem 3. Use the settings 1)
T- and a
b and notice that from equation (4) T
that with probability 6, for all t:
GKb
Z
|| ||- < (1 - a)
-
a P-1
T
T
Further, we have
log(1 - < -a
log (G(1 (G))) < to + log(G)
Therefore, if 2 T, then we have with probability at least 1 - 5:
2Z
1/++1/l- <
P-1
This implies the bound on |VF (w.) II when ||m,]|+ is smaller than the threshold described in the Theorem statement. Next, again from Lemma 1,
La7
F( F(w.) - v/ //l +
21/lee +
2
Lnd
F(w.) m/lmmell- 31/lée/l+
6%
Ls
F( - m/llmell- +
+
T
2T 2n-1
from which conclude that if ||mael= 2
62
Ls
F( S F( m +l+
T
2T
Now, for the final part of the Theorem, observe that since G-Lipschitz, we must have that
F(ØT) - F(wi) S GT O(lo((T)
so that F(Ør)= ((1++1) S (F ( F( + log(())) Now consider two cases, either
6z
Ls
min 227
| <
+
Tp-T
30-7
27
In the first case, by bound || the desired bound |[V[()))]] follows. In the latter case, we have
T
F(0r+1) $ F(Wr)
1
Σ
mell
2 T
T

---

So that
F( F(Ør++ 1) + log(T) =Ö
min t27
0o
Tn
Tp
And so again the result follows from our bound on ||
In addition to providing method for identifying critical points, Theorem also provides intuitively desirable guarantee about the last-iterate of the algorithm wr. Specifically, one can easily check that the following Corollary:
Corollary 4. Under the notation and assumptions of Theorem 3, with probability at least 6, there exists some in such that both of the the following inequalities hold:
(oog(T/))
|VF 00
T
F(WT) S F())
where WT is the last iterate of Algorithm
Proof. Set & be the last index such that Wiy does not satisfy (2). Set w.g. Then Theorem implies the result.
In most applications we are not actually interested in finding critical points, but instead wish to actually minimize the objective F. This observation tells that objective value is at least never worse than it would have been had indeed searched specifically for a critical point. Moreover, by providing a guarantee about Wh, have closer match practical use-cases: many of the theoretical analyses of non-convex stochastic gradient methods [13, 38, 20] only show that randomly selected iterate has small gradient norm. In contrast, Corollary provides non-asymptotic guarantee for the final iterate, which is the iterate most likely to be deployed after training in practice.
3 Second-Order Smooth Losses
In this section, we provide our extension to losses that are second order smooth. The overall method very similar: we employ normalized SGD with momentum. However, in order to take advantage of the second-order smoothness, we will need to use a more advanced form of momentum. Specifically, we use implicit gradient transport [4], as implemented by [9] in their NIGT algorithm. This algorithm replaces the standard momentum update with an extrapolation procedure:
B W w.-1)
m, --, (1-3)V T7,
B
By evaluating the gradient at the "shifted" point W. + (( and performing a second-order Taylor expansion enabled by the second-order smoothness assumption, it is possible to show that mh is less-biased estimate of V F(,,) than it would be using the standard momentum update. We augment this procedure with gradient clipping in Algorithm 2 below, and provide its analysis in Theorems 5 and 6, which are directly analogous Theorems 2 and 3.
1-8
Theorem 5. Suppose Ez [xllV P)] GP for all t for some G. Suppose F is L-smooth and p-second-order smooth. Set = a, a ip and T 9p-r for arbitrary constants b and satisfying 5 1, and set at/5.
T
G
Then with probability at least 1 8:
T
(og(T/5)
Σ [VF((B)]+ 50
2p-2
T 3
where the big- Oh hides constant that depend on L, P. G, b, C, and F(i)) F(ur+1), but not T or 8.
We also have direct analog of Theorem

---

Algorithm 2 NIGT with Clipping
Input: Initial Point Wi learning rate 7), momentum parameter clipping parameter 7, time
horizon T:
Set mi =.
for 1. = 1....Td do
Sample 2, Pz.
T = w; + -(01-01-1)
1-8
Set 9t = IV (BB.a)l) min(r, +
f(
Set m, = 3m + (1 3)m(" :
Set 1h+11 Wh m)()))) end for
Theorem 6. Under the assumptions of Theorem 5, define the constants:
K = 10
+ 4 Cmxx(1, (ogxx(1)oo(33711) (3)))) +
Z ps²
F 2
b
Then, with probability at least
8, if
3
2pp 2
1 12T
log(T) + log(G)) - log( 2))
b
5p 3
Then we have
2Z
mm VF(+.)| <
T 3
And so long
6%
Ls
M |1+ 2 2
(5)
2p-2 5p-3
Te-T up-s
27
we have F (We+1) < F(w,) - 2 1 Moreover, (5) is ever not satisfied, we must have
14Z
L
VF((1)) <
2p-2
3p-1
T
5p
T
50-
Finally, ift the iteration with smallest value of || such that 2 T, we have with probability at least 1 8:
log (T/5)
[V(())) 50
T 2
Finally, note that we have a statement about the last iterate of Algorithm 2 that is directly analogous to Corollary 4:
Corollary 7. Under the notation and assumptions of Theorem 6, with probability at least - 5, there exists some in such that both of the following inequalities hold:
log (T5)
|VF( 50
3
F(ur) F(6)
where Wi is the last derate of Algorithm

---

0.70
0.6
0.65
Eval Accuracy
bumin _etepeee
Accuracy
0.4
warmup 3125
0.60
bumin 20mps=250
0.3
burnin stops=300
buniia stmps- 1000
0.55
werstup, 3125
bumtnn steps=2000 2000
buminn steps=0
0.2
burtia
0.50
buminn 250
buminn 500
burain
bumin stepa=2000
buminn pirpe
0.00
0.40
20
100
1k training steps
1K training steps
Figure 1: Comparison between warm-up and burn-in strategies on Masked language modeling BERT- Large pre-training task. Plot (a) shows the eval accuracy when trained with Algorithm I and plot (b) does the same when trained with LAMB.
Experiments
Our theoretical analysis suggests that there is short "burn-in" phase in algorithms during which the momentum estimate m, may be far from the true gradient V and after this period, m, will always be close to (..,). As result, during the burn-in phase it is possible for the optimizer to make slow or even negative progress, while afterwards we should expect more steady improvement. might set the learning rate 77 O throughout the burn-in phase and only begin
This suggests that taking steps after my has accumulated enough gradients to be come a good estimate. This procedure reminiscent of the common practice of learning rate "warm-up", in which the learning rate starts O and is linearly increased to some desired value during the initial phase of training, and then later decreased. Thus, we sought answer how much of the success of the warm-up schedule can be attributed to burn-in.
To answer this question, conducted experimental study the BERT pretraining task. We chose to experiment on BERT since it has been shown to have heavy tails empirically [43] and the common practice is already clip the gradients. Figure 3 plots the masked language model accuracy of BERT-Large model when trained with Algorithm 1. We also include experimental results using the LAMB optimizer instead of Algorithm 1. LAMB optimizer has been shown to obtain state of the art performance on BERT task [40]. Moreover, both our analysis and the LAMB optimizer relies on normalized updates. Note that we also include experimental results using the smaller BERT-Base model in the appendix. All our experiments were conducted using the Tensorflow framework on TPUv3 architecture.
When using Algorithm 1, we employ base learning rate To of 0.3 and batch size of 512. We came up with this learning rate by performing grid search ?o [1.0, 0.5, 0.1, 0.05, 0.01, 0.005, 0.001, 0.0005, 0.0001] and choosing the one which attains the best eval accuracy. To obtain the optimal base learning rate of 0.3, again ran grid search To E (0.1,.2. 0.3, 0.4, 0.5). Our warmup baseline employs the standard practice linear warm-up for 3125 steps and polynomical decay of π V/ (1 $) for the rest of the steps. For fair comparison with the baseline, we start the polynomial decay after 3125 steps regardless of how many steps were used for burn-in. We set the learning rate to the constant base learning rate value of 0.3 for BERT-Large for steps in between burn-in and polynomial decay. For the models trained using burn-in in, instead of linear warm-up, we set the base learning rate 7/0 zero for the amount of steps we run the burn-in for. As shown in Fig 1, we perform a grid search across [0, 250, 500, 1000, 2000, 3125) steps for the initial burn-in period and compare their results. Notice that, without warm-up or burn-in (i.e. O burn-in steps), we see very poor performance. Further, this poor performance persists for smaller values of burn-in, but with large enough burn-in values we suddenly see very good results. Once this critical level of burn-in is achieved, we continue to see small gains as the amount of burn-in is increased. This suggests that without burn-in or warm-up, the first iterations are taking extremely poor steps in the non-convex landscape that may destroy some favorable properties of the initialization point. With enough burn-in, each step is well-aligned with the true gradient and so we may avoid this bad behavior. Further, see that employing burn-in recovers almost all of the advantage gained

---

from employing warm-up (final accuracy of 62.95% with burn-in vs. 63.45% with warm-up). These results lend empirical support to our theoretical intuition that the initial steps of the optimization may be unhelpful, or even hurtful.
For the LAMB optimizer, we reproduce all the hyperparameters from You al. [40]. Most impor- tantly, we train all our models using the base learning rate 7/o of 6.25e-4 and batch size of 512. As shown in Fig 1, again the warmup baseline employs the standard practice of linear warm-up for 3125 steps [40] and polynomical decay of = 100 * (1 - +) for the rest of the steps. Again, for fair comparison with the baseline, we start the polynomial decay after 3125 steps regardless of how many steps were used for burn-in. We set the learning rate to the constant base learning rate value of 6.25e-4 for steps between burn-in and polynomial decay. Note that when using the LAMB optimizer, warm-up baseline does considerably better than models where burn-in was employed. However, notice that the gap between employing warm-up and using neither warm-up nor burn-in is considerably smaller than with A Algorithm 1, for which not using either technique results in extremely poor accuracy. This suggests that the additional heuristics employed by the LAMB optimizer over the simple normalized SGD approach of Algorithm I are providing some non-trivial robustness that remains to be understood. Further, there is still a noticeable gap between the linear learning-rate increase in warm-up and the coarser burn-in policy. It remains to be seen exactly what mechanism this smoother policy employs to aid optimization.
Conclusion
We have provided algorithms based normalized SGD with momentum and gradient clipping that provide high probability- gradients. We also show that by- by-product of our high- sish-probbbility bounds, the algorithms provide a pleasing result for the last iterate of the optimization procedure. Specifically, the last iterate has function value smaller than the function value of critical point.
for finding critical points of non-convex objectives with heavy-
Finally, we identify critical "burn-in" phase early in optimization in which the optimizer may be making poor decisions. Empirically, we show that keeping the learning rate zero for the first iterations indeed does help training compared to simply running the optimizer with no special treatment during this phase. It is interesting question to precisely identify what other ad advantages are provided by the increasing schedule used by warm-up during this phase.
The authors thank the anonymous reviewers for many helpful suggestions. AC was a visiting researcher at Google when this work was completed.
Acknowledgments:
References
[1] Allen-Zhu. Z. Natasha 2:
Faster non-convex optimization than sgd. Advances in neural information processing systems, pp. 2675-2686, 2018.
[2] Arjevani. Y., Carmon, Y., Duchi, J. C., Foster, D. J.. Srebro, N., and Woodworth, B. Lower bounds for non-convex stochastic optimization. arXiv preprint 1rXiv:1912.02365,, 2019.
[3] Arjevani, Y., Carmon, Y., Duchi, J. C., Foster, D. J., Sekhari, A., and Sridharan, K. Second-order information in non-convex stochastic optimization: Power and limitations. In Conference on Learning Theory, pp. 242-2922 2020.
[4] Amold, S. M., Manzagol, P.-.,, Babanezhad, R., Mitliagkas, I,, and Roux, N. L. Reducing the variance in online optimization by transporting past gradients. arXiv preprint /rXiv:/906.084 2019.
[5] Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A.,
Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert- Voss, A., Krueger, G., Henighan, T. Child, R., Ramesh, A., Ziegler, D. M., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners, 2020.
[6] Bubeck, S., BisacBiancch,, N., and Lugosi, G. Bandits with heavy tail. IEEE Transactions on Information Theory, 59(11):7111-7117 2013.
10

---

[7] Chen, M. X., Firat, O., Bapna, A., Johnson, M., Macherey, W., Foster, G., Jones, L., Schuster, M., Shazeer, N., Parmar. N., and al. The best of both worlds: Combining recent advances in neural machine translation. Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), 2018, doi: 10.18653/v1/p18- 1000//8888 URL. http://dx. doi org/ 10. 18653/v1/P18-100 1008.
[8] Cutkosky, A. Algorithms and Lower Bounds for Parameter-free Online Learning. PhD thesis, Stanford University, 2018.
[9] Cutkosky, A. and Mehta, H. Momentum improves normalized sgd. arXiv:2002.03305, 2020.
arXiv preprint
[10] Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume (Long and Short Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019, Association for Computational Linguistics. doi: 10. 10.18853/v1/N19-140 1423. URL https: / //ww.//wwww/o//// aclweb. org/ herg/anthology//ee -
[11] Duchi, J., Hazan, E.,, and Singer, Y. Adaptive subgradient methods for online learning and stochastic optimization. In Conference on Learning Theory (COLT), 2010.
[12] Fang, C., Lin, Z., and Zhang, T. Sharp analysis for nonconvex sgd escaping from saddle points. In Conference on Learning Theory, pp. 1192-1234, 2019.
[13] Ghadimi, S. and Lan, G. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM Journal on Optimization, 23(4):2341-2366 2013.
[14] Gorbunov, E., Danilova, M., and Gasnikov, A. Stochastic optimization with heavy-tailed noise via accelerated gradient clipping. In Neural Information Processing Systems, 2020.
[15] Goyal, P., Dollár, P., Girshick, R., Noordhuis, P., Wesolowski, L., Kyrola, A., Tulloch, A., Jia, Y., and He, K. Accurate, large minibatch sgd: Training imagenet in hour, 2017.
[16] He, K., Zhang, X., Ren, S., and Sun, Deep residual learning for image recognition. 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2015.
[17] Howard, S. R., Ramdas, A., McAuliffe, J., and Sekhon, J. Time-uniform. nonparametrie, nonasymptotic confidence sequences. arXiv preprint arXiv:1810.08240, 2018
[18] Huang, X. S., Perez, F., Ba, J., and Volkovs, M. Improving transformer optimization through better initialization. In III, H. D. and Singh, A. (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 4475 4483. PMLR, 13-18 Jul 2020. URL http://prooceeding: mlr h.r.peess//v1//
huang20f html.
[19] Kingma, D. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint 1411v9141.66980, 2014.
[20] Li, X. and Orahona, F. On the convergence of stochastic gradient descent with adaptive stepsizes. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 983-992. PMLR, 2019.
[21] Li, X. and Orabona, F. A high probability analysis of adaptive sgd with momentum. arXiv preprint arXiv:2007. 14294, 2020.
[22] Liu, L., Jiang, H., He, P., Chen, W., Liu, X., Gao, J., and Han, J. On the variance of the adaptive learning rate and beyond. In International Conference on Learning Representations, 2020. URL httpp://openrevier. n wet.://orrrmm/ddrrrr//22200000
[23] Lugosi, G. and Mendelson, S. Mean estimation and regression under heavy-tailed distributions:
A survey. Foundations of Computational Mathematics, (15)::145-11999 2019.
[24] Minsker, S. On some extensions of bernstein's incquality for self-adjoint operators. arXiv preprint 111.5:1112.5448, 2011.

---

[25] Nguyen, T. Q. and Salazar, J. Transformers without tears: Improving the normalization of self-attention. ArXiv, abs/1910.05895. 2019.
[26] Pascanu, R., Mikolov, T., and Bengio, Y. Understanding the exploding gradient problem. CoRR, abs/1211.5063, 2(417):1, 2012.
[27] Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, J. Exploring the limits of transfer learning with unified text-to-fext transformer. Journal of Machine Learning Research, 21(140):1-67, 2020. URL http: http:://mmmmrrer org/ httgg/pppps/////// / 20-074. html.
[28] Reddi, S. J., Kale, S., and Kumar, S. On the convergence of adam and beyond. In International Conference on Learning Representations, 2018.
[29] Scaman, K. and Malherbe, C. Robustness analysis of non-convex stochastic gradient descent using biased expectations. Advances in Neural Information Processing Systems, 33, 2020.
[30] Simşekli, U., Gürbüzbalaban, M., Nguyen, T. H., Richard, G., and Sagun, L. On the heavy-tailes theory stochastic gradient descent for deep neural networks. arXiv preprint ww1Xiv:912.000 2019.
[31] Simsekli, U., Sagun, I.., and Gurbuzbalaban, M. A tail-index analysis of stochastic gradient noise in deep neural networks. In International Conference on Machine Learning pp. 5827- 5837. PMLR, 2019.
[32] Tripuraneni, N., Stern, M. Jin, C.,, Regier, J., and Jordan, M. I. Stochastic cubic regularization for fast nonconvex optimization. In Advances in neural information processing systems, pp. 1899-2908. 2018.
[33] Tropp, J. Freedman's 's inequality for matrix martingales. Electronic Communications in Proba- bility, 16:262-270, 2011.
[34] Tropp, J. A. An introduction to matrix concentration inequalities. arXiv preprint arXiv:/501.00 1501.0151, 2015.
[35] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.. Gomez, A. N.. Kaiser, E., and Polosukhin, Attention all you need. Advances in Neural Information Processing Systems, 30:5998-6000 2017.
[36] Wang, H., Gürbüzbalaban, M., Zhu, L., Şimsekli, U., and Erdogdu, M. A. Convergence rates of stochastic gradient descent under infinite noise variance. arXiv preprint 21X1v:2102.10346. 2021.
[37] Wang, Q. Li, B., Xiao, T., Zhu, J., Li, C., Wong, D. F., and Chao, L. S. Learning deep transformer models for machine translation. Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019. doi: 10.18653/v1/p19-111 http:/ doi. . org/ 10. 11663//111199111777
URL
[38] Ward, R., Wu, X., and Bottou, L.. Adagrad stepsizes: Sharp convergence nonconvex landscapes. In International Conference on Machine Learning, pp. 6677-6686. PMLR, 2019.
[39] You, Y., Gitman, L., and Ginsburg, B. Large batch training of convolutional networks. arXiv preprint arXiv: arXiv:/708. 03888, 2017.
[40] You, Y., Li, J., Reddi, S., Hseu, J., Kumar, S., Bhojanapalli, S., Song, X., Demmel, J., Keutzer, K., and Hsieh, C.-J. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations, 2020. URL https: //ap:::///ppeeee /ppnreviee... net/forum?an idd SyxissttE........
[41] Zhang, B., Jin, J., Fang, C., and Wang. L. Improved analysis of clipping algorithms for non-convex optimization. arXiv preprint arXiv:2010.02579, 2020.
[42] Zhang, J., He, T., Sra, S., and Jadbabaie, A. Why gradient clipping accelerates training: A theoretical justification for adaptivity. In International Conference on Learning Representations, 2019.
12

---

[43] Zhang, J., Karimireddy, S. P.. Veit, A., Kim. S., Reddi, S., Kumar, S., and Sra, S. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33, 2020.
13


Reviews:
[{"review_id": "j76n3qQkW7", "reviewer": "Reviewer_WuQp", "summary": "This paper focuses on high probability convergence behavior of nnormalized SGD with momentum and gradient clipping under heavy-tailed noise but finite moments. Specifically, without the light-tail assumption, e.g. sub-Gaussian noise, and with bounded pth moment of the gradient for $p \\in (1,2]$,\n-\tThey propose an algorithm that achieves $\\mathcal O (\\log(T/\\delta) / T^{\\frac{p-1}{3p-2}})$ convergence with high probability. They also characterize the behavior of last iterate and corresponding function value.\n-\tUnder additional second-order smoothness, they propose a modified version of the aforementioned algorithm, which achieves convergence rate of order $\\mathcal O (\\log(T/\\delta) / T^{\\frac{2p-2}{5p-3}})$. To my knowledge, this is the first high probability results under heavy tails for second-order smooth functions.\nAuthors also demonstrate the effect of initial burn-in period as discussed in Theorem 3 (and Theorem 6 for second order smoothness) through some experiments.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Prior work on high probability analysis of first-order methods usually rely on light-tail assumptions or almost-sure boundedness of gradient. This work specifically avoids those two assumptions, and achieve the same convergence rates with the in-expectation counterparts with heavy tailed gradients. Additionally, to the best of my knowledge, under heavy tails assumptions, they show high-probability convergence for second-order smooth functions.\nAuthors also show some complementary result on:\n-\tBehavior of the last iterates generated by Algorithm 1 and 2. \n-\tQuality of first-order estimate after burn-in period\n-\tConditions under which the monotonic descent in objective is achieved\n\nThese results both apply to first and second-order smooth settings. One small criticism I have is that both Algorithm 1 and 2 require a fixed time horizon, however, it doesn’t decrease the impact of these results in my opinion.\nI think the related work is adequately cited up to some relevant results from the convex realm. It might be helpful to motivate the removal of light-tails and almost-sure boundedness.\n\nThis paper is theoretically oriented, and authors provide theoretical analysis for their main theorems both in the main text and the appendix. To complete some of their theoretical findings, authors run benchmark experiments for the demonstration of burn-in idea. According to my interpretation, this burn-in period is also a parameter to be tuned, and if picked properly, it helps with the optimization process by bypassing the initial steps during which the algorithm might perform poorly. The authors are honest with their results and explain the implications of their findings with intuitions and examples, which strengthens their paper.\n\nI believe their results are significant from a theoretical point of view as this the first result of its kind for second-order smoothness. Theorem 3 and 6 provide a detailed view for the behavior of the gradient estimate, as well as the bound on full gradient when the algorithm runs beyond a time period called burn-in. As a complementary result, Corollary 4 and 7 validate the behavior of last iterate and reinforces the practice of using the last iterate generated by the algorithm.\nThe paper is well-organized and has a clear presentation of the problem setting and the main results. There are some typos and incomplete sentences, which should be corrected for the final version. Authors provide explanations for their algorithmic choices and explain what each theorem implies both theoretically and in practice.\n\nFinally, I have some particular question/comments regarding the paper:\n-\tIncomplete sentence at line 109\n-\tCould the authors please explain how they obtain the inequality right after line 173? Also, using the definitions of $\\mathcal T$ and $\\eta$, this quantity is upper bounded by $O ( \\log(T) / T^{\\frac{p-1}{3p-2}} )$, which implies a sublinearly decreasing bound. If what I computed is correct, this is an unexpected result.\n-\tIf I didn’t miss it, G-Lipschitz assumption does not exist in the Theorem 3 but in its proof.\n-\tIn the conclusion, it reads “Specifically, the last iterate has function value smaller than the function value of a critical point. ” I believe it should be $\\epsilon$-critical point, as otherwise it might imply that $w_T$ is not near a stationary point but has smaller function value than some stationary point.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors don't talk about the negative societal impact of their work, but given that this is a theoretical paper it is probably less likely that it would have any immediate negative societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "TpT3jtAN_W", "reviewer": "Reviewer_KPvy", "summary": "The paper contains novel strong theoretical results on finding stationary points for smooth non-convex stochastic problems with heavy-tailed noise. To achieve such results, the authors propose a new method based on normalized Momentum-SGD with clipping of stochastic gradients. The key part of the analysis is in an accurate application of Freedman's inequality for bounded martingale difference sequences with bounded conditional variances. The derived rates have $\\log(1/\\delta)$ dependence on the probability margin/confidence level $\\delta$ that is a desirable property for high-probability results; the dependence on the number of iterations $T$ coincide in the obtained rates is the same as in the state-of-the-art in-expectation convergence results under the same assumptions. Moreover, the paper contains an analysis under second-order smoothness and results for last-iterate convergence that partially explain the success of the \"warm-up\" stepsize schedules.\n\nAlthough the paper contains a number of mathematical inaccuracies and typos, I believe that all of them can be easily fixed (see the details below). If the authors address all the comments properly, I will increase my score to 8.", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths\n\n1. **The paper focuses on a poorly investigated but very important setup.** In contrast, to almost all previous high-probability results for smooth non-convex stochastic problems, the results derived in this paper hold even in the case when the variance of the stochastic gradient is unbounded, i.e., the authors assume that $\\mathbb{E}\\|\\nabla f(w,z)\\|_*^{\\mathfrak{p}} \\leq G^{\\mathfrak{p}}$ for some ${\\mathfrak{p}} \\in (1,2]$. When ${\\mathfrak{p}} < 2$ the noise in the stochastic gradient is heavy-tailed and this setup was never considered in the literature on high-probability convergence. However, obtaining tight high-probability guarantees in the heavy-tailed regime is an important problem. Moreover, for Theorems 5 and 6 there are no known analogs for the convergence in expectation.\n\n2. **Tightness of the obtained results.** The result of Theorem 2 establishes $O\\left(\\frac{\\log(T/\\delta)}{T^{\\frac{\\mathfrak{p}-1}{3\\mathfrak{p}-2}}}\\right)$ bound for the average of gradients norms after $T$ iterations with probability at least $1-\\delta$. This result is obtained under smoothenss of the objective $F$ and $\\mathbb{E}\\|\\nabla f(w,z)\\|_*^{\\mathfrak{p}} \\leq G^{\\mathfrak{p}}$ for some ${\\mathfrak{p}} \\in (1,2]$. In this settings, thereare no known high-probability guarantees in the literature. Moreover, the obtained result has good dependence on $\\delta$, and in terms of the dependence on $T$ it matches the \"in-expectation\" result from [42] that was proven to be optimal (see [42]).\n\n3. **Last-iterate guarantees and the connections with warm-up stepsize schedules.** The authors propose a simple but elegant way to obtain last-iterate convergence guarantees and formalize it as Theorem 3: the function value in the last point produced by the algorithm is not bigger than the function value at the point where the gradient is small enough with high probability. In contrast, almost all previous high-probability results for finding stationary points of non-convex functions establish upper bounds only for the average of norms or squared norms of the gradients along the trajectory. Moreover, Theorem 3 implies that after $\\mathcal{T} \\sim T^{\\frac{\\mathfrak{p}}{3\\mathfrak{p}-2}} = o(T)$ iterations momentum $m_t$ approximates the true gradient $\\nabla f(x_t)$ with accuracy $~ \\frac{1}{T^{\\frac{\\mathfrak{p}-1}{3\\mathfrak{p}-2}}}$ (which is quite good) with high probability. Intuitively, it implies that the updates become less noisy and more similar to the updates normalized GD that guarantees a decrease in the objective. This result partially explains the empirical success of the warm-up stepsize schedule.\n\n4. **First results under second-order smoothness in such generality.** The authors generalize the results of Theorem 2 and 3 to the case when the objective function has Lipschitz Hessian in addition to the Lipschitzness of the gradients. The derived rate of convergence is $O\\left(\\frac{\\log(T/\\delta)}{T^{\\frac{2\\mathfrak{p}-2}{5\\mathfrak{p}-3}}}\\right)$ that is slightly better than the one derived without second-order smoothness. There are no analogs of these results for the convergence in expectation.\n\n## Weaknesses\n\n1. **Proofs are hard to read.** I have checked all the proofs and noticed a number of mathematical inaccuracies, typos, and poorly explained derivations (see the comments below). Although all of them can be easily corrected, it is vital to fix these issues.\n\n2. **Boundedness of the gradients.** The assumption on the stochastic gradient implies boundedness of $\\|F(w)\\|_*$ for all $w$ that is quite strong. However, almost all high-probability results for finding stationary points in the non-convex case rely on this assumption. Therefore, this drawback is not that significant given the contribution of the paper.\n\n## Questions and Comments\n\n1. **line 24:** Even if the result is well-known the reference is required.\n\n2. **line 50, $\\tilde{O}()$:** The authors should either show the dependence on $\\delta$ or write explicitly that $\\tilde{O}()$ hides **polylogarithmical factors** of $1/(\\varepsilon\\delta)$.\n\n3. **lines 89-90, about $\\ell_p$-norm:** The authors should either prove this result somewhere in the appendix or add the reference to the source where it is proven.\n\n4. **line 92, \"In this paper, we will exclusively consider the case $p = 2$ and $C\\ge 1$\":** this statement should be either highlighted or repeated explicitly in the statements of the main result. The reason why I ask the authors to do this is that in the section on concentration bounds main results for Banach spaces are proven for general $p$. It is misleading when in the proof of Theorem 2 concentration inequalities are applied with $p=2$.\n\n5. **line 95:** Appendix D $\\to$ Appendix E.\n\n6. **line 96, incorrect definition of smoothness:** one should use primal norm in the right-hand side of the inequality.\n\n7. **line 98, \"We assume that $\\nabla f(w,z)$ satisfies ...\":** In fact, the authors assume boundedness of the $\\mathfrak{p}$th moment of the stochastic gradient by $G^{\\mathfrak{p}}$, not the central $\\mathfrak{p}$th moment.\n\n8. **line 109:** Banac $\\to$ Banach, In $\\to$ in\n\n9. **line 110, the last word:** normalization?\n\n10. **line 111, Lemma 1:** the statement is strange, some word was omitted. Moreover, sometimes authors use $w$, sometimes -- $\\vec{w}$. The notation should be unified.\n\n11. **Algorithm 1, Input:** one should write that $\\beta = 1-\\alpha$ or replace $\\alpha$ by $\\beta$. Next, there is a typo in the update rule for $w_{t+1}$: one should have $w_{t+1} = w_t - \\eta d(m_t)$.\n\n12. **line 120** matches $\\to$ matches up to logarithmical factors.\n\n13. **line 129:** must have $\\to$ must have with probability at least $1-\\delta$.\n\n14. **Formulas from Theorem 2, 3 and almost all other formulas in the paper:** there are a lot of mistakes with the punctuation after the formulas. That is, if the sentence ends with a formula, it is needed to add a full stop after the formula. Please, fix it everywhere in the paper.\n\n15. **line 147, \"By smoothness we have ...\"**: it should be primal norm in the right-hand side.\n\n16. **line 149, the inequality below:** $(1-\\alpha)^t\\hat\\epsilon_0 \\to (1-\\alpha)^{t+1}\\hat\\epsilon_0$. Please fix it here and everywhere below. The same problem is in Theorem 3, 5, and 6 (as well as incorrect indexes).\n\n17. **line 177, the power in the denominator:** $\\frac{2\\mathfrak{p}-1}{3\\mathfrak{p}-2} \\to \\frac{\\mathfrak{p}-1}{3\\mathfrak{p}-2}$.\n\n18. **Theorem 5, definition of $\\eta$**: $a \\to s$.\n\n19. **Theorem 6**: a number of mistakes in the statement, that should be fixed (incorrect bounds and rates).\n\n20. **lines 440-441 and the formula behind:** please, do not skip the steps in the derivations. It becomes hard to read the proof. Moreover, there are too many typos in these 6 lines: incorrect definition of $\\epsilon_t$ (should be $x_t$ instead of $w_t$), in the formula for $m_{t+1}$ it should be $g_{t+1}^{\\text{clip}}$ instead of $g_{t}^{\\text{clip}}$, wrong sign in front of $\\nabla^2 F(w_{t+1})$ and superfluous right round bracket after $\\epsilon_{t+1}$ in the last line of the formula.\n\n21. **line 442, the formulas behind:** the indexes in the sums are incorrect.\n\n22. **line 443, $\\|S(a,b)\\|_{*} \\le \\frac{\\rho}{2}\\|a-b\\|^2$**: one should add either a proof or a reference.\n\n23. **line 443, formulas behind**: the authors skipped a lot of steps of derivations of these inequalities. It is inconvenient for readers. Moreover, the constants are not tight: one can get 2 times smaller factor in front of the sum.\n\n24. **lines 444-445 and inequality behind, \"by exactly the same argument as in Theorem 2\"**: why numerical constants are different then? The authors should write a complete proof or fix the issue. In the next formula, the numerical constants differ again...\n\n25. **Lemma 8:** It should be mentioned explicitly in the statement that this result is only for Hilbert spaces.\n\n26. **line 483:** Lemma 14 $\\to$ Lemma 12\n\n27. **line 484:** $\\hat\\mu_t \\to \\hat X_t$.\n\n28. **line 496, application of Lemma 14**: In Lemma 14 one needs to have bounded variance while you show only the boundedness of the central $\\mathfrak{p}$th moments. Moreover, all steps in obtaining the inequality after line 497 should be added (it is not just an application of Lemma 14). Although it can be derived, I strongly encourage the authors to add all missing derivations to the proofs.\n\n29. **Appendix C:** Feedman $\\to$ Freedman\n\n30. **lines 509 - 514, definition of $s_t$:** this definition implies that $s_1 = 0$, and, as a consequence, $s_t = 0$ for all $t$. This is because for $t=1$ the sum $\\sum_{i=1}^{t-1}s_i = 0$ since it does not contain summands (it is a common convention). Therefore, formally, all the proofs in the paper are incorrect. But this problem can be easily corrected by changing the definition of $s_1$. Moreover, in the definition of $s_t$ it should be an inner product of $X_t$ and the gradient. In the current form, it is incorrect.\n\n31. **inequality below line 518:** sums of $s_t$ and $X_t$ should be swapped.\n\n32. **line 524, \"the function $x \\to x^p$ is convex...\":** this assumes that $p \\geq 1$. This assumption was never mentioned explicitly.\n\n33. **lines 527-529 and inequality behind:** again, too many typos and too many derivations are omitted. First of all, from the explanation, it does not follow what is written in the inequality. To get this, the authors should change the definition of $b+c$ (take maximum over $i\\leq T-1$ and upper bound it after, in the second line of the inequality). Next, in the second line of the inequality, it should be summation from $i=1$ to $T-1$ for $s_i$. To conclude, I repeat once again that the proofs are not reader-friendly and should be filled with detailed explanations. All the inaccuracies should be corrected.\n\n34. **Lemma 11:** the reference is missing. It is not a classical form of the Freedman's inequality (for the classical one, see Theorem 1.1 from Tropp, J. (2011). Freedman's inequality for matrix martingales. Electronic Communications in Probability, 16, 262-270.).\n\n35. **line 542:** $1-\\delta \\to 1-3\\delta$\n\n36. **line 544, \"$|s_t| = \\|X_t\\|$\"**: it is not true. The authors should add a simple derivation via Cauchy-Schwarz showing that $|s_t| < \\|X_t\\|$.\n\n37. **line 545:** Why is it needed to use $2\\delta$ instead of $\\delta$ here?\n\n38. **inequality after line 550, $R^2$ under the square root:** it should be $R^4$.\n\n39. **inequality after line 551:** $\\sigma_t^2 \\to \\|X_t\\|^2$ in the first line, $\\sigma_t \\to \\sigma_t^2$ in the last line.\n\n40. **line 560:** Banach space $\\to$ $(p,C)$-smooth Banach space.\n\n41. **line 562:** $1-\\delta \\to 1-3\\delta$\n\n42. **expression after line 563:** What is the statement here? Is it the right-hand side of the previous inequality under additional assumptions on $p$ and $C$?\n\n43. **line 564:** when you apply Jensen's inequality you use $p \\leq 2$ implicitly. So, it should be added in the statement of the lemma even for the first bound.\n\n44. **Proof of Proposition 16:** the proof is incorrect (though the statement of the proposition is correct).\n\n## Comment after rebuttal\nI thank the authors for their response. Since the authors properly addressed my main criticism, I am increasing my score to 8 in the hope that the authors will apply all necessary corrections in the final version of the paper.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately addressed the limitations and potential negative societal impact of their work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "12 hours"}, "extra_text": {}}, {"review_id": "FWMh88dtDv", "reviewer": "Reviewer_5L1D", "summary": "The paper studies non-convex stochastic optimization using normalized SGD with clipping and momentum. They obtain convergence of the average to critical points in high-probability with best-known rates for smooth losses when the gradients have heavy tails, that is, with finite p-th moments for some p\\in(1,2], and convergence of the last iterate is also obtained. The paper also considers normalized implicit gradient with transport with clipping in the case of second-order smooth losses, which is new and they obtain convergence results for the average as well as the last iterate. ", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is original and the main results are clear. It is known that the analysis for gradients with heavy tails can be more challenging, and it is nice that the authors can carry out the analysis when the gradient has finite p-th moments for p possibly smaller than 2. The normalized implicit gradient with transport with clipping for second-oder smooth losses is also new. \n\nWhen I'm reading your proof of Theorem 2 and Theorem 3 for example, it seems you are able to keep track of the constants, and I'm wondering if your results of Theorem 2 and Theorem 3 can be made non-asymptotic with explicit constants instead of big O? That can strengthen the contribution of the paper.\n\nAlso, the paper seems to be written in a rush, and there are numerous typos and places with confusions, and I suggest the authors check the paper carefully and paper would be in a better shape if the typos and minor mistakes are corrected. For example,\n\n(1) Page 1. Line 24. for some \\sigma. I don't understand this. What's the relation between \\sigma and G?\n\n(2) Page 2. Line 45. What do you mean by convergence rate itself is a random variable? For SGD, even though you have randomness, using criteria like Wasserstein metric or others, there are many results in the literature with deterministic convergence rates.\n\n(3) Page 3. In the last paragraph of Section 1.1. What is \\nabla F(w,z)? Should it be \\nabla F(w) or?\n\n(4) Page 3. In the first paragraph in Section 2, Banac should be Banach.\n\n(5) In Theorem 2, is G the same as the G on page 1?\n\n(6)  Page 4. In the paragraph before Proof of Theorem 2, is burnin-in period a typo?\n\n(7) Page 6. In the second display of Proof of Theorem 3, there is an extra ) for \\log(G(1-\\apha)^{t}))\n\n(8) You wrote in the proof of Theorem 3, \"observe that since F is G-Lipschitz\". Where did you assume F is G-Lipschitz?\n\n(9) In the statement of Theorem 5, you wrote \"for all x for some G\". Should it be \"for all w for some G\" instead?\n\n(10) In references, some words should be capitalized. For example, sgd should be SGD in [9].", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I could not find out where the authors mention limitations and societal impact in their paper. One limitation I can see is that in the main results, they seem to be asymptotic instead of non-asymptotic. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2.5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper considers non-convex SO using first-order algorithms for which the gradient estimates may have heavy tail distributions. Using gradient clipping normalized GD with momentum is shown to convergence to near critical points in high-probability with best-known rates for smooth losses when the gradients only have bounded p'th moments for some p between 1 and 2. Next high-probability bounds are obtained for the case of second-order smooth losses.

The reviewers are unanimously strongly in favor of accepting the paper and I concur. This seems like a nice contribution.
Decision: Accept (Oral)