Decentralized Nonconvex Optimization under Heavy-Tailed Noise: Normalization and Optimal Convergence

TL;DR

This paper introduces GT-NSGDm, a normalization-based decentralized algorithm for nonconvex optimization under heavy-tailed noise, achieving optimal convergence rates matching centralized lower bounds.

Contribution

It proposes a novel normalization and gradient tracking method that guarantees optimal convergence rates in decentralized heavy-tailed noise settings, a first in the literature.

Findings

GT-NSGDm achieves the optimal non-asymptotic convergence rate of O(1/T^{(p-1)/(3p-2)})

The method is topology-independent when the tail index p is unknown, with a convergence rate of O(1/T^{(p-1)/(2p)})

Experiments show GT-NSGDm outperforms baselines in robustness and efficiency on real-world tasks.

Abstract

Heavy-tailed noise in nonconvex stochastic optimization has garnered increasing research interest, as empirical studies, including those on training attention models, suggest it is a more realistic gradient noise condition. This paper studies first-order nonconvex stochastic optimization under heavy-tailed gradient noise in a decentralized setup, where each node can only communicate with its direct neighbors in a predefined graph. Specifically, we consider a class of heavy-tailed gradient noise that is zero-mean and has only $p$-th moment for $p \in (1, 2]$. We propose GT-NSGDm, Gradient Tracking based Normalized Stochastic Gradient Descent with momentum, that utilizes normalization, in conjunction with gradient tracking and momentum, to cope with heavy-tailed noise on distributed nodes. We show that, when the communication graph admits primitive and doubly stochastic weights, GT-NSGDm guarantees, for the \textit{first} time in the literature, that the expected gradient norm converges at an optimal non-asymptotic rate $O\big(1/T^{(p-1)/(3p-2)}\big)$, which matches the lower bound in the centralized setup. When tail index $p$ is unknown, GT-NSGDm attains a non-asymptotic rate $O\big( 1/T^{(p-1)/(2p)} \big)$ that is, for $p < 2$, topology independent and has a speedup factor $n^{1-1/p}$ in terms of the number of nodes $n$. Finally, experiments on nonconvex linear regression with tokenized synthetic data and decentralized training of language models on a real-world corpus demonstrate that GT-NSGDm is more robust and efficient than baselines.