Efficient Distributed Optimization under Heavy-Tailed Noise

TL;DR

This paper introduces TailOPT, a novel distributed optimization framework that effectively handles heavy-tailed stochastic gradient noise, improving training efficiency and performance in large-scale machine learning models.

Contribution

It proposes TailOPT with adaptive clipping techniques, providing convergence guarantees under heavy-tailed noise and introducing a memory-efficient variant, $Bi^2Clip.

Findings

TailOPT outperforms existing methods on language tasks.

$Bi^2Clip achieves adaptive-like performance without extra gradient statistics.

The framework guarantees convergence under unbounded gradient variance.

Abstract

Distributed optimization has become the default training paradigm in modern machine learning due to the growing scale of models and datasets. To mitigate communication overhead, local updates are often applied before global aggregation, resulting in a nested optimization approach with inner and outer steps. However, heavy-tailed stochastic gradient noise remains a significant challenge, particularly in attention-based models, hindering effective training. In this work, we propose TailOPT, an efficient framework designed to address heavy-tailed noise by leveraging adaptive optimization or clipping techniques. We establish convergence guarantees for the TailOPT framework under heavy-tailed noise with potentially unbounded gradient variance and local updates. Among its variants, we highlight a memory and communication efficient instantiation which we call $Bi^2Clip$, which performs coordinate-wise clipping at both the inner and outer optimizers, achieving adaptive-like performance (e.g., Adam) without the cost of maintaining or transmitting additional gradient statistics. Empirically, TailOPT, including $Bi^2Clip$, demonstrates superior performance on several language tasks and models, outperforming state-of-the-art methods.