Asynchronous Stochastic Optimization Robust to Arbitrary Delays
Year: 2022

Abstract:
We consider stochastic optimization with delayed gradients where, at each time step t, the algorithm makes an update using a stale stochastic gradient from step t − dt for some arbitrary delay dt . This setting abstracts asynchronous distributed optimization where a central server receives gradient updates computed by worker machines. These machines can experience computation and communication loads that might vary significantly over time. In the general non-convex smooth optimization setting, we give a simple and efficient algorithm that requires O (σ2/ε4 + τ/ε2) steps for finding an ε-stationary point x, where τ is the average delay 1 T ∑T t=1 dt and σ2 is the variance of the stochastic gradients. This improves over previous work, which showed that stochastic gradient decent achieves the same rate but with respect to the maximal delay maxt dt , that can be significantly larger than the average delay especially in heterogeneous distributed systems. Our experiments demonstrate the efficacy and robustness of our algorithm in cases where the delay distribution is skewed or heavy-tailed.

text: We consider stochastic optimization with delayed gradients where, at each time step 𝑡, the algorithm makes an update using a stale stochastic gradient from step 𝑡 − 𝑑𝑡 for some arbitrary delay 𝑑𝑡 . This setting abstracts asynchronous distributed optimization where a central server receives gradient updates computed by worker machines. These machines can experience computation and communication loads that might vary significantly over time. In the general non-convex smooth optimization setting, we give a simple and efficient algorithm that requires 𝑂 (σ2/ϵ4 + τ/ϵ2) steps for finding an ϵ-stationary point 𝑥, where τ is the average delay 1
𝑇
∑𝑇 𝑡=1 𝑑𝑡 and σ2 is
the variance of the stochastic gradients. This improves over previous work, which showed that stochastic gradient decent achieves the same rate but with respect to the maximal delay max𝑡 𝑑𝑡 , that can be significantly larger than the average delay especially in heterogeneous distributed systems. Our experiments demonstrate the efficacy and robustness of our algorithm in cases where the delay distribution is skewed or heavy-tailed.

heading: 1 Introduction
text: Gradient-based iterative optimization methods are widely used in large-scale machine learning applications as they are extremely simple to implement and use, and come with mild computational requirements. On the other hand, in their standard formulation they are also inherently serial and synchronous due to their iterative nature. For example, in stochastic gradient descent (SGD), each step involves an update of the form 𝑥𝑡+1 = 𝑥𝑡 − η𝑔𝑡 where 𝑥𝑡 is the current iterate, and 𝑔𝑡 is a (stochastic) gradient vector evaluated at 𝑥𝑡 . To progress to the next step of the method, the subsequent iterate 𝑥𝑡+1 has to be fully determined by the end of step 𝑡 as it is required for future gradient queries. Evidently, this scheme has to wait for the computation of the gradient 𝑔𝑡 to complete (this is often the most computationally intensive part in SGD) before it can evaluate 𝑥𝑡+1. In modern large scale machine learning applications, a direct serial implementation of gradient methods like SGD is overly costly, and parallelizing the optimization process over several cores or machines is desired. Perhaps the most common parallelization approach is via mini-batching, where computation of stochastic gradients is distributed across several worker machines that send updates to a parameter server. The parameter server is responsible for accruing the individual updates into a single averaged gradient, and consequently, updating the optimization parameters using this gradient.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
While mini-batching is well understood theoretically [e.g., 16, 9, 8, 10], it is still fundamentally synchronous in nature and its performance is adversely determined by the slowest worker machine: the parameter server must wait for all updates from all workers to arrive before it can update the model it maintains. This could cause serious performance issues in heterogeneous distributed networks, where worker machines may be subject to unpredictable loads that vary significantly between workers (due to different hardware, communication bandwidth, etc.) and over time (due to varying users load, power outages, etc.). An alternative approach that has recently gained popularity is to employ asynchronous gradient updates [e.g., 21, 2, 7, 18, 11]; namely, each worker machine computes gradients independently of the other machines, possibly on different iterates, and sends updates to the parameter server in an asynchronous fashion. This implies the parameter server might be making stale updates based on delayed gradients taken at earlier, out-of-date iterates. While these methods often work well in practice, they have proven to be much more intricate and challenging to analyze theoretically than synchronous gradient methods, and overall our understanding of asynchronous updates remains lacking. Recently, Arjevani et al. [4] and subsequently Stich and Karimireddy [26] have made significant progress in analyzing delayed asynchronous gradient methods. They have shown that in stochastic optimization, delays only affect a lower-order term in the convergence bounds. In other words, if the delays are not too large, the convergence rate of SGD may not be affected by the delays. (4 first proved this for quadratic objectives; 26 then proved a more general result for smooth functions.) More concretely, Stich and Karimireddy [26] showed that SGD with a sufficiently attenuated step size to account for the delays attains an iteration complexity bound of the form
𝑂
( σ2
ϵ4 + τmax ϵ2
) (1)
for finding an ϵ-stationary point of a possibly non-convex smooth objective function (namely, a point at which the gradient is of norm ≤ ϵ). Here σ2 is the variance of the noise in the stochastic gradients, and τmax is the maximal possible delay, which is also needed to be known a-priori for properly tuning the SGD step size. Up to the τmax factor in the second term, this bound is identical to standard iteration bounds for stochastic non-convex SGD without delays [12]. While the bound in Eq. (1) is a significant improvement over previous art, it is still lacking in one important aspect: the dependence on the maximal delay could be excessively large in truly asynchronous environments, making the second term in the bound the dominant term. For example, in heterogeneous or massively distributed networks, the maximal delay is effectively determined by the single slowest (or less reliable) worker machine—which is precisely the issue with synchronous methods we set to address in the first place. Moreover, as Stich and Karimireddy [26] show, the step size used to achieve the bound in Eq. (1) could be as much as τmax-times smaller than that of without delays, which could severely impact performance in practice.

