Stability and Generalization of Nonconvex Optimization with Heavy-Tailed Noise

TL;DR

This paper develops a framework to analyze the generalization bounds of stochastic optimization algorithms under heavy-tailed gradient noise, extending understanding beyond convergence to model generalization.

Contribution

It introduces a truncation-based approach for deriving generalization bounds under heavy-tailed noise and analyzes popular algorithms like clipped and normalized SGD.

Findings

Established generalization bounds using stability with heavy-tailed noise

Analyzed stability of clipped and normalized stochastic gradient descent

Provided insights into the effects of heavy-tailed noise on algorithm performance

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.