High Probability Analysis for Non-Convex Stochastic Optimization with Clipping
This work addresses a theoretical gap for researchers and practitioners using gradient clipping to stabilize training in neural networks, offering rigorous guarantees under weaker assumptions, though it is incremental as it builds on existing clipping methods.
The paper tackles the lack of theoretical guarantees for gradient clipping in non-convex stochastic optimization by providing high-probability analysis for optimization and generalization bounds under heavy-tailed gradient assumptions, showing that clipping enables convergence with bounded α-th moments (α ∈ (1, 2]).
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 $α$-th moments for some $α\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.