Improved Convergence in High Probability of Clipped Gradient Methods with Heavy Tails

TL;DR

This paper introduces a new analysis method for clipped gradient algorithms with heavy-tailed noise, improving high-probability convergence guarantees and allowing adaptive parameters without prior problem knowledge.

Contribution

We develop a novel supermartingale-based analysis that reduces failure probability dependence on iteration count, enabling adaptive step sizes and clipping in heavy-tailed settings.

Findings

Improved high-probability convergence bounds for clipped gradient methods.

Applicable to both convex and nonconvex stochastic optimization.

Eliminates need for prior problem constants in parameter setting.

Abstract

In this work, we study the convergence \emph{in high probability} of clipped gradient methods when the noise distribution has heavy tails, ie., with bounded $p$th moments, for some $1<p\le2$. Prior works in this setting follow the same recipe of using concentration inequalities and an inductive argument with union bound to bound the iterates across all iterations. This method results in an increase in the failure probability by a factor of $T$, where $T$ is the number of iterations. We instead propose a new analysis approach based on bounding the moment generating function of a well chosen supermartingale sequence. We improve the dependency on $T$ in the convergence guarantee for a wide range of algorithms with clipped gradients, including stochastic (accelerated) mirror descent for convex objectives and stochastic gradient descent for nonconvex objectives. This approach naturally allows the algorithms to use time-varying step sizes and clipping parameters when the time horizon is unknown, which appears impossible in prior works. We show that in the case of clipped stochastic mirror descent, problem constants, including the initial distance to the optimum, are not required when setting step sizes and clipping parameters.