Gap-Aware Mitigation of Gradient Staleness

1 INTRODUCTION . The steady growth of deep neural networks over the years has made it impractical to train them from scratch on a single worker ( i.e. , computational device ) . Distributing the computations over several workers can drastically reduce the training time . However , due to the sequential nature of the widely used stochastic gradient descent ( SGD ) method , distributing the process is not an easy task . Synchronous SGD ( SSGD ) is the most common method used to distribute the learning process across multiple workers . Several recent works ( Mikami et al. , 2018 ; Ying et al. , 2018 ; Yamazaki et al. , 2019 ; Goyal et al. , 2017 ) have shown that SSGD can achieve large speedups while maintaining high accuracy . The major drawback of SSGD is that its speed is confined to the slowest worker in every iteration . This shortcoming is magnified in non-dedicated1 environments such as cloud computing . For this reason , all the above mentioned works were forced to use homogeneous workers in a dedicated network , which serves to reduce the variance in the workers ’ iteration times . Unlike cloud computing , dedicated networks are expensive and therefore not available to most users . In asynchronous SGD ( ASGD ) , each worker communicates independently of the others , thereby addressing the major drawback of SSGD . ASGD enjoys linear speedup in terms of the number of workers , even on non-dedicated networks . This makes ASGD a potentially better alternative to SSGD when using cloud computing . Unfortunately , ASGD also has a significant weakness known as gradient staleness ; the gradients used to update the parameter server ’ s ( master ) parameters are ∗Equal contribution . 1An environment of computation nodes who are not specifically optimized to work together often based on older parameters and therefore are inaccurate . Prior works have shown that gradient staleness severely hinders the convergence process by reaching reduced final accuracy ( Chen et al. , 2016 ; Cui et al. , 2016 ) . Mitliagkas et al . ( 2016 ) showed that gradient staleness also induces implicit momentum , thus the momentum coefficient γ must be decayed when scaling up the number of workers . Most research works measure the gradient staleness of a worker according to the delay : the number of master updates since the worker began calculating the stochastic gradient g , until g is used to update the master . To overcome gradient staleness , Zhang et al . ( 2015b ) proposed Staleness-Aware ( SA ) , which penalizes the step size of stale gradients linearly to their delay . This method was later embraced by other works ( Jiang et al. , 2017 ; Hardy et al. , 2017 ; Giladi et al. , 2019 ) and is currently the common method for penalizing stale gradients . Unfortunately , this method suffers from a degradation of the final accuracy , especially when scaling up the number of workers . In Section 4.1 , we show that the main reason for this degradation is the over-penalization and under-penalization caused by SA . SSGD and ASGD rely on hyperparameter tuning for every different number of workers ( Shallue et al. , 2018 ) . Tuning is extremely time-consuming , thus avoiding it is beneficial , whenever possible . Our contribution : To mitigate gradient staleness while minimizing the degradation of final accuracy , we define a measure of gradient staleness we refer to as the Gap . The Gap is based on the difference between the parameters used to calculate the gradient and the parameters on which the gradient is applied . We propose a new method called Gap-Aware ( GA ) that penalizes the step size of stale gradients linearly to their Gap , while eliminating the over-penalization or under-penalization of SA . No new hyperparameters are introduced using the GA method . • We show that GA out-performs SA , especially as the number of workers scales up . • We prove that the convergence rate of the GA-ASGD algorithm with a non-convex loss function is consistent with SGD : O ( 1√ BK ) where K is the total number of steps and B is the batch size . • We show that penalizing the gradient itself rather than the step-size , eliminates under-penalization . • Our results suggest that GA can be used without re-tuning the hyperparameters . • As opposed to conclusions by Mitliagkas et al . ( 2016 ) , we show that applying momentum in an asynchronous environment is advantageous ( using GA ) , even when multiple workers are used . • We combine GA with Adam ( Kingma & Ba , 2015 ) ( Adam-GA ) , and show that Adam-GA achieves almost two orders of magnitude better perplexity than Adam or Adam-SA ( which combines Adam with SA ) using several workers on the Transformers-XL model . Our results establish GA as a superior gradient-penalizing option to SA and suggest that using GA is a preferable alternative to SSGD in non-dedicated networks such as cloud computing , even when scaling to large numbers of workers . To validate our claims , we performed experiments on the CIFAR10 , CIFAR100 ( Krizhevsky , 2012 ) , ImageNet ( Russakovsky et al. , 2015 ) , and WikiText-103 ( Merity et al. , 2016 ) datasets , using several state-of-the-art architectures . A version of GA has reached 72.18 % final test accuracy on the ImageNet dataset using 128 simulated asynchronous workers . As far as we know , this is the largest number of asynchronous workers reported to converge on ImageNet . 2 RELATED WORK . Eliminating gradient staleness is a challenging task and several papers suggested techniques to reduce its detrimental effects . Zheng et al . ( 2017 ) proposed DC-ASGD , which uses a Taylor expansion to mitigate the gradient staleness . EASGD ( Zhang et al. , 2015a ) uses a center force to pull the workers ’ parameters toward the master ’ s parameters . Both DC-ASGD and EASGD achieve high accuracy on small numbers of workers , but fall short when trained on large clusters . Chan & Lane ( 2014 ) proposed penalizing stale gradients by reducing their size and thus limiting their effect on the learning process . They suggest decaying the learning rate exponentially to the delay ; this makes the step size arbitrarily small when the amount of workers grows , virtually ceasing the learning process . As part of their convergence analysis , Dutta et al . ( 2018 ) suggest a penalizing method that is linear to the norm between the master and worker ’ s parameters . However , their method introduces another hyperparameter , which requires additional time to tune . As opposed to other methods , GA performs well even when the number of workers is large , without introducing new hyperparameters . Gupta et al . ( 2016 ) as well as Aji & Heafield ( 2019 ) suggest collecting several gradients before updating the master to reduce the effects of gradient staleness . Wen et al . ( 2017 ) propose minimizing the size of the gradients to reduce communication times . GA is orthogonal to both of these approaches . 3 BACKGROUND . The goal of an optimization procedure is to minimize f ( θ ) , where f is a smooth , but not necessarily convex , objective function ( a.k.a . loss ) and the vector θ ∈ Rd is the model ’ s parameters : θ∗ = arg min θ∈Rd f ( θ ) : = Eξ [ F ( θ ; ξ ) ] ( 1 ) where ξ ∈ Ξ is a random variable from Ξ , the entire set of training samples Ξ = { 1 , 2 , · · · , M } . F ( θ ; ξ ) is the stochastic loss function with respect to the training sample indexed by ξ. SGD is commonly the workhorse behind the optimization of deep neural networks . Denoting η as the learning rate and k as the step number , SGD ’ s iterative update rule is : θk+1 = θk − ηk∇f ( θk ) . We denote Xk as the variable X at the kth step , where X is any variable . Momentum Momentum ( Polyak , 1964 ) is a widely adopted optimization technique due to its accelerated convergence and oscillation reduction ( Sutskever et al. , 2013 ) . Instead of simply using the gradient , the momentum iterative update rule2 uses an exponentially-weighted moving average of gradients called the update vector : vk+1 = γvk +∇f ( θk ) . The update rule is : θk+1 = θk − ηkvk+1 . Nesterov ’ s Accelerated Gradient ( NAG ) ( Nesterov , 1983 ) is a well-used variation of momentum that has been proven to achieve quadratic speedup in convergence rate compared to SGD . 4 ASYNCHRONOUS SGD ( ASGD ) . We consider the commonly used ASGD , which operates with a parameter-server ( master ) , used to keep the model ’ s most up-to-date parameters . Each worker maintains a replica of the model ’ s parameters . The workers run in parallel and synchronize with the master independently from each other at the beginning of each batch iteration . We denote τk as the delay at the kth step . The worker and master algorithms are given by Algorithm 1 and 2 , respectively , where B is the batch-size and ξk , b denotes the bth sample in the batch sampled at the kth iteration . Gradient staleness appears when a gradient gk is computed on parameters θk−τk but applied to different parameters θk . Algorithm 1 Momentum-ASGD : worker i Always do : Receive parameters θk−τk from the master Get B training samples ξk−τk , [ 1 ... B ] Compute gradient : gik ← ∑B b=1 ∇F ( θk−τk ; ξk−τk , b ) B Send gik to the master Algorithm 2 Momentum-ASGD : master For k = 1 ... K do : Receive gradient gik from worker i Update momentum vk+1 ← γvk + gik Update master ’ s weights : θk+1 ← θk − ηk · vk+1 Send θk+1 to worker i 4.1 STALENESS-AWARE . Zhang et al . ( 2015b ) proposed a gradient penalization method called Staleness-Aware ( SA ) . SA aims to reduce the effects of gradient staleness by dividing the step size by its corresponding τ ( Algorithm 3 ) . The worker algorithm remains unchanged . The two ideas behind SA is that stale gradients should be penalized to reduce their impact and that gradient staleness scales up with τ . SA successfully mitigates the gradient staleness when N is small . Although commonly used , SA can potentially over-penalize as well as under-penalize stale gradients as we show below . 2We consider the version of momentum without dampening Algorithm 3 Staleness-Aware : master Initialize an iteration array : iter = [ 0 ] ∗N For k = 1 ... K do : Receive gradient gik from worker i Calculate worker i ’ s delay τk ← k − iter [ i ] Update momentum vk+1 ← γvk + gik Update master θk+1 ← θk − ηkτk vk+1 Send θk+1 to worker i Save current iteration iter [ i ] ← k Algorithm 4 Gap-Aware : master For k = 1 ... K do : Receive gradient gik from worker i Calculate Gap : Gk = |θk−θk−τk | C + 1 d Update momentum vk+1 ← γvk + ( 1 Gk ) gik Update master θk+1 ← θk − ηkvk+1 Save and send current parameters θk+1 to worker i Over-Penalizing Let us assume that at some step k , after τ master updates we get θk = θk−τ just as the gradient calculated on θk−τ is applied . This means there is no gradient staleness for the next gradient update since it was computed using the same parameters on which it is applied . Unfortunately , the delay remains τ > 0 , thereby causing over-penalization when SA is used . Additionally , τ scales linearly with N , which dramatically reduces η when N is large . Consequently , on large numbers of workers , the convergence rate of SA is sluggish and its accuracy plummets . Under-Penalizing The SA method doesn ’ t take into account that when using momentum , the update step also contains past gradients . To emphasize the importance of this issue , let ’ s examine a fictional example : assume that some gradient gk is very stale ( τk is large ) . Following the SA technique , the update rule is : θk+1 = θk − ηkτk · vk+1 = θk − ηk τk · ( γvk + gk ) . This means the stale gradient gk is indeed penalized by being multiplied by ηkτk , which is small ; thus gk doesn ’ t change the parameters much . We further assume the next iteration is very fast ( τk+1 is small ) . The next update will be : θk+2 = θk+1 − ηk+1τk+1 · ( γ 2vk + γgk + gk+1 ) . The stale gradient gk is multiplied by γ ηk+1 τk+1 , which is large ( assuming ηk ≈ ηk+1 ) . Despite the fact that gk was stale , it still has a significant impact on the learning process . In other words , gk is under-penalized . To eliminate this possibility we penalize the stale gradient itself rather than the learning rate . Using this method , the staleness of each gradient is accounted for within the update vector v .

Summary: The paper introduces a new variant of asynchronous SGD, GA-ASGD, for distributed training. The goal is to mitigate the gradient staleness issue caused by asynchronously applying gradients to an old version of parameters. Prior work addresses this issue by penalizing the learning step of a worker linearly to its missed updates since getting a replica of parameters. However, that approach does not consider the differences between the old and new versions of parameters, which can cause over-penalization and under-penalization. The main contribution of this paper is to introduce a new way of measuring weight staleness and to explore the idea of penalizing the gradient itself to mitigate the staleness issue. The paper does a good job of discussing how GA can be applied to existing optimizers such as Adam. It performs empirical studies on ImageNet and Transformers-XL to conclude that GA-ASGD outperforms ASGD and the prior staleness aware approach. It also demonstrates the scalability of GA by scaling up to training with 128 asynchronous workers.