heading: 1.1 Contribution
text: We propose a new algorithm for stochastic optimization with asynchronous delayed updates, we call “Picky SGD,” that is significantly more robust than SGD, especially when the (empirical) distribution of delays is skewed or heavy-tailed and thus the maximal delay could be very large. For general smooth possibly non-convex objectives, our algorithm achieves a convergence bound of the form
𝑂
( σ2
ϵ4 + τavg ϵ2
) ,
where now τavg is the average delay in retrospect. This is a significant improvement over the bound in Eq. (1) whenever τavg ≪ τmax, which is indeed the case with heavy-tailed delay distributions. Moreover, Picky SGD is very efficient, extremely simple to implement, and does not require to know the average delay τavg ahead of time for optimal tuning. In fact, the algorithm only relies on a single additional hyper-parameter beyond the step-size. Notably, and in contrast to SGD as analyzed in previous work [26], our algorithm is able to employ a significantly larger effective step size, and thus one could expect it to perform well in practice compared to SGD. Indeed, we show in experiments that Picky SGD is able to converge quickly on large image classification tasks with a relatively high learning rate, even when very large delays are
introduced. In contrast, in the same setting, SGD needs to be configured with a substantially reduced step size to be able to converge at all, consequently performing poorly compared to our algorithm. Finally, we also address the case where 𝑓 is smooth and convex, in which we give a close variant of our algorithm with an iteration complexity bound of the form
𝑂
( σ2
ϵ2 + τavg ϵ ) for obtaining a point 𝑥 with 𝑓 (𝑥) − 𝑓 (𝑥∗) ≤ ϵ (where 𝑥∗ is a minimizer of 𝑓 over ℝ𝑑). Here as well, our rate matches precisely the one obtained by the state-of-the-art [26], but with the dependence on the maximal delay being replaced with the average delay. For consistency of presentation, we defer details on the convex case to the full version of the paper [? ] and focus here on our algorithm for non-convex optimization. Concurrently to this work, Aviv et al. [5] derived similar bounds that depend on the average delay. Compared to our contribution, their results are adaptive to the smoothness and noise parameters, but on the other hand, are restricted to convex functions and their algorithms are more elaborate and their implementation is more involved.

heading: 1.2 Additional related work
text: For general background on distributed asynchronous optimization and basic asymptotic convergence results, we refer to the classic book by Bertsekas and Tsitsiklis [6]. Since the influential work of Niu et al. [24], there has been significant interest in asynchronous algorithms in a related model where there is a delay in updating individual parameters in a shared parameter vector (e.g., [25, 19, 28, 17]). This is of course very different from our model, where steps use the full gradient vector in atomic, yet delayed, updates. Also related to our study is the literature on Local SGD (e.g., 27 and references therein), which is a distributed gradient method that perform several local (serial) gradient update steps before communicating with the parameter server or with other machines. Local SGD methods have become popular recently since they are used extensively in Federated Learning [20]. We note that the theoretical study in this line of work is mostly concerned with analyzing existing distributed variants of SGD used in practice, whereas we aim to develop and analyze new algorithmic tools to help with mitigating the effect of stale gradients in asynchronous optimization. A related yet orthogonal issue in distribution optimization, which we do not address here, is reducing the communication load between the workers and servers. One approach that was recently studied extensively is doing this by compressing gradient updates before they are transmitted over the network. We refer to [3, 14, 26] for further discussion and references.

heading: 2 Setup and Basic Definitions

heading: 2.1 Stochastic non-convex smooth optimization
text: We consider stochastic optimization of a β-smooth (not necessarily convex) non-negative function 𝑓 defined over the 𝑑-dimensional Euclidean space ℝ𝑑 . A function 𝑓 is said to be β-smooth if it is differentiable and its gradient operator is β-Lipschitz, that is, if ∥∇ 𝑓 (𝑥) − ∇ 𝑓 (𝑦)∥ ≤ β∥𝑥 − 𝑦∥ for all 𝑥, 𝑦 ∈ ℝ𝑑 . This in particular implies (e.g., [22]) that for all 𝑥, 𝑦 ∈ ℝ𝑑 ,
𝑓 (𝑦) ≤ 𝑓 (𝑥) + ∇ 𝑓 (𝑥) · (𝑦 − 𝑥) + β 2 ∥𝑦 − 𝑥∥2. (2)
We assume a stochastic first-order oracle access to 𝑓 ; namely, 𝑓 is endowed with a stochastic gradient oracle that given a point 𝑥 ∈ ℝ𝑑 returns a random vector ?̃?(𝑥), independent of all past randomization, such that 𝔼[?̃?(𝑥) | 𝑥] = ∇ 𝑓 (𝑥) and 𝔼[∥?̃?(𝑥) − ∇ 𝑓 (𝑥)∥2 | 𝑥] ≤ σ2 for some variance bound σ2 ≥ 0. In this setting, our goal is to find an ϵ-stationary point of 𝑓 , namely, a point 𝑥 ∈ ℝ𝑑 such that ∥∇ 𝑓 (𝑥)∥ ≤ ϵ, with as few samples of stochastic gradients as possible.

