Stability and Generalization of Nonconvex Optimization with Heavy-Tailed Noise

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.

💡 Research Summary

The paper addresses the gap in theoretical understanding of generalization for non‑convex stochastic optimization when the gradient noise exhibits heavy tails. Instead of the classical bounded‑variance assumption, the authors adopt the p‑bounded‑centered‑moment (p‑BCM) condition, requiring that for some p∈(1,2] the p‑th moment of the deviation between stochastic and population gradients is uniformly bounded by σₚᵖ. This model captures a wide range of heavy‑tailed behaviors observed in deep learning and reinforcement learning.

A central contribution is a new truncation technique that controls the heavy‑tailed noise. By clipping each stochastic gradient at a threshold γ and analyzing the clipped variable’s p‑th moment, the authors are able to bound the deviation between the population gradient ∇F(A(S)) and the empirical gradient ∇F_S(A(S)) in terms of algorithmic stability. They prove that if a randomized algorithm A is ϵ‑uniformly stable in gradients, then

 E