Stability and Generalization of Nonconvex Optimization with Heavy-Tailed Noise
Hongxu Chen, Ke Wei, Xiaoming Yuan, Luo Luo
TL;DR
The paper addresses generalization in nonconvex stochastic optimization under heavy-tailed gradient noise by developing a truncation-based stability framework under the $p$-BCM condition with $p\in(1,2]$. It couples algorithmic stability in gradients with a ghost-sample argument to bound the population gradient error by $4\epsilon + C_p\sigma_p n^{-(p-1)/p}$, where $\epsilon$ is a stability parameter and $C_p$ a $p$-dependent constant. The authors then apply the framework to clipped SGD and normalized SGD (including mini-batch and momentum variants), deriving explicit stability and rate bounds that improve upon prior results in the heavy-tailed regime. The results illuminate trade-offs between clipping and normalization and provide concrete guidance on parameter choices to achieve favorable generalization in practice, with potential impact on training regimes where gradient noise is heavy-tailed.
Abstract
The empirical evidence indicates that stochastic optimization with heavy-tailed gradient noise is more appropriate to characterize the training of machine learning models than that with standard bounded gradient variance noise. Most existing works on this phenomenon focus on the convergence of optimization errors, while the analysis for generalization bounds under the heavy-tailed gradient noise remains limited. In this paper, we develop a general framework for establishing generalization bounds under heavy-tailed noise. Specifically, we introduce a truncation argument to achieve the generalization error bound based on the algorithmic stability under the assumption of bounded $p$th centered moment with $p\in(1,2]$. Building on this framework, we further provide the stability and generalization analysis for several popular stochastic algorithms under heavy-tailed noise, including clipped and normalized stochastic gradient descent, as well as their mini-batch and momentum variants.