Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Optimal Asynchronous Stochastic Nonconvex Optimization under Heavy-Tailed Noise

Yidong Wu, Luo Luo

TL;DR

This work tackles asynchronous stochastic nonconvex optimization with heavy-tailed gradient noise across heterogeneous workers. It introduces RANSGDm, an asynchronous normalized SGD with momentum and a delay-threshold mechanism that discards stale gradients, achieving optimal time complexity under $p$-bounded central moment with $p\in(1,2]$. Theoretical results establish matching lower bounds for both fixed and universal computation models, while experiments on a synthetic quadratic problem and GPT-2 training validate superior performance under heavy-tailed noise. The approach advances practical asynchronous optimization in distributed systems where delays and noise deviate from standard bounded-variance assumptions, with implications for large-scale machine learning and NLP model training.

Abstract

This paper considers the problem of asynchronous stochastic nonconvex optimization with heavy-tailed gradient noise and arbitrarily heterogeneous computation times across workers. We propose an asynchronous normalized stochastic gradient descent algorithm with momentum. The analysis show that our method achieves the optimal time complexity under the assumption of bounded $p$th-order central moment with $p\in(1,2]$. We also provide numerical experiments to show the effectiveness of proposed method.

Optimal Asynchronous Stochastic Nonconvex Optimization under Heavy-Tailed Noise

TL;DR

This work tackles asynchronous stochastic nonconvex optimization with heavy-tailed gradient noise across heterogeneous workers. It introduces RANSGDm, an asynchronous normalized SGD with momentum and a delay-threshold mechanism that discards stale gradients, achieving optimal time complexity under -bounded central moment with . Theoretical results establish matching lower bounds for both fixed and universal computation models, while experiments on a synthetic quadratic problem and GPT-2 training validate superior performance under heavy-tailed noise. The approach advances practical asynchronous optimization in distributed systems where delays and noise deviate from standard bounded-variance assumptions, with implications for large-scale machine learning and NLP model training.

Abstract

This paper considers the problem of asynchronous stochastic nonconvex optimization with heavy-tailed gradient noise and arbitrarily heterogeneous computation times across workers. We propose an asynchronous normalized stochastic gradient descent algorithm with momentum. The analysis show that our method achieves the optimal time complexity under the assumption of bounded th-order central moment with . We also provide numerical experiments to show the effectiveness of proposed method.
Paper Structure (30 sections, 17 theorems, 134 equations, 3 figures, 1 table, 1 algorithm)

This paper contains 30 sections, 17 theorems, 134 equations, 3 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. The Algorithm and Main Results
  4. The Complexity Analysis
  5. Discussion
  6. The Lower Complexity Bounds
  7. The Fixed Computation Model
  8. The Universal Computation Model
  9. Experiments
  10. The Quadratic Problem
  11. Training GPT-2
  12. Conclusion
  13. The Proofs for Results in Section \ref{['sec:alg-main']}
  14. The Proof of Lemma \ref{['lem:staleness']}
  15. The Proof of Lemma \ref{['lem:descent_mom']}
  16. ...and 15 more sections

Key Result

Theorem 1

Under Assumptions ass:smooth and ass:pBCM, we run our RANSGDm (Algorithm alg:ringmaster_nsgd_mom_server) by taking where Then the output $\hat{x}\in{\mathbb{R}}^d$ holds $\mathbb{E}\|\nabla f(\hat{x})\|\le \epsilon$.

Figures (3)

  • Figure 1: The results of solving the quadratic problem, where the expected delay of slow workers is 0.02s.
  • Figure 2: The results of solving the quadratic problem, where the expected delay of slow workers is 0.005s.
  • Figure 3: The results of training GPT-2 model with different delay settings. The sub-figures (a) and (b) set exponential distributed delays with expectations $\{0,0.5,0.9,1.2,1.5,2.0,2.5\}$ and $\{1,2,3,4,5,6,7\}$, respectively. The sub-figure (c) set Pareto distributed delays with expectations $\{1,2,3,4,5,6,7\}$.

Theorems & Definitions (31)

  • Remark 1
  • Theorem 1
  • Corollary 1
  • Corollary 2
  • Remark 2
  • Lemma 1
  • Lemma 2
  • Lemma 3
  • Definition 1
  • Theorem 2
  • ...and 21 more