Stability and Generalization of Stochastic Nonconvex Optimization under Heavy-Tailed Noise: From Gradient Norm to Goldstein Stationarity

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 stability-based frameworks to establish dimension-independent generalization bounds of nonconvex optimization under heavy-tailed noise for both smooth and nonsmooth settings. Specifically, we introduce a truncation argument to achieve the generalization error bound based on the algorithmic stability under finite pth moment conditions with p∈(1,2]. For the smooth case, we apply our framework to derive population gradient norm bounds for several popular stochastic optimization algorithms, including clipped and normalized stochastic gradient descent, as well as their mini-batch and momentum variants. For the nonsmooth case, we establish a population Goldstein stationarity guarantee for online-to-nonconvex conversion methods. It is worth noting that our theory for the nonsmooth setting is established by introducing an on-average gradient stability for the smooth surrogate function, which does not require the additional weak convexity assumption used in existing analysis based on the Moreau envelope.