Asynchronous Heavy-Tailed Optimization

TL;DR

This paper investigates asynchronous optimization algorithms in the presence of heavy-tailed gradient noise, proposing modifications that improve convergence, delay tolerance, and robustness in training transformer models.

Contribution

It introduces delay-aware learning rate scheduling and delay compensation techniques with theoretical guarantees and empirical improvements over existing methods.

Findings

Convergence rates match synchronous algorithms under heavy-tailed noise.

Enhanced delay tolerance compared to prior asynchronous approaches.

Better accuracy and runtime trade-offs in image and language tasks.

Abstract

Heavy-tailed stochastic gradient noise, commonly observed in transformer models, can destabilize the optimization process. Recent works mainly focus on developing and understanding approaches to address heavy-tailed noise in the centralized or distributed, synchronous setting, leaving the interactions between such noise and asynchronous optimization underexplored. In this work, we investigate two communication schemes that handle stragglers with asynchronous updates in the presence of heavy-tailed gradient noise. We propose and theoretically analyze algorithmic modifications based on delay-aware learning rate scheduling and delay compensation to enhance the performance of asynchronous algorithms. Our convergence guarantees under heavy-tailed noise match the rate of the synchronous counterparts and improve delay tolerance compared with existing asynchronous approaches. Empirically, our approaches outperform prior synchronous and asynchronous methods in terms of accuracy/runtime trade-offs and are more robust to hyperparameters in both image and language tasks.