Asynchronous Stochastic Optimization Robust to Arbitrary Delays

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 . 1 Introduction . 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 . 1.1 Contribution . 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 . 1.2 Additional related work . 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 . 2 Setup and Basic Definitions . 2.1 Stochastic non-convex smooth optimization . 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 .

Summary: 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.