Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Shuhua Yu, Dusan Jakovetic, Soummya Kar

TL;DR

This work tackles decentralized nonconvex stochastic optimization under heavy-tailed gradient noise on graphs. It introduces GT-NSGDm, a gradient-tracking based normalization method with momentum to mitigate heavy-tailed noise while addressing heterogeneity across nodes. The authors prove that, with known tail index p, the algorithm achieves an optimal non-asymptotic rate of O(1/T^{(p-1)/(3p-2)}) and, when p is unknown, a topology-robust rate of O(1/T^{(p-1)/(2p)}) with a node-count speedup of n^{1-1/p} for p<2, supported by empirical results on robust linear regression and decentralized Transformer training. The results demonstrate robustness and efficiency of GT-NSGDm across topologies, suggesting practical impact for large-scale decentralized training under realistic heavy-tailed gradient noise.

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.

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

TL;DR

This work tackles decentralized nonconvex stochastic optimization under heavy-tailed gradient noise on graphs. It introduces GT-NSGDm, a gradient-tracking based normalization method with momentum to mitigate heavy-tailed noise while addressing heterogeneity across nodes. The authors prove that, with known tail index p, the algorithm achieves an optimal non-asymptotic rate of O(1/T^{(p-1)/(3p-2)}) and, when p is unknown, a topology-robust rate of O(1/T^{(p-1)/(2p)}) with a node-count speedup of n^{1-1/p} for p<2, supported by empirical results on robust linear regression and decentralized Transformer training. The results demonstrate robustness and efficiency of GT-NSGDm across topologies, suggesting practical impact for large-scale decentralized training under realistic heavy-tailed gradient noise.

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 -th moment for . 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 , which matches the lower bound in the centralized setup. When tail index is unknown, GT-NSGDm attains a non-asymptotic rate that is, for , topology independent and has a speedup factor in terms of the number of nodes . 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.
Paper Structure (19 sections, 10 theorems, 55 equations, 5 figures, 3 tables, 1 algorithm)

This paper contains 19 sections, 10 theorems, 55 equations, 5 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Contributions
  3. Related Work
  4. Paper Outline and Notation
  5. Problem Formulation
  6. Algorithm Development: GT-NSGDm
  7. Main Results
  8. Experiments
  9. Robust linear regression on synthetic tokenized data
  10. Decentralized training of Transformers
  11. Conclusion and Future Work
  12. Proof of Claim \ref{['cl:normaldsgd']}
  13. Proofs of Theorems
  14. Preliminaries
  15. Intermediate Lemmas
  16. ...and 4 more sections

Key Result

Theorem 1

Let Assumptions as:ldd, as:smooth, as:noise, as:network hold. Denote $f(\bar{\boldsymbol{x}}^0) - f_* = \Delta_0, [\nabla f_1(\bar{\boldsymbol{x}}^0), \ldots, \nabla f_n(\bar{\boldsymbol{x}}^0)]^\top = \nabla F(\mathbf{1}_n \otimes \bar{\boldsymbol{x}}^0)$. Take and $1 - \beta = 1/T^{\frac{p}{3p-2}}$. Assume $\beta \ge 1/10$, then the sequence generated from GT-NSGDm satisfies that

Figures (5)

  • Figure 1: Comparisons of the empirical density of gradient noise norm in different epochs of training a Transformer model with a synthetic Lévy $\alpha$-stable distribution.
  • Figure 2: Comparison of performance on a ring graph under various types of injected stochastic gradient noise, measured by the average estimation error $(1/n) \sum_{i=1}^n \| \boldsymbol{w}_i^t - \boldsymbol{w}_* \|$ over step count $t$.
  • Figure 3: Empirical studies on GT-NSGDm's dependence on problem parameters $\lambda, \sigma, n$.
  • Figure 4: Comparison of graph-wide average validation losses in decentralized training of Transformer models, over ring, directed exponential, and complete graphs.
  • Figure 5: Comparison of graph-wide average training losses in decentralized training of Transformer models, over ring, directed exponential, and complete graphs.

Theorems & Definitions (26)

  • Remark 1: Heavy-tailed distributions
  • Remark 2: Empirical evidence
  • Remark 3: Why vanilla gradient normalization fails?
  • Claim 1
  • Theorem 1
  • Remark 4: Order-optimal rate
  • Remark 5: Speedup in $n$
  • Theorem 2
  • Remark 6: Speedup in $n$ and topology independent rate
  • proof
  • ...and 16 more