Bringing Order to Asynchronous SGD: Towards Optimality under Data-Dependent Delays with Momentum
This work provides a theoretical breakthrough for asynchronous distributed training by establishing optimal rates under realistic data-dependent delays, addressing a long-standing gap in the literature.
The paper proposes a momentum-based asynchronous SGD framework that preserves information from delayed gradients, achieving the first optimal convergence rates for data-dependent delays in both convex and non-convex smooth settings.
Asynchronous stochastic gradient descent (SGD) enables scalable distributed training but suffers from gradient staleness. Existing mitigation strategies, such as delay-adaptive learning rates and staleness-aware filtering, typically attenuate or discard delayed gradients, introducing systematic bias: updates from simpler or faster-to-process samples are overrepresented, while gradients from more complex samples are delayed or suppressed. In contrast, prior approaches to data-dependent delays rely on a Lipschitz assumption that yields suboptimal rates or leave the smooth, convex case unaddressed. We propose a momentum-based asynchronous framework designed to preserve information from delayed gradients while mitigating the effects of staleness. We establish the first optimal convergence rates for data-dependent delays in both convex and non-convex smooth setups, providing a new result for asynchronous optimization under standard assumptions. Additionally, we derive robust learning-rate schedules that simplify hyperparameter tuning in practice.