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