Stability and Generalization of Nonconvex Optimization with Heavy-Tailed Noise

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.

• The paper provides a generalization error bound for stochastic nonconvex optimization under heavy-tailed noise, demonstrating that any algorithm that maintains stability in gradients achieves this bound.
• The method employs a truncation argument without reliance on SDE-based arguments, showing how to manage gradient noise with unbounded variance effectively.
• The framework favorably compares normalized Stochastic Gradient Descent (SGD) variants to clipped SGD, highlighting the superiority of normalization combined with mini-batching or momentum.
• Clipped SGD shows a rate of $O(n^{-rac{p-1}{3(3p-2)})$ compared to the better rate of normalized SGD methods which all scale as $O(n^{-rac{p-1}{7p-6)}. See the sue of NSGD with Clipping and Normalization for population risk

Motivation and problem setting

The paper studies generalization guarantees for stochastic nonconvex optimization when the gradient noise is heavy-tailed. The standard assumption in stability-based generalization analysis is bounded variance, i.e., a bounded second centered moment of the stochastic gradients. However, empirical evidence from deep neural network training and reinforcement learning indicates that stochastic gradient noise is better modeled by heavy-tailed distributions, formalized here through the bounded $p$th centered moment ($p$-BCM) condition: $$\mathbb{E}\bigl[\|\nabla f(x;\xi)-\nabla F(x)\|^p\bigr] \le \sigma_p^p$$ for some $p \in (1,2]$. This condition is weaker and more general than the $\alpha$-stable model used in Lévy-driven SDE analyses, which implicitly assumes identically distributed noise across coordinates.

The paper's objective is the population gradient norm $\mathbb{E}_{S,A}\bigl[\|\nabla F(A(S))\|\bigr]$ at the output $A(S)$ of an algorithm trained on $n$ i.i.d. samples, decomposed into an optimization error term $\mathbb{E}\bigl[\|\nabla F_S(A(S))\|\bigr]$ plus a generalization error term $$\mathbb{E}\bigl[\|\nabla F(A(S))-\nabla F_S(A(S))\|\bigr]$$. Prior work had established convergence rates under $p$0-BCM for clipped SGD and normalized SGD variants, but generalization bounds under this noise model remained limited: existing results either relied on SDE-based continuous-time arguments restricted to vanilla SGD, or assumed noise conditions not equivalent to $p$1-BCM.

Main result: generalization via truncation

The central contribution is a generalization error bound under $p$2-BCM for any algorithm that is $p$3-uniformly stable in gradients (i.e., replacing one training sample changes expected squared gradient norms by at most $p$4). The bound reads:

$p$5

where $p$6 at $p$7 and $p$8 for $p$9, with $$\mathbb{E}\bigl[\|\nabla f(x;\xi)-\nabla F(x)\|^p\bigr] \le \sigma_p^p$$0. At $$\mathbb{E}\bigl[\|\nabla f(x;\xi)-\nabla F(x)\|^p\bigr] \le \sigma_p^p$$1 the bound exactly recovers the bounded-variance result of Lei et al., so the extension interpolates continuously across the tail index.

The proof is nontrivial because the ghost-sample decomposition produces terms $$\mathbb{E}\bigl[\|\nabla f(x;\xi)-\nabla F(x)\|^p\bigr] \le \sigma_p^p$$2 whose variance is unbounded under heavy-tailed noise, so the standard Jensen-plus-squared-norm expansion fails. The key technique is a truncation argument: each $$\mathbb{E}\bigl[\|\nabla f(x;\xi)-\nabla F(x)\|^p\bigr] \le \sigma_p^p$$3 is split into a truncated part built from clipped gradient noise $$\mathbb{E}\bigl[\|\nabla f(x;\xi)-\nabla F(x)\|^p\bigr] \le \sigma_p^p$$4 (centered to be mean-zero) and a residual tail part. The truncated part is controlled via $$\mathbb{E}\bigl[\|\nabla f(x;\xi)-\nabla F(x)\|^p\bigr] \le \sigma_p^p$$5, which converts the $$\mathbb{E}\bigl[\|\nabla f(x;\xi)-\nabla F(x)\|^p\bigr] \le \sigma_p^p$$6th moment into a variance-like quantity scaling as $$\mathbb{E}\bigl[\|\nabla f(x;\xi)-\nabla F(x)\|^p\bigr] \le \sigma_p^p$$7; cross terms are handled by an additional dataset perturbation argument yielding $$\mathbb{E}\bigl[\|\nabla f(x;\xi)-\nabla F(x)\|^p\bigr] \le \sigma_p^p$$8. The residual tail satisfies $$\mathbb{E}\bigl[\|\nabla f(x;\xi)-\nabla F(x)\|^p\bigr] \le \sigma_p^p$$9, giving a bias of order $p \in (1,2]$0. Optimizing over the truncation level $p \in (1,2]$1 balances the two terms and yields the stated rate.

Compared with prior nonconvex bounds, the rate degrades gracefully with heavier tails: $p \in (1,2]$2 at $p \in (1,2]$3 versus $p \in (1,2]$4 for $p \in (1,2]$5, while uniform-convergence alternatives such as Mei et al.'s $p \in (1,2]$6 depend on dimension and require sub-Gaussian tails. An implication is that stability-based generalization remains viable in high dimensions under heavy-tailed noise, provided the algorithm itself is stable.

Applications: clipping versus normalization

The framework is applied to four algorithms, all analyzed under $p \in (1,2]$7-smoothness, $p \in (1,2]$8, and a moment bound $p \in (1,2]$9 along the trajectory. The resulting population risk rates after tuning $\alpha$0 are summarized below.

Algorithm	Stability driver	Population risk rate
Clipped SGD	$\alpha$1	$\alpha$2
Mini-batch NSGD (NSGD-B)	$\alpha$3	$\alpha$4
NSGD with momentum (NSGD-M)	$\alpha$5	$\alpha$6
NSGD with clipping and momentum (NSGD-CM)	$\alpha$7	$\alpha$8

Two structural findings emerge. First, the stability of clipped SGD scales linearly in the clipping threshold $\alpha$9, whereas normalized methods have stability independent of $\mathbb{E}_{S,A}\bigl[\|\nabla F(A(S))\|\bigr]$0; balancing against the optimization error (which improves as $\mathbb{E}_{S,A}\bigl[\|\nabla F(A(S))\|\bigr]$1 grows, since clipping bias is $\mathbb{E}_{S,A}\bigl[\|\nabla F(A(S))\|\bigr]$2) forces a worse overall rate for clipped SGD. Second, vanilla normalized SGD can fail to converge — the paper gives a concrete one-dimensional example where normalization turns the iterates into a simple random walk that never reaches the unique minimizer — which motivates either mini-batching or momentum. Both remedies achieve the same rate $\mathbb{E}_{S,A}\bigl[\|\nabla F(A(S))\|\bigr]$3, though through different mechanisms: NSGD-B reduces the gradient estimation error as $\mathbb{E}_{S,A}\bigl[\|\nabla F(A(S))\|\bigr]$4 via a martingale difference inequality, while NSGD-M reduces it via the momentum bias $\mathbb{E}_{S,A}\bigl[\|\nabla F(A(S))\|\bigr]$5, with the momentum parameter ultimately limiting further improvement. Adding clipping on top of momentum (NSGD-CM) does not improve the risk bound over NSGD-M, indicating that clipping contributes nothing once normalization and momentum are present — consistent with recent convergence-side findings that clipping can be dispensed with.

Limitations and open questions

Several caveats qualify these results. All algorithmic analyses assume $\mathbb{E}_{S,A}\bigl[\|\nabla F(A(S))\|\bigr]$6-smoothness and $\mathbb{E}_{S,A}\bigl[\|\nabla F(A(S))\|\bigr]$7; the nonsmooth setting is explicitly left open. The population risk bounds additionally require a trajectory-wise moment bound $\mathbb{E}_{S,A}\bigl[\|\nabla F(A(S))\|\bigr]$8, which is an auxiliary condition beyond $\mathbb{E}_{S,A}\bigl[\|\nabla F(A(S))\|\bigr]$9-BCM itself. The results are in-expectation rather than high-probability guarantees, and no lower bounds are established for the generalization problem under heavy-tailed noise, so it remains unknown whether the exponent $A(S)$0 is optimal. Finally, the comparison between NSGD-B and NSGD-M rests on specific hyperparameter scalings; whether other tunings yield sharper rates is not addressed.

Conclusion

This paper establishes the first systematic link between algorithmic stability and generalization for nonconvex stochastic optimization under the $A(S)$1-BCM heavy-tailed noise condition, via a truncation-and-tail-decomposition proof that recovers the bounded-variance bound at $A(S)$2. Applied to clipped and normalized SGD variants, the framework shows that normalization combined with mini-batching or momentum yields strictly better population risk rates than clipping alone, and that clipping adds no benefit atop normalization with momentum. The main open problems are nonsmooth objectives and matching lower bounds for learning under heavy-tailed noise.