High-probability Bounds for Non-Convex Stochastic Optimization with Heavy Tails

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 ] ≤ 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 . 1.1 Setup and assumptions . 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 . 2 Normalized SGD with Momentum . 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 .

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