Ordered Momentum for Asynchronous SGD

TL;DR

The paper addresses incorporating momentum into asynchronous SGD (ASGD) without sacrificing convergence by introducing ordered momentum (OrMo). OrMo organizes gradients into iteration-based groups and updates momentum accordingly to maintain coherent updates despite out-of-order arrivals, with a theoretical convergence guarantee for non-convex objectives under both constant and delay-adaptive learning rates. The main results show a nonconvex convergence rate of $\mathcal{O}\big(\sqrt{L\sigma^2/T} + (KLG/T)^{2/3} + KL/T\big)$ for constant rate and $\mathcal{O}\big(\sqrt{L\sigma^2/T} + KL/T\big)$ for OrMo-DA, notably without dependence on the maximum delay. Empirically, OrMo and OrMo-DA achieve superior convergence and faster wall-clock performance compared to ASGD and momentum-based baselines across CIFAR-10/100 with homogeneous and heterogeneous worker settings, demonstrating robustness to delays and slow workers.

Abstract

Distributed learning is essential for training large-scale deep models. Asynchronous SGD (ASGD) and its variants are commonly used distributed learning methods, particularly in scenarios where the computing capabilities of workers in the cluster are heterogeneous. Momentum has been acknowledged for its benefits in both optimization and generalization in deep model training. However, existing works have found that naively incorporating momentum into ASGD can impede the convergence. In this paper, we propose a novel method called ordered momentum (OrMo) for ASGD. In OrMo, momentum is incorporated into ASGD by organizing the gradients in order based on their iteration indexes. We theoretically prove the convergence of OrMo with both constant and delay-adaptive learning rates for non-convex problems. To the best of our knowledge, this is the first work to establish the convergence analysis of ASGD with momentum without dependence on the maximum delay. Empirical results demonstrate that OrMo can achieve better convergence performance compared with ASGD and other asynchronous methods with momentum.

Figures (9)

• Figure 1: An example of the momentum term ${\bf u}_{10}$ in SSGDm when $K=4$. The gradients shown in red indicate those having not arrived at the server. In this case, ${\bf u}_{10}=\beta^2\times \left(\eta{\bf g}_0^{0}+\eta{\bf g}_0^{1}+\eta{\bf g}_0^{2}+\eta{\bf g}_0^{3}\right)+\beta^1\times\left(\eta{\bf g}_4^0+\eta{\bf g}_4^1+\eta{\bf g}_4^2+\eta{\bf g}_4^3\right)+\beta^0\times \left(\eta{\bf g}_8^0+\eta{\bf g}_8^3\right)$.
• Figure 2: An example of the momentum term ${\bf u}_{10}$ in OrMo when $K=4$. The gradients shown in red indicate those having not arrived at the server. In this case, ${\bf u}_{10}=\beta^3\times \left(\eta{\bf g}_0^{0}+\eta{\bf g}_0^{1}+\eta{\bf g}_0^{2}+\eta{\bf g}_0^{3}\right)+\beta^2\times\left(\eta{\bf g}_1^{k_0}+\eta{\bf g}_2^{k_1}+\eta{\bf g}_3^{k_2}\right)+\beta^1\times \left(\eta{\bf g}_6^{k_5}+\eta{\bf g}_8^{k_7}\right)+\beta^0\times\left(\eta{\bf g}_9^{k_8}\right)$.
• Figure 3: Test accuracy curves on CIFAR10 with different numbers of workers.
• Figure 4: Test accuracy curves on CIFAR100 with different numbers of workers.
• Figure 5: Training curves with respect to wall-clock time on CIFAR10 when $K=16$.

Theorems & Definitions (31)

• Remark 1
• Proposition 1
• Remark 2
• Remark 3
• Lemma 1
• Lemma 2
• Lemma 3
• Theorem 1
• Theorem 2
• Remark 4