High-probability bounds for Non-Convex Stochastic Optimization with Heavy Tails
Year: 2021

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 ∈ (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 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 whenever the current iterate is not a critical point, which provides intuition behind the popular practice of learning rate “warm-up” and also yields a last-iterate guarantee. 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 as optimizing the following stochastic optimization problem: min ~ w F (~ w) = E z [f(~ w, z)] where f(~ w, z) is some loss function taking model parameters ~ w and a random example z, where z has some distribution Pz . Then SGD consists of the update: ~ wt+1 = ~ wt − η∇f(~ wt, zt) where z1, . . . , zT are i.i.d. random variables distributed according to Pz , and η is some real number called the learning rate. Assuming that F is L-smooth and that∇f(~ w, z) satisfies E[‖∇f(~ w, z)‖2] ≤ G for some G for all ~ w, SGD satisfies [13]:

heading: 1 SGD and Heavy Tails
text: 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 as optimizing the following stochastic optimization problem:
min ~w F (~w) = E z [f(~w, z)]
where f(~w, z) is some loss function taking model parameters ~w and a random example z, where z has some distribution Pz . Then SGD consists of the update:
~wt+1 = ~wt − η∇f(~wt, zt) where z1, . . . , zT are i.i.d. random variables distributed according to Pz , and η is some real number called the learning rate. Assuming that F is L-smooth and that∇f(~w, z) satisfies E[‖∇f(~w, z)‖2] ≤ G2 for some G for all ~w, SGD satisfies [13]:
1 T E
[ T∑
t=1
‖∇F (~wt)‖2 ] ≤ O ( L∆
T + G2√ T
)
where ∆ = F (~w1)− inf F (~w). This implies that SGD takes O( −4) iterations to find an -critical point, which is a point ~w such that ‖∇F (~w)‖ ≤ [13]. When the stochastic gradients are further assumed to have light sub-gaussian tails, [21] shows that SGD with momentum can achieve this same result in high probability as well.
35th Conference on Neural Information Processing Systems (NeurIPS 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] or various versions of convexity [36].
A standard tool when working with heavy-tailed random variables is truncation, or clipping. Succinctly, if X is a random variable such that E[‖X‖2] =∞, X itself may be poorly behaved due to the heavy tail, but a truncated value Xclip = X|X| min(|X|, τ) 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 SGD combined with clipping can find an -critical point in expectation in O ( − 3p−2 p−1 ) iterations.
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 -critical point in Õ ( − 3p−2 p−1 ) iterations with
high probability.
Second, the vast majority of first-order optimization methods for stochastic non-convex optimization consider exclusively the standard L2 or Hilbert-space norm. Such analyses can take advantage of the fact that E[〈∇F (x), ~g〉] = ‖∇F (x)‖2 when ~g is an unbiased estimate of ∇F (x). 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 ~g (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. It turns out the key to operating with different norms is the popular practice of augmenting SGD iterates with momentum. When using momentum-based updates we do not rely on the same unbiasedness argument, and so the techniques can be extended to more general norms. This provides more concrete evidence of the theoretical advantages of momentum.
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 second-order 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 -critical point in O( 3.5) iterations using a 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 an -critical point in Õ ( − 5p−3 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 an intriguing connection to the popular warm-up learning rate schedule [16, 15, 10], 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 to 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 to the use of layer-normalization, 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 2 and 3, we 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.

heading: 1.1 Setup and assumptions
text: We assume that the objective F (~w) = E[f(~w, z)] is a Frechet-differentiable function mapping a Banach space B to the reals R. Let ‖ · ‖ 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 a dual vector v ∈ B? to a vector w ∈ B. We assume that ‖ · ‖p? is Frechet differentiable and satisfies for some p, C for all x, y ∈ B:
‖x+ y‖p? ≤ ‖x‖p? + 〈∇‖x‖p?, y〉+ C‖y‖p? (1) Such a Banach space is called (p, C)-smooth. Note that if ‖ · ‖ is the standard p norm in Rn for p ∈ (1, 2], then ‖ · ‖ is the corresponding q norm, and the dual space is (2, 1p−1 )-smooth. In the common case that B is a Hilbert space (e.g. Rn with the euclidean norm), then B = B?, ‖ · ‖ = ‖ · ‖? and p = 2 and C = 1. In this paper, we will exclusively consider the case p = 2 and C ≥ 1. Finally, given any v ∈ B?, we will write d(v) ∈ B to indicate a unit vector satisfying both ‖d(v)‖ = 1 and 〈v, d(v)〉 = ‖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 ‖∇F (~w) − ∇F (~x)‖? ≤ L‖~w − ~z‖?. In Section 3, we will also assume that F is ρ-second-order smooth, which means that ‖(∇2F (~w) − ∇2F (~x))v‖? ≤ ρ‖~w − ~x‖‖v‖ for all ~w, ~x, v ∈ B. We assume that ∇f(~w, z) satisfies both E[‖∇f(~w, z) − ∇F (~w)‖p?] ≤ Gp and E[‖∇f(~w, z)‖p?] ≤ Gp for some G and some p ∈ (1, 2]. We will design algorithms such that after T stochastic gradient evaluations, we can output a point ~w such that ‖∇F (~w)‖? is small with high probability.

heading: 2 Normalized SGD with Momentum
text: In this section, we describe our algorithms for finding critical points with high probability in nonconvex objectives. Our analysis makes use of gradient clipping, momentum, and normalized updates. Clipping is a standard technique for mitigating the effect of heavy tails, essentially by throwing out outlier data. Momentum is useful as it intuitively averages together many clipped gradients, resulting in a 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, as 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 : B → R be a Frechet-differentiable L-Smooth function from a Banach space B to the reals. Let w ∈ B. Let g? ∈ B? and let g ∈ B be a unit-vector satisfying 〈g?, g〉 = ‖g‖?. Define w′ = w − ηg. Define = g? −∇F (w). Then:
F (w′) ≤ F (w)− η‖∇F (w)‖? + 2η‖ ‖? + Lη2
2
Theorem 2. Suppose Ez[‖∇f(~w, z)‖p?] ≤ Gp for all ~x for some G. Suppose F is L-smooth. Set β = 1 − α, α = b
T p 3p−2 and η = s T 2p−1 3p−2 for arbitrary constant b and s satisfying α ≤ 1 and set τ = G/α1/p. Then with probability at least 1− δ:
1
T
T∑
t=1
‖∇F (~wt)‖? ≤ O ( log(T/δ)
T p−1 3p−2
)
where the big-Oh hides constant that depend on L, G b, s, C and F (~w1)− F (~wT+1), but not T or δ.
Algorithm 1 Normalized SGD with Clipping and Momentum Input: Initial Point ~w1, learning rate η, momentum parameter β, clipping parameter τ , time horizon T : Set ~m0 = 0. for t = 1 . . . T do
Sample zt ∼ Pz . Set ~gclipt = ∇f(~wt,zt) ‖∇f(~wt,zt)‖? min(τ, ‖∇f(~w1, zt)‖?). Set ~mt = β ~mt−1 + (1− β)~gclipt . Set ~wt+1 = ~wt − ηd(~mt).
end for
This Theorem provides an analog of the standard in-expectation result for non-convex optimization, but now the result holds with high-probability. Note that this matches the in-expectation convergence rates from Zhang et al. [43] up to logarithmic factors. The next Theorem takes this theme one step further: it shows that after a 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, log(3T/δ)) + 4C1/2 √ max(1, log(3T/δ)) + 1
Z = sL
b +GKb
p−1 p
Then, with probability at least 1− δ, for all t satisfying:
t ≥ T = T p 3p−2
b ( p− 1 3p− 2 log(T ) + log(G)− log(Z) )
we have
‖~mt −∇F (~wt)‖? ≤ 2Z
T p−1 3p−2
Further, so long as
‖~mt‖? ≥ 2 ( 6Z
T p−1 3p−2
+ Ls
2T 2p−1 3p−2
) (2)
we have F (~wt+1) < F (~wt)− η2‖~mt‖?. Moreover, whenever (2) is not satisfied for t ≥ T , we must have
‖∇F (~wt)‖? ≤ 14Z
T p−1 3p−2
+ Ls
T 2p−1 3p−2
Finally, if t is the iteration with smallest value of ‖~mt‖? such that t ≥ T , we have:
‖∇F (~wt)‖? ≤ O ( log(T/δ)
T p−1 3p−2
)
Theorem 3 shows that after a small burn-in period T = o(T ), ~mt is very likely to be an extremely good estimate of the true gradient ∇F (~xt). 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 a “warm-up” component, in which the learning rate actually increases during the first iterations. Our result seems to hint at a 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 Õ(T p 3p−2 ) iterates, we can simply return the point ~xt with the smallest value of ‖~mt‖. Intuitively, this works because the objective value can only increase by a total of O(T p 3p−2 η) = O(T 1−p 3p−2 ) = O(1) during the burn-in period so even without a 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, ~mt is always very close to∇F (~wt), so we make steady progress until a critical point is encountered.
Proof of Theorem 2. Define ̂t = ~mt −∇F (~wt). and t = ~gclipt −∇F (~wt). Also, define S(a, b) = ∇F (a)−∇F (b). By smoothness, we have ‖S(~wt, ~wt+1)‖? ≤ L‖~wt − ~wt+1‖? From the definition of ~mt+1, we have the following recursive formulation for any t ≥ 1:
~mt+1 = (1− α)(∇F (~wt) + ̂t) + α~g clipt+1 = ∇F (~wt+1) + (1− α)(S(~wt, ~wt+1) + ̂t) + α t+1
̂t+1 = (1− α)S(~wt, ~wt+1) + (1− α)̂t + α t+1 (3) Further, by defining ~m0 = 0 and ~w0 = ~w1 so that ̂0 = −∇F (~w1), we have that (3) holds for t = 0 as well. Now, we unravel the recursion for all iterations:
̂t+1 = (1− α)t+1̂0 + α t∑
t′=0
(1− α)t′ t+1−t′ + (1− α) t∑
t′=0
(1− α)t′S(~wt−t′ , ~wt+1−t′)
Next, take the magnitude of both sides and use triangle inequality (using the fact that ‖S(~wt, ~wt+1)‖? ≤ Lη):
‖̂t+1‖? ≤ (1− α)t+1‖̂0‖? + α ∥∥∥∥∥ t∑
t′=0 (1− α)t′ t+1−t′ ∥∥∥∥∥ ? + (1− α)ηL t∑ t′=0 (1− α)t′
≤ (1− α)t+1G+ (1− α)ηL α + α ∥∥∥∥∥ t∑
t′=0 (1− α)t′ t+1−t′ ∥∥∥∥∥ ?
where in the second line we have used ~m0 = 0 and ‖∇F (~w0)‖? ≤ G, which follows since ‖∇F (~w0)‖? ≤ Ez[‖∇f(~w0, z)‖p?]1/p = G. Now, we invoke Lemma 9, a concentration inequality in Banach space (see appendix B). This implies that with probability at least 1− δ/T : ∥∥∥∥∥ t∑
t′=0 (1− α)t′ t+1−t′ ∥∥∥∥∥ ? ≤ 10Cτ max(1, log(3T/δ)) + t∑ t′=0 (1− α)t′Gp τp−1
+ 4 ( C t∑
t′=0
(1− α)2t′Gpτ2−p )1/2√ max(1, log(3T/δ))
≤ 10Cτ max(1, log(3T/δ)) + G p
ατp−1 +
4 ( CGpτ2−p )1/2√ max(1, log(3T/δ))
α1/2
Therefore, with probability at least 1− δ/T :
‖̂t+1‖? ≤ (1− α)t+1G+ (1− α)ηL
α + 10Cατ max(1, log(3T/δ))
+ 4 ( αCGpτ2−p )1/2√ max(1, log(3T/δ)) + Gp
τp−1
Now, with τ = G α1/p , and Dδ = max(1, log(3T/δ)) this becomes:
≤ (1− α)t+1G+ (1− α)ηL α
+Gα1−1/p [ 10CDδ + 4C 1/p √ Dδ + 1 ]
= (1− α)t+1G+ (1− α)ηL α +Gα1−1/pK (4)
where we define the constant K = 10CDδ + 4C1/2 √ Dδ + 1.
Now, we have from Lemma 1: T∑
t=1
‖∇F (~wt)‖? ≤ F (~w1)− F (~wT ) η + LTη 2 + 2
T∑
t=1
‖̂t‖?
so with probability at least 1− δ:
≤ F (~w1)− F (~wT ) η + LTη 2 + 2
T∑
t=1
(1− α)tG+ (1− α)ηL α +Gα1−1/pK
≤ F (~w1)− F (~wT ) η + LTη 2 + 2G(1− α) α + 2(1− α)ηTL α + 2GTα1−1/pK
Now set α = b T p 3p−2 and η = s T 2p−1 3p−2 . This yields:
1
T
T∑
t=1
‖∇F (~wt)‖ ≤ O ( T 1−p 3p−2 log(T/δ) )
The proof of Theorem 3 follows essentially the same idea as that of Theorem 2. However, this time we use the explicit bound on ‖̂t‖? and observe that after the burn-in period, the contribution from ‖ 1‖? has decreased exponentially to be insignificant, and so ~mt will be a very high-quality estimate of∇F (~xt) at every single iteration.
Proof of Theorem 3. Use the settings η = s T 2p−1 3p−2 and α = b T p 3p−2 and notice that from equation (4) that with probability 1− δ, for all t:
‖̂t+1‖? ≤ (1− α)tG+ sL
bT p−1 3p−2
+ GKb
p−1 p
T p−1 3p−2 = (1− α)tG+ Z T p−1 3p−2
Further, we have
log(1− α) ≤ −α log ( G(1− α)t ) ≤ −tα+ log(G)
Therefore, if t ≥ T , then we have with probability at least 1− δ:
‖̂t+1‖? ≤ 2Z
T p−1 3p−2
This implies the bound on ‖∇F (~wt)‖? when ‖~mt‖? is smaller than the threshold described in the Theorem statement. Next, again from Lemma 1,
F (~wt+1) ≤ F (~wt)− η‖∇F (~wt)‖? + 2η‖̂t‖? + Lη2
2
≤ F (~wt)− η‖~mt‖? + 3η‖̂t‖? + Lη2
2
≤ F (~wt)− η‖~mt‖? + η ( 6Z
T p−1 3p−2
+ Ls
2T 2p−1 3p−2
)
from which we can conclude that if ‖~mt‖? ≥ 2 ( 6Z
T p−1 3p−2 + Ls 2T 2p−1 3p−2
) , F (~wt+1) ≤ F (~wt)− η2‖~mt‖?.
Now, for the final part of the Theorem, observe that since F is G-Lipschitz, we must have that
F (~wT )− F (~w1) ≤ GT η = O(log(T )) so that F (~wT )− F (~wT+1) ≤ O(F (~w1)− F (~wT+1) + log(T )). Now consider two cases, either
min t≥T ‖~mt‖? ≤ 2
( 6Z
T p−1 3p−2
+ Ls
2T 2p−1 3p−2
)
or not. In the first case, by our bound on ‖̂t‖?, the desired bound on ‖∇F (~wt)‖? follows. In the latter case, we have
F (~wT+1) ≤ F (~wT )− η
2
T∑ t=T ‖~mt‖?
So that
min t≥T ‖~mt‖? ≤ O
( F (~w1)− F (~wT+1) + log(T )
Tη
) = Õ ( 1
T 2p−1 3p−2
)
And so again the result follows from our bound on ‖̂t‖?.
In addition to providing a method for identifying critical points, Theorem 3 also provides an intuitively desirable guarantee about the last-iterate of the algorithm ~wT . Specifically, one can easily check that the following Corollary:
Corollary 4. Under the notation and assumptions of Theorem 3, with probability at least 1− δ, there exists some ŵ such that both of the following inequalities hold:
‖∇F (ŵ)‖? ≤ O ( log(T/δ)
T p−1 3p−2
)
F (~wT ) ≤ F (ŵ) where ~wT is the last iterate of Algorithm 1
Proof. Set t̂ to be the last index such that ~wt̂ does not satisfy (2). Set ŵ = ~wt̂. Then Theorem 3 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 us that our objective value is at least never worse than it would have been if had indeed searched specifically for a critical point. Moreover, by providing a guarantee about ~wT , we have a closer match to practical use-cases: many of the theoretical analyses of non-convex stochastic gradient methods [13, 38, 20] only show that a randomly selected iterate has a small gradient norm. In contrast, Corollary 4 provides a non-asymptotic guarantee for the final iterate, which is the iterate most likely to be deployed after training in practice.

heading: 3 Second-Order Smooth Losses
text: In this section, we provide our extension to losses that are second order smooth. The overall method is 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:
~mt = β ~mt−1 + (1− β)∇f ( ~wt +
β(~wt − ~wt−1) 1− β , zt
)
By evaluating the gradient at the “shifted” point ~wt + β(~wt−~wt−1)
1−β and performing a second-order Taylor expansion enabled by the second-order smoothness assumption, it is possible to show that ~mt is a less-biased estimate of ∇F (~wt) 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 to Theorems 2 and 3.
Theorem 5. Suppose Ez[‖∇f(~w, z)‖p] ≤ Gp for all ~w for some G. Suppose F is L-smooth and ρ-second-order smooth. Set β = 1− α, α = b
T 2p 5p−3 and η = a T 3p−1 5p−3 for arbitrary constants b and c
satisfying α ≤ 1, and set τ = G α1/p . Then with probability at least 1− δ:
1
T
T∑
t=1
‖∇F (~wt)‖? ≤ O ( log(T/δ)
T 2p−2 5p−3
)
where the big-Oh hides constant that depend on L, ρ, G, b, s, C, and F (~w1)− F (~wT+1), but not T or δ.
We also have a direct analog of Theorem 3
Algorithm 2 NIGT with Clipping Input: Initial Point ~w1, learning rate η, momentum parameter β, clipping parameter τ , time horizon T : Set ~m0 = 0. for t = 1 . . . T do
Sample zt ∼ Pz . ~xt = ~wt + β(~wt−~wt−1) 1−β Set ~gclipt = ∇f(~xt,zt) ‖∇f(~xt,zt)‖? min(τ, ‖∇f(~x1, zt)‖?).
Set ~mt = β ~mt−1 + (1− β)~gclipt . Set ~wt+1 = ~wt − ηd(~mt).
end for
Theorem 6. Under the assumptions of Theorem 5, define the constants:
K = 10 max(1, log(3T/δ)) + 4 √ C max(1, log(3T/δ)) + 1
Z = ρs2
b2 +GKb
p−1 p
Then, with probability at least 1− δ, if
t ≥ T = T 2p 5p−3
b ( 2p− 2 5p− 3 log(T ) + log(G)− log(Z) )
Then we have
‖~mt −∇F (~xt)‖ ≤ 2Z
T 2p−2 5p−3
And so long as
‖~mt‖? ≥ 2 ( 6Z
T 2p−2 5p−3
+ Ls
2T 3p−1 5p−3
) (5)
we have F (~wt+1) < F (~wt)− η2‖~mt‖. Moreover, if (5) is ever not satisfied, we must have
‖∇F (~wt)‖? ≤ 14Z
T 2p−2 5p−3
+ Ls
T 3p−1 5p−3
Finally, if t is the iteration with smallest value of ‖~mt‖ such that t ≥ T , we have with probability at least 1− δ:
‖∇F (~wt)‖? ≤ O ( log(T/δ)
T 2p−2 5p−3
)
Finally, we 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 1− δ, there exists some ŵ such that both of the following inequalities hold:
‖∇F (ŵ)‖? ≤ O ( log(T/δ)
T 2p−2 5p−2
)
F (~wT ) ≤ F (ŵ)
where ~wT is the last iterate of Algorithm 2

heading: 4 Experiments
text: Our theoretical analysis suggests that there is a short “burn-in” phase in our algorithms during which the momentum estimate ~mt may be far from the true gradient∇F (~wt), and after this period, ~mt will always be close to ∇F (~wt). As a 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. This suggests that one might set the learning rate η = 0 throughout the burn-in phase and only begin taking steps after ~mt has accumulated enough gradients to be come a good estimate. This procedure is reminiscent of the common practice of learning rate “warm-up”, in which the learning rate starts at 0 and is linearly increased to some desired value during the initial phase of training, and then later decreased. Thus, we sought to answer how much of the success of the warm-up schedule can be attributed to burn-in.
To answer this question, we conducted an experimental study on 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 to 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 η0 of 0.3 and batch size of 512. We came up with this learning rate by performing grid search η0 ∈ [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, we again ran a grid search η0 ∈ [0.1, 0.2, 0.3, 0.4, 0.5]. Our warmup baseline employs the standard practice of linear warm-up for 3125 steps and polynomical decay of ηt = η0 ∗ (1 − tT ) for the rest of the steps. For a 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, instead of linear warm-up, we set the base learning rate η0 to 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. 0 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, we 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 et al. [40]. Most importantly, we train all our models using the base learning rate η0 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 ηt = η0 ∗ (1− tT ) for the rest of the steps. Again, for a 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 in 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 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 1 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.

heading: 5 Conclusion
text: We have provided algorithms based on normalized SGD with momentum and gradient clipping that provide high probability-rates for finding critical points of non-convex objectives with heavy-tailed gradients. We also show that as a by-product of our high-probability 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 a critical point.
Finally, we identify a 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 an interesting question to precisely identify what other advantages are provided by the increasing schedule used by warm-up during this phase.
Acknowledgments: The authors thank the anonymous reviewers for many helpful suggestions. AC was a visiting researcher at Google when this work was completed.

Review summaries:
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. 

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.  Although 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.

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]$, -	They 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. -	Under 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. Authors also demonstrate the effect of initial burn-in period as discussed in Theorem 3 (and Theorem 6 for second order smoothness) through some experiments.