heading: 2.2 Asynchronous delay model
text: We consider an abstract setting where stochastic gradients (namely, outputs for invocations of the stochastic first-order oracle) are received asynchronously and are subject to arbitrary delays. The asynchronous model can be abstracted as follows. We assume that at each step 𝑡 of the optimization,
the algorithm obtains a pair (𝑥𝑡−𝑑𝑡 , 𝑔𝑡 ) where 𝑔𝑡 is a stochastic gradient at 𝑥𝑡−𝑑𝑡 with variance bounded by σ2; namely, 𝑔𝑡 is a random vector such that 𝔼𝑡𝑔𝑡 = ∇ 𝑓 (𝑥𝑡−𝑑𝑡 ) and 𝔼𝑡 ∥𝑔𝑡 − ∇ 𝑓 (𝑥𝑡−𝑑𝑡 )∥2 ≤ σ2 for some delay 0 ≤ 𝑑𝑡 < 𝑡. Here and throughout, 𝔼𝑡 [·] denotes the expectation conditioned on all randomness drawn before step 𝑡. After processing the received gradient update, the algorithm may query a new stochastic gradient at whatever point it chooses (the result of this query will be received with a delay, as above). Few remarks are in order: • We stress that the delays 𝑑1, 𝑑2, . . . are entirely arbitrary, possibly chosen by an adversary; in
particular, we do not assume they are sampled from a fixed stationary distribution. Nevertheless, we assume that the delays are independent of the randomness of the stochastic gradients (and of the internal randomness of the optimization algorithm, if any).1
• For simplicity, we assumed above that a stochastic gradient is received at every round 𝑡. This is almost without loss of generality:2 if at some round no feedback is observed, we may simply skip the round without affecting the rest of the optimization process (up to a re-indexing of the remaining rounds).
• Similarly, we will also assume that only a single gradient is obtained in each step; the scenario that multiple gradients arrive at the same step (as in mini-batched methods) can be simulated by several subsequent iterations in each of which a single gradient is processed.

heading: 3 The Picky SGD Algorithm
text: We are now ready to present our asynchronous stochastic optimization algorithm, which we call Picky SGD; see pseudo-code in Algorithm 1. The algorithm is essentially a variant of stochastic gradient descent, parameterized by a learning rate η as well as a target accuracy ϵ.
Algorithm 1: Picky SGD 1: input: learning rate η, target accuracy ϵ. 2: for 𝑡 = 1, . . . , 𝑇 do 3: receive delayed stochastic gradient 𝑔𝑡 and point 𝑥𝑡−𝑑𝑡 such that 𝔼𝑡 [𝑔𝑡 ] = ∇ 𝑓 (𝑥𝑡−𝑑𝑡 ). 4: if ∥𝑥𝑡 − 𝑥𝑡−𝑑𝑡 ∥ ≤ ϵ/(2β) then 5: update: 𝑥𝑡+1 = 𝑥𝑡 − η𝑔𝑡 . 6: else 7: pass: 𝑥𝑡+1 = 𝑥𝑡 . 8: end if 9: end for
Picky SGD maintains a sequence of iterates 𝑥1, . . . , 𝑥𝑇 . At step 𝑡, the algorithm receives a delayed stochastic gradient 𝑔𝑡 that was computed at an earlier iterate 𝑥𝑡−𝑑𝑡 (line 3). Then, in line 4, the algorithm tests whether ∥𝑥𝑡 − 𝑥𝑡−𝑑𝑡 ∥ ≤ ϵ/2β. Intuitively, this aims to verify whether the delayed (expected) gradient ∇ 𝑓 (𝑥𝑡−𝑑𝑡 ) is “similar” to the gradient ∇ 𝑓 (𝑥𝑡 ) at the current iterate 𝑥𝑡 ; due to the smoothness of 𝑓 , we expect that if 𝑥𝑡−𝑑𝑡 is close to 𝑥𝑡 , then also the corresponding gradients will be similar. If this condition holds true, the algorithm takes a gradient step using 𝑔𝑡 with step size η. Our main theoretical result is the following guarantee on the success of the algorithm. Theorem 1. Suppose that Algorithm 1 is initialized at 𝑥1 ∈ ℝ𝑑 with 𝑓 (𝑥1) ≤ 𝐹 and ran with
𝑇 ≥ 500β𝐹 ( σ2
ϵ4 + τ + 1 ϵ2
) , η =
1 4β
min { 1, ϵ2
σ2
} ,
where τ be the average delay, i.e., τ = (1/𝑇) ∑𝑇
𝑡=1 𝑑𝑡 . Then, with probability at least 1 2 , there is some

heading: 1 ≤ 𝑡 ≤ 𝑇 for which ∥∇ 𝑓 (𝑥𝑡 )∥ ≤ ϵ.
text: Observe that the optimal step size in Theorem 1 is independent of the average delay τ. This is important for two main reasons: (i) implementing the algorithm does not require knowledge about
1One can thus think of the sequence of delays as being fixed ahead of time by an oblivious adversary. 2We may, in principle, allow to query the stochastic gradient oracle even on rounds where no feedback is received, however this would be redundant in most reasonable instantiations of this model (e.g., in a parameter server architecture).
future, yet-to-be-seen delays; and (ii) even with very large delays, the algorithm can maintain a high effective step size. We note that the guarantee of Theorem 1 is slightly different from typical bounds in non-convex optimization (e.g., the bounds appearing in the previous work [14]): our result claims about the minimal gradient norm of any iterate rather than the average gradient norm over the iterates. Arguably, this difference does not represent a very strong limitation: the significance of convergence bounds in non-convex optimization is, in fact, in that they ensure that one of the iterates along the trajectory of the algorithm is indeed an approximate critical point, and the type of bound we establish is indeed sufficient to ensure exactly that. We further note that while the theorem above only guarantees a constant success probability, it is not hard to amplify this probability to an arbitrary 1 − δ simply by restarting the algorithm 𝑂 (log(1/δ)) times (with independent stochastic gradients); with high probability, one of the repetitions will be successful and run through a point with gradient norm ≤ ϵ, which would imply the guarantee in the theorem with probability at least 1 − δ.

