Distributed Stochastic Gradient Descent with Staleness: A Stochastic Delay Differential Equation Based Framework

TL;DR

This work addresses the bottleneck of gradient staleness in asynchronous distributed SGD by introducing a stochastic delay differential equation (SDDE) framework that accommodates non-memoryless computation delays via a Poisson approximation. By linking gradient noise and gradient staleness to SDDE dynamics, it derives convergence conditions from characteristic roots and analyzes first hitting times to quantify training time under various delay scenarios. The study provides practical parameter-optimization strategies for different ASGD paradigms, including optimal worker counts and scheduling in ideal and bandwidth-constrained settings, as well as event-triggered communication schemes. Numerical experiments on synthetic models and MNIST demonstrate the framework’s ability to predict when increased parallelism helps or hinders convergence, and to guide design choices that balance update frequency, staleness, and bandwidth. Overall, the SDDE-based framework offers a principled approach to scheduling and protocol design for scalable, asynchronous distributed learning with staleness considerations.

Abstract

Distributed stochastic gradient descent (SGD) has attracted considerable recent attention due to its potential for scaling computational resources, reducing training time, and helping protect user privacy in machine learning. However, the staggers and limited bandwidth may induce random computational/communication delays, thereby severely hindering the learning process. Therefore, how to accelerate asynchronous SGD by efficiently scheduling multiple workers is an important issue. In this paper, a unified framework is presented to analyze and optimize the convergence of asynchronous SGD based on stochastic delay differential equations (SDDEs) and the Poisson approximation of aggregated gradient arrivals. In particular, we present the run time and staleness of distributed SGD without a memorylessness assumption on the computation times. Given the learning rate, we reveal the relevant SDDE's damping coefficient and its delay statistics, as functions of the number of activated clients, staleness threshold, the eigenvalues of the Hessian matrix of the objective function, and the overall computational/communication delay. The formulated SDDE allows us to present both the distributed SGD's convergence condition and speed by calculating its characteristic roots, thereby optimizing the scheduling policies for asynchronous/event-triggered SGD. It is interestingly shown that increasing the number of activated workers does not necessarily accelerate distributed SGD due to staleness. Moreover, a small degree of staleness does not necessarily slow down the convergence, while a large degree of staleness will result in the divergence of distributed SGD. Numerical results demonstrate the potential of our SDDE framework, even in complex learning tasks with non-convex objective functions.

Figures (5)

• Figure 1: System model.
• Figure 2: The impact of $\tau$ on the convergence of the SGD algorithm.
• Figure 3: The impact of the gradient noise and learning rate on the convergence of the gradient descent algorithms.
• Figure 4: Run-time and step staleness in $B$-ASGD.
• Figure 5: ASGD with gradient staleness on the MNIST dataset.