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

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.

References in corpus (11)

• Language Models are Few-Shot Learners
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• On the Convergence of Adam and Beyond
• Large Batch Training of Convolutional Networks
• AdaGrad stepsizes: Sharp convergence over nonconvex landscapes
• Why gradient clipping accelerates training: A theoretical justification for adaptivity
• A Tail-Index Analysis of Stochastic Gradient Noise in Deep Neural Networks
• On the Heavy-Tailed Theory of Stochastic Gradient Descent for Deep Neural Networks
• A High Probability Analysis of Adaptive SGD with Momentum
• Reducing the variance in online optimization by transporting past gradients
• Convergence Rates of Stochastic Gradient Descent under Infinite Noise Variance