heading: 4 Analysis
text: In this section we analyze Algorithm 1 and prove our main result. Throughout, we denote 𝑥 ′𝑡 = 𝑥𝑡−𝑑𝑡 and let 𝑁𝑡 denote the noise vector at step 𝑡, namely 𝑁𝑡 = 𝑔𝑡 − ∇ 𝑓 (𝑥 ′𝑡 ). Note that 𝔼[𝑁𝑡 | 𝑥𝑡 , 𝑥 ′𝑡 ] = 0 and 𝔼[∥𝑁𝑡 ∥2 | 𝑥𝑡 , 𝑥 ′𝑡 ] ≤ σ2, since the iterates 𝑥𝑡 , 𝑥 ′𝑡 are conditionally independent of the noise in 𝑔𝑡 as this gradient is obtained by the algorithm only at step 𝑡, after 𝑥𝑡 , 𝑥 ′𝑡 were determined. To prove Theorem 1, we will analyze a variant of the algorithm that will stop making updates once it finds a point with ∥∇ 𝑓 (𝑥)∥ ≤ ϵ (and eventually fails otherwise). That is, if ∥𝑥𝑡 − 𝑥 ′𝑡 ∥ > ϵ/2β or ∥∇ 𝑓 (𝑥𝑡 )∥ ≤ ϵ then 𝑥𝑡+1 = 𝑥𝑡 . Else, 𝑥𝑡+1 = 𝑥𝑡 − η𝑔𝑡 . This variant is impossible to implement (since it needs to compute the exact gradient at each step), but the guarantee of Theorem 1 is valid for this variant if and only if it is valid for the original algorithm: one encounters an ϵ-stationary point if and only if the other does so. First, we prove a simple technical lemma guaranteeing that whenever the algorithm takes a step, a large gradient norm implies a large decrease in function value. It is a variant of the classical “descent lemma,” adapted to the case where the gradient step is taken with respect to a gradient computed at a nearby point. Lemma 2. Fix 𝑥, 𝑥 ′ ∈ ℝ𝑑 with ∥𝑥 − 𝑥 ′∥ ≤ ϵ/2β and ∥∇ 𝑓 (𝑥 ′)∥ > ϵ. Let 𝑁 ∈ ℝ𝑑 be a random vector with 𝔼[𝑁 | 𝑥, 𝑥 ′] = 0 and 𝔼[∥𝑁 ∥2 | 𝑥, 𝑥 ′] ≤ σ2. Then,
𝔼[ 𝑓 (𝑥 − η(∇ 𝑓 (𝑥 ′) + 𝑁))] − 𝔼 𝑓 (𝑥) ≤ −η 2 𝔼∥∇ 𝑓 (𝑥 ′)∥2 + η
2β 2 (σ2 + 𝔼∥∇ 𝑓 (𝑥 ′)∥2).
In particular, for our choice of η, we have η
4 𝔼∥∇ 𝑓 (𝑥 ′)∥2 ≤ 𝔼 𝑓 (𝑥) − 𝔼[ 𝑓
( 𝑥 − η(∇ 𝑓 (𝑥 ′) + 𝑁) ) ] . (3)
Proof. Using the smoothness of 𝑓 (Eq. (2)), we have
𝑓 (𝑥 − η(∇ 𝑓 (𝑥 ′) + 𝑁)) − 𝑓 (𝑥) ≤ −η∇ 𝑓 (𝑥) · (∇ 𝑓 (𝑥 ′) + 𝑁) + 12η 2β∥∇ 𝑓 (𝑥 ′) + 𝑁 ∥2.
Taking expectation over 𝑁 conditioned on 𝑥, 𝑥 ′, we get
𝔼[ 𝑓 (𝑥 − η(∇ 𝑓 (𝑥 ′) + 𝑁)) − 𝑓 (𝑥) | 𝑥, 𝑥 ′] ≤ −η∇ 𝑓 (𝑥) · ∇ 𝑓 (𝑥 ′) + 12η
2β(∥∇ 𝑓 (𝑥 ′)∥2 + σ2) = −η∇ 𝑓 (𝑥 ′) · ∇ 𝑓 (𝑥 ′) − η∇ 𝑓 (𝑥 ′) · (∇ 𝑓 (𝑥) − ∇ 𝑓 (𝑥 ′)) + 12η
2β(∥∇ 𝑓 (𝑥 ′)∥2 + σ2) ≤ −η∥∇ 𝑓 (𝑥 ′)∥2 + ηβ∥∇ 𝑓 (𝑥 ′)∥∥𝑥 − 𝑥 ′∥ + 12η
2β(∥∇ 𝑓 (𝑥 ′)∥2 + σ2) = η(β∥∇ 𝑓 (𝑥 ′)∥∥𝑥 − 𝑥 ′∥ − ∥∇ 𝑓 (𝑥 ′)∥2) + 12η 2β(∥∇ 𝑓 (𝑥 ′)∥2 + σ2).
Since ϵ ≤ ∥∇ 𝑓 (𝑥 ′)∥ then
∥𝑥 − 𝑥 ′∥ ≤ ϵ 2β ≤ 1 2β ∥∇ 𝑓 (𝑥 ′)∥,
and we have 𝔼 [ 𝑓 (𝑥 − η(∇ 𝑓 (𝑥 ′) + 𝑁)) − 𝑓 (𝑥) | 𝑥, 𝑥 ′ ] ≤ −η
2 ∥∇ 𝑓 (𝑥 ′)∥2 + 12η 2β(σ2 + ∥∇ 𝑓 (𝑥 ′)∥2).
If ϵ ≥ σ then σ2 ≤ ∥∇ 𝑓 (𝑥 ′)∥2. This, with η = 1/4β, yields Eq. (3). If ϵ < σ and η = ϵ2/4σ2β, then η2 ≤ ϵ2/16σ2β2. Plugging that in instead, using ∥∇ 𝑓 (𝑥 ′)∥ ≥ ϵ, and taking expectations (with respect to 𝑥, 𝑥 ′) gets us Eq. (3). ■
We next introduce a bit of additional notation. We denote by 𝐼𝑡 the indicator of event that the algorithm performed an update at time 𝑡. Namely, 𝐼𝑡 = 𝐼 { ∥𝑥𝑡 − 𝑥 ′𝑡 ∥ ≤ ϵ/2β and ∥∇ 𝑓 (𝑥𝑡 )∥ > ϵ } .
Note that 𝐼𝑡 = 1 implies that ∥∇ 𝑓 (𝑥𝑠)∥ ≥ ϵ for all 𝑠 = 1, . . . , 𝑡. Further, we denote by ∆𝑡 = 𝑓 (𝑥𝑡 ) − 𝑓 (𝑥𝑡+1) the improvement at time 𝑡. Since 𝑓 is non-negative and 𝑓 (𝑥1) ≤ 𝐹, we have that for all 𝑡,
𝑡∑︁ 𝑖=1 ∆𝑖 = 𝑓 (𝑥1) − 𝑓 (𝑥𝑡+1) ≤ 𝐹.
Note that by Lemma 2 we have that 𝔼∆𝑡 ≥ 0. The rest of the proof is split into two cases: σ ≤ ϵ, and σ ≥ ϵ.

heading: 4.1 Case (i): σ ≤ ϵ
text: This regime is intuitively the “low noise” regime in which the standard deviation of the gradient noise, σ, is smaller than the desired accuracy ϵ. We prove the following. Lemma 3. Suppose that σ ≤ ϵ and the algorithm fails with probability ≥ 12 . Then 𝑇 ≤ 128β𝐹 (τ + 1)/ϵ2.
To prove the lemma above, we first show that the algorithm must make a significant number of updates, as shown by the following lemma. Lemma 4. If the algorithm fails, then the number of updates that it makes is at least 𝑇/4(τ + 1).
Proof. Consider 𝑈2τ, the number of steps 𝑡 for which the delay 𝑑𝑡 is at least 2τ. We must have 𝑈2τ ≤ 𝑇/2 (otherwise the total sum of delays exceeds τ𝑇 , contradicting the definition of τ). On the other hand, let 𝑘 be the number of updates that the algorithm makes. Let 𝑡1 < 𝑡2 < ... < 𝑡𝑘 be the steps in which an update is made. Denote 𝑡0 = 0 and 𝑡𝑘+1 = 𝑇 . Now, fix 𝑖 and consider the steps at times 𝑠𝑛 = 𝑡𝑖 + 𝑛 for 𝑛 ∈ [1, 2, . . . , 𝑡𝑖+1 − 𝑡𝑖 − 1]. In all those steps no update takes place and 𝑥𝑠𝑛 = 𝑥𝑡𝑖 . We must have 𝑑𝑠𝑛 > 𝑛 for all 𝑛 (otherwise 𝑥𝑡 = 𝑥𝑡−𝑑𝑡 for 𝑡 = 𝑠𝑛 and an update occurs). In particular we have that 𝑑𝑠𝑛 ≥ 2τ in at least 𝑡𝑖+1 − 𝑡𝑖 − 1 − 2τ steps in [𝑡𝑖 , 𝑡𝑖+1]. Hence,
𝑈2τ ≥ 𝑘−1∑︁ 𝑖=0 (𝑡𝑖+1 − 𝑡𝑖 − 1 − 2τ) = 𝑇 − 𝑘 (1 + 2τ).
Finally, it follows that 𝑇 − 𝑘 (1 + 2τ) ≤ 𝑇/2 which implies 𝑘 ≥ 𝑇4(τ+1) . ■
Given the lemma above, we prove Lemma 3 by showing that if the algorithm fails, it makes many updates in all of which we have ∥∇ 𝑓 (𝑥𝑡 )∥ > ϵ. By Lemma 2, this means that in the 𝑇 time steps of the algorithm, it must decrease the value of 𝑓 significantly. Since we start at a point in which 𝑓 (𝑥1) ≤ 𝐹, we must conclude that 𝑇 cannot be too large.
Proof of Lemma 3. Combining Eq. (3) with η = 1/(4β) and Lemma 4, we get that if the algorithm fails with probability ≥ 12 then
𝐹 ≥ 𝑇∑︁ 𝑡=1 𝔼∆𝑡 ≥ 1 16β 𝑇∑︁ 𝑡=1 𝔼[𝐼𝑡 ∥∇ 𝑓 (𝑥𝑡 )∥2] ≥ 1 16β 𝔼 [ 𝑇∑︁ 𝑡=1 𝐼𝑡 ∥∇ 𝑓 (𝑥𝑡 )∥2 ]
≥ 1 32β 𝔼 [ 𝑇∑︁ 𝑡=1 𝐼𝑡 ∥∇ 𝑓 (𝑥𝑡 )∥2 algorithm fails ] ≥ ϵ 2 32β 𝔼 [ 𝑇∑︁ 𝑡=1 𝐼𝑡 algorithm fails ] ≥ ϵ 2 32β 𝑇 4(τ + 1) .
This yields the lemma’s statement. ■

heading: 4.2 Case (ii): σ > ϵ
text: This is the “high noise” regime. For this case, we prove the following guarantee for the convergence of our algorithm. Lemma 5. Assume that σ > ϵ and the algorithm fails with probability ≥ 12 . Then,
𝑇∑︁ 𝑡=1 𝔼∆𝑡 ≥ 𝑇 500β min
{ ϵ2
τ , ϵ4 σ2
} .
In particular,
𝑇 ≤ 500β𝐹 ( τ
ϵ2 + σ
2
ϵ4
) .
This result is attained using the following observation. Consider the iterate of algorithm at time 𝑡, 𝑥𝑡 , and the point at which the gradient was computed 𝑥 ′𝑡 = 𝑥𝑡−𝑑𝑡 . We claim that if the algorithm has not decreased the function value sufficiently during the interval [𝑡 − 𝑑𝑡 , 𝑡 − 1], then it is likely to trigger a large decline in the function value at time 𝑡. Formally, either 𝔼∆𝑡 is large, or ∑𝑡−1 𝑖=𝑡−𝑑𝑡 𝔼∆𝑖 is large. To
show the claim, we first upper bound the distance ∥𝑥𝑡 − 𝑥 ′𝑡 ∥ in terms of ∑𝑡−1
𝑖=𝑡−𝑑𝑡 𝔼∆𝑖 , as shown by the following technical lemma. Lemma 6. For all 𝑡 and 𝑘 , it holds that
𝔼∥𝑥𝑡 − 𝑥𝑡+𝑘 ∥ ≤ √√ 1 β 𝑡+𝑘−1∑︁ 𝑖=𝑡 𝔼∆𝑖 + 4 ϵ 𝑡+𝑘−1∑︁ 𝑖=𝑡 𝔼∆𝑖 .
Proof. We have
𝔼∥𝑥𝑡 − 𝑥𝑡+𝑘 ∥ = η𝔼 𝑡+𝑘−1∑︁
𝑖=𝑡
𝐼𝑖 (∇ 𝑓 (𝑥 ′𝑖) + 𝑁𝑖) ≤ η𝔼 𝑡+𝑘−1∑︁ 𝑖=𝑡 𝐼𝑖∇ 𝑓 (𝑥 ′𝑖) + η𝔼 𝑡+𝑘−1∑︁ 𝑖=𝑡 𝐼𝑖𝑁𝑖
. We continue bounding the second term above as follows:
𝔼 𝑡+𝑘−1∑︁ 𝑖=𝑡 𝐼𝑖𝑁𝑖
≤ √√ 𝔼
𝑡+𝑘−1∑︁ 𝑖=𝑡 𝐼𝑖𝑁𝑖 2
= √√ 𝔼
𝑡+𝑘−1∑︁ 𝑖=𝑡 𝑡+𝑘−1∑︁ 𝑗=𝑡 𝐼𝑖 𝐼 𝑗𝑁𝑖 · 𝑁 𝑗
= √√ 𝔼
𝑡+𝑘−1∑︁ 𝑖=𝑡 𝐼𝑖 ∥𝑁𝑖 ∥2 (𝔼[𝑁𝑖 | 𝐼𝑖 , 𝐼 𝑗 , 𝑁 𝑗 ] = 0 for 𝑖 > 𝑗)
≤ σ √√ 𝔼
𝑡+𝑘−1∑︁ 𝑖=𝑡 𝐼𝑖
≤ σ ϵ
√√ 𝔼
𝑡+𝑘−1∑︁ 𝑖=𝑡 𝐼𝑖 ∥∇ 𝑓 (𝑥 ′𝑖)∥ 2 (∥∇ 𝑓 (𝑥 ′ 𝑖 )∥ ≥ ϵ when 𝐼𝑖 = 1)
≤ σ ϵ √√ 16σ2β ϵ2 𝑡+𝑘−1∑︁ 𝑖=𝑡 𝔼∆𝑖 (Eq. (3), η = ϵ2/4βσ2)
= 4σ2
ϵ2
√√ β
𝑡+𝑘−1∑︁ 𝑖=𝑡 𝔼∆𝑖
= 1 η √√ 1 β 𝑡+𝑘−1∑︁ 𝑖=𝑡 𝔼∆𝑖 , (η = ϵ2/4βσ2)
and
𝔼 𝑡+𝑘−1∑︁ 𝑖=𝑡 𝐼𝑖∇ 𝑓 (𝑥 ′𝑖) ≤ 𝑡+𝑘−1∑︁ 𝑖=𝑡 𝔼𝐼𝑖 ∥∇ 𝑓 (𝑥 ′𝑖)∥
≤ 1 ϵ 𝑡+𝑘−1∑︁ 𝑖=𝑡 𝔼𝐼𝑖 ∥∇ 𝑓 (𝑥 ′𝑖)∥2 (∥∇ 𝑓 (𝑥 ′𝑖)∥ ≥ ϵ when 𝐼𝑖 = 1)
≤ 4 ϵη 𝑡+𝑘−1∑︁ 𝑖=𝑡 𝔼∆𝑖 . (Eq. (3))
This completes the proof. ■ Given the lemma above, it is now clear that if ∑𝑡−1
𝑖=𝑡−𝑑𝑡 𝔼∆𝑖 is sufficiently small, then 𝔼∥𝑥𝑡 − 𝑥 ′ 𝑡 ∥ ≪ ϵ/β
which means that the algorithm is likely (with constant probability) to take a step at time 𝑡. This argument yields the following. Corollary 7. Assume that the algorithm fails with probability ≥ 12 . If ∑𝑡−1 𝑖=𝑡−𝑑𝑡 𝔼∆𝑖 < ϵ
2/125β then 𝔼∆𝑡 ≥ ϵ4/64σ2β. In particular,
𝔼∆𝑡 + 1 2τ 𝑡−1∑︁ 𝑖=𝑡−𝑑𝑡 𝔼∆𝑖 ≥ 1 250β min
{ ϵ2
τ , ϵ4 σ2
} .
Proof. If ∑𝑡−1
𝑖=𝑡−𝑑𝑖 𝔼∆𝑖 < ϵ 2/125β, then 𝔼∥𝑥𝑡−𝑑𝑡 − 𝑥𝑡 ∥ ≤ ϵ/8β by Lemma 6. By a Markov inequality,
with probability ≥ 34 , we have ∥𝑥𝑡−𝑑𝑡 − 𝑥𝑡 ∥ ≤ ϵ/2β. Since the probability that ∥∇ 𝑓 (𝑥𝑡−𝑑𝑡 )∥ > ϵ is at least 12 , we get that 𝔼𝐼𝑡 ≥ 1 4 . By Lemma 2 this implies that
𝔼∆𝑡 ≥ 1 4 · ϵ 2 · ϵ2 16σ2β = ϵ4 64σ2β ,
which yields our claim. ■
We now prove our main claim. We show that if the algorithm fails, then in all time steps in which 𝑑𝑡 ≤ 2τ (of which there are at least 𝑇/2), either the algorithm makes a substantial step, or it has made significant updates in the interval [𝑡 − 𝑑𝑡 , 𝑡 − 1]. In any case, the function value must necessarily decrease overall in the 𝑇 time steps of the algorithm, concluding that 𝑇 cannot be too large.
Proof of Lemma 5. We have, 𝑇∑︁ 𝑡=1 𝔼∆𝑡 ≥ ∑︁ 𝑡:𝑑𝑡 ≤2τ 1 2τ 𝑡−1∑︁ 𝑖=𝑡−𝑑𝑡 𝔼∆𝑖 .
Hence, using Corollary 7, 𝑇∑︁ 𝑡=1 𝔼∆𝑡 ≥ 1 2 ∑︁ 𝑡:𝑑𝑡 ≤2τ ( 𝔼∆𝑡 + 1 2τ 𝑡−1∑︁ 𝑖=𝑡−𝑑𝑡 𝔼∆𝑖 ) ≥
{𝑡 : 𝑑𝑡 ≤ 2τ} 1250β min{ ϵ2τ , ϵ4σ2 } ≥ 𝑇
2 1 250β min
{ ϵ2
τ , ϵ4 σ2 } = 𝑇
500β min
{ ϵ2
τ , ϵ4 σ2
} ,
where we used Markov’s inequality to show that |{𝑡 : 𝑑𝑡 ≤ 2τ}| ≥ 12𝑇 . ■

heading: 4.3 Concluding the proof
text: Proof of Theorem 1. In the case σ ≤ ϵ, Lemma 3 implies that if 𝑇 > 128β𝐹 (τ + 1)/ϵ2 then the algorithms succeeds with probability greater than 1/2, which yields the theorem in this case. Similarly, Lemma 5 gives our claim in the case when σ > ϵ. ■

heading: 5 Experiments
text: To illustrate the robustness and efficacy of Picky SGD, we present a comparison between the performance of SGD versus Picky SGD under various delay distributions. In particular, we show that Picky SGD requires significantly less iterations to reach a fixes goal and is more robust to varying delay distributions.

heading: 5.1 Setup
text: The main goal of our experimental setup is to be reproducible. For that end, the experimentation is done in two phases. First, we perform a simulation to determine the delay 𝑑𝑡 at each iteration without actually computing any gradients:3 this is done by simulating 𝑁 concurrent worker threads sharing and collectively advancing a global iteration number, where each worker repeatedly records the current global iteration number 𝑡start, waits a random amount of time from a prescribed Poisson distribution, then records the new global iteration number 𝑡 = 𝑡end and the difference 𝑑𝑡 = 𝑡end − 𝑡start, and increases the global iteration number. This information (a delay schedule) is calculated once for each tested scheme (differing in the number of workers and random distribution, as detailed below), and is stored for use in the second phase. In the second phase of the experiments, the algorithms SGD and Picky SGD are executed for each delay schedule. Here, at every iteration the gradient is computed (if needed) and is kept until its usage as dictated by the schedule (and then applied at the appropriate global iteration number). As a result of this configuration, we get a fully reproducible set of experiments, where the algorithms performance may be compared as they are executed over identical delay series of identical statistical properties. We created four different delay schedules: A baseline schedule (A) using 𝑁 = 10 workers and sampling the simulated wait from a Poisson distribution (this schedule serves to compare Picky SGD and SGD in a setting of relatively small delay variance) and schedules (B) (C) and (D) all using 𝑁 = 75 workers and sampling the simulated wait from bi-modal mixtures of Poisson distributions of similar mean but increasing variance respectively.4 See Figure 2 in the the full version of the paper [? ] for an illustration of the delay distributions of the four delay schedules used. All training is performed on the standard CIFAR-10 dataset [15] using a ResNet56 with 9 blocks model [13] and implemented in TensorFlow [1]. We compare Picky SGD (Algorithm 1) to the SGD algorithm which unconditionally updates the state 𝑥𝑡 given the stochastic delayed gradient 𝑔𝑡 (recall that 𝑔𝑡 is the stochastic gradient at state 𝑥𝑡−𝑑𝑡 ). For both algorithms, instead of a constant learning rate η we use a piecewise-linear learning rate schedule as follows: we consider a baseline η0 piecewise-linear learning rate schedule5 that achieves optimal performance in a synchronous distributed optimization setting (that is, for 𝑑𝑡 ≡ 0)6 and search (for each of the four delay schedules and each algorithm – to compensate for the effects of delays) for the best multiple of the baseline rate and the best first rate-change point. Alternatively, we also used a cosine decay learning rate schedule (with the duration of the decay as meta parameters). Another meta-parameter we optimize is the threshold ϵ/(2β) in line 4 of Picky SGD. Batch size 64 was used throughout the experiments. Note that although use chose the threshold value ϵ/2β by an exhaustive search, in practice, a good choice can be found by logging the distance values during a typical execution and choosing a high percentile value. See the full version of the paper [? ] for more details.
3Note that up to the training data ordering a computation of 𝑇 steps of Picky SGD or SGD is uniquely determined by the starting state 𝑥1 and the sequence {𝑡 − 𝑑𝑡 }𝑡=1...𝑇 .
4See the the full version of the paper [? ] for specific parameter values and implementation details. 5With rate changes at three achieved accuracy points 0.93, 0.98, and 0.99. 6This is also the best performance achievable in an asynchronous setting.

heading: 5.2 Results
text: The accuracy trajectory for the best performing combination of parameters of each algorithm for each of the four delay schedules is shown in Fig. 1 and summarized in Table 1. Clearly, Picky SGD significantly outperforms SGD in terms of the final accuracy and the number of epochs it takes to achieve it. We also emphasize that the generalization performance (that is, the evaluation accuracy as related to the training accuracy) was not observed to vary across delay schedules or the applied algorithms (see e.g., Fig. 4 in the the full version of the paper [? ]), and that the nature of the results is even more pronounced when using the alternative cosine decay learning rate schedule (see Fig. 5 in the the full version of the paper [? ]). Specific details of the meta parameters used, and additional performance figures are reported in the full version of the paper [? ].

heading: 5.3 Discussion
text: We first observe that while the number of epochs it takes Picky SGD to reach the target accuracy mark is almost the same across the delay schedules (ranging from 288 to 344), SGD requires significantly more epochs to attain the target accuracy (ranging from 350 up to 466 for the highest variance delay schedule)—this is consistent with the average-delay bound dependence of Picky SGD (as stated in Theorem 1) compared to the max-delay bound dependence of SGD. Furthermore, the best baseline learning rate multiplier meta-parameter for Picky SGD is the same (0.2) across all high-variance delay schedules, while the respective meta parameter for SGD is significantly smaller (0.05) and sometimes varying, explaining the need for more steps to reach the target and evidence of Picky SGD superior robustness.

heading: Acknowledgements
text: AD is partially supported by the Israeli Science Foundation (ISF) grant no. 2258/19. TK is partially supported by the Israeli Science Foundation (ISF) grant no. 2549/19, by the Len Blavatnik and the Blavatnik Family foundation, and by the Yandex Initiative in Machine Learning.

Review summaries:
This paper improves the iteration complexity to get $\epsilon$ small gradients for smooth non-convex functions using SGD with delayed gradients. This setting, in particular, captures asynchronous distributed SGD. The previously known best algorithm [1] accumulates error feedback to provide a $\sigma^2/\epsilon^4 + \tau_{max}/\epsilon^2$ rate, where $\tau_{max}$ is the maximum delay in any update. This can be catastrophic in cross-device federated learning with edge devices, where a device gets disconnected for multiple rounds, and then sends its gradient for the update.   This paper deals with such a scenario in a very simple manner, by measuring the distance between the current iterate and the old iterate at which the gradient was computed, and making the update only if it is smaller than a threshold value. This algorithm is aptly called *picky SGD*. Surprisingly, this removes the need for error feedback and is able to guarantee a $\sigma^2/\epsilon^4 + \tau_{avg}/\epsilon^2$ rate where $\tau_{avg}$ is the average delay over all the rounds. Picky SGD with a different threshold also offers state-of-the-art convergence rates for convex smooth functions, with the dependence on the maximum delay changed to the one on average delay.  The authors also study the effect of sampling the delays from distributions with different levels of variance in their experiments. The results demonstrate that picky SGD requires comparable or fewer steps to reach a fixed level of accuracy than vanilla asynchronous SGD. This is attributed to the fact that picky SGD tolerates a much larger step size than asynchronous SGD.    **References**  [1] Stich, Sebastian U., and Sai Praneeth Karimireddy. "The error-feedback framework: Better rates for SGD with delayed gradients and compressed communication." arXiv preprint arXiv:1909.05350 (2019).

This paper gives an alternative to SGD, called picky SGD, for stochastic optimization with delays. While Stich & Karimireddy proved that SGD was robust to delays as long as the maximum delay is controled, here picky SGD is proved to be robust to much larger delays, as it requires only the average delay (and not the maximum delay) to be controled. This is supported by both theoretical convergence to a stationary point (in the non-convex case) or to a minimum (in the convex case) and with simulations with artificially introduced delays. 

This paper considers the problem of distributed stochastic asynchronous optimization, in which workers send delayed stochastic gradients to a server. The focus is on non-convex smooth objectives, although results for convex objectives are provided in the supplementary material.   Standard asynchronous distributed SGD algorithms are first discussed, and the authors highlight that their step-size (and thus, also their iteration complexity) depends on a bound over the maximum possible delay over all iterations. Instead, this paper introduces an algorithm that does not require to know this bound, and only depends on the average observed delay.  This algorithm is called Picky-SGD, and its main difference with standard delayed SGD is that the update is only performed if the delayed iterate is close enough to the current iterate. Theoretical results are then given to show that Picky-SGD achieves the standard complexity results for these problems, but replacing the maximum delay by the average one. The proof is then presented in details, and a set of experiments is given to illustrate the superiority of Picky-SGD over delayed SGD.