High Probability Convergence Bounds for Non-convex Stochastic Gradient Descent with Sub-Weibull Noise

TL;DR

This paper establishes high probability convergence bounds for stochastic gradient descent in non-convex settings with sub-Weibull noise, extending theoretical guarantees beyond sub-Gaussian assumptions and including a novel concentration inequality.

Contribution

It introduces high probability convergence bounds for non-convex SGD under sub-Weibull noise and develops a new concentration inequality for heavy-tailed martingale differences.

Findings

Convergence bounds hold without convexity assumptions.

Extended concentration inequality for sub-Weibull martingale differences.

A post-processing method for selecting a convergent iterate.

Abstract

Stochastic gradient descent is one of the most common iterative algorithms used in machine learning and its convergence analysis is a rich area of research. Understanding its convergence properties can help inform what modifications of it to use in different settings. However, most theoretical results either assume convexity or only provide convergence results in mean. This paper, on the other hand, proves convergence bounds in high probability without assuming convexity. Assuming strong smoothness, we prove high probability convergence bounds in two settings: (1) assuming the Polyak-{\L}ojasiewicz inequality and norm sub-Gaussian gradient noise and (2) assuming norm sub-Weibull gradient noise. In the second setting, as an intermediate step to proving convergence, we prove a sub-Weibull martingale difference sequence self-normalized concentration inequality of independent interest. It extends Freedman-type concentration beyond the sub-exponential threshold to heavier-tailed martingale difference sequences. We also provide a post-processing method that picks a single iterate with a provable convergence guarantee as opposed to the usual bound for the unknown best iterate. Our convergence result for sub-Weibull noise extends the regime where stochastic gradient descent has equal or better convergence guarantees than stochastic gradient descent with modifications such as clipping, momentum, and normalization.