High Probability Analysis for Non-Convex Stochastic Optimization with Clipping

TL;DR

This paper provides a high probability theoretical analysis of gradient clipping in non-convex stochastic optimization, addressing heavy-tailed gradient behaviors and deriving bounds for optimization and generalization.

Contribution

It offers the first high probability analysis for stochastic optimization with gradient clipping under weak heavy-tailed assumptions, covering multiple algorithms.

Findings

Derived optimization bounds with gradient clipping

Established generalization bounds under heavy-tailed gradients

Applicable to SGD, momentum, and adaptive methods

Abstract

Gradient clipping is a commonly used technique to stabilize the training process of neural networks. A growing body of studies has shown that gradient clipping is a promising technique for dealing with the heavy-tailed behavior that emerged in stochastic optimization as well. While gradient clipping is significant, its theoretical guarantees are scarce. Most theoretical guarantees only provide an in-expectation analysis and only focus on optimization performance. In this paper, we provide high probability analysis in the non-convex setting and derive the optimization bound and the generalization bound simultaneously for popular stochastic optimization algorithms with gradient clipping, including stochastic gradient descent and its variants of momentum and adaptive stepsizes. With the gradient clipping, we study a heavy-tailed assumption that the gradients only have bounded $\alpha$-th moments for some $\alpha \in (1, 2]$, which is much weaker than the standard bounded second-moment assumption. Overall, our study provides a relatively complete picture for the theoretical guarantee of stochastic optimization algorithms with clipping.