Understanding Top-k Sparsification in Distributed Deep Learning

1 INTRODUCTION . Training large-scale deep neural networks ( DNNs ) generally exploits distributed synchronous stochastic gradient descent ( SGD ) optimization algorithms to reduce the overall training time . Let P be the number of workers in a distributed setting , and x ∈ Rd denotes the model parameters with d dimensions . At the t-th iteration , distributed synchronous SGD updates the model parameters by xt+1 = xt − ηt 1 P P∑ p=1 gpt , ( 1 ) where gpt ∈ Rd is the stochastic gradient with its locally selected data for the loss function fp ( x ) : Rd → R and ηt is the learning rate . The aggregation of d-dimension gradients from P workers requires a communication complexity of O ( d ) in terms of communication traffics1 , which generally limits the system scalability . Gradient sparsification ( Strom , 2015 ; Dryden et al. , 2016 ; Aji & Heafield , 2017 ; Chen et al. , 2018 ; Lin et al. , 2018 ) is a promising technique for distributed SGD , which can significantly reduce the communication traffic while reserving the model convergence . In gradient sparsification , a compressor Compk is applied on each worker to locally select k , k ≤ d , gradients for aggregation and Compk ∈ { Topk , Randk } ( Stich et al. , 2018 ) . Compk ( g p t ) ∈ Rd zeros out ( d− k ) elements of gpt and keeps k elements unchanged . The zeroed-out d− k elements are stored as residual pt for the next iteration . Formally , the model parameters are updated by xt+1 = xt − ηt 1 P P∑ p=1 Compk ( g p t + p t ) and p t+1 = g p t + p t − Compk ( g p t + p t ) , ( 2 ) where pt ∈ Rd and p 0 = 0 . In theory , distributed SGD with gradient sparsification ( e.g. , Topk , Randk and any other k-contraction operators ) with error compensation has been proved to have the 1The ring-based AllReduce collective can achieve the bandwidth optimal performance that is not related to the number of workers , but there exist latency terms that will increase with increased number of workers . same order of convergence rate as vanilla SGD for both convex and non-convex problems if the number of iterations is large ( Wangni et al. , 2018 ; Stich et al. , 2018 ; Alistarh et al. , 2018 ; Jiang & Agrawal , 2018 ; Karimireddy et al. , 2019 ; Tang et al. , 2019 ; Zheng et al. , 2019 ) . The convergence rates are derived with a key contraction property of the sparsification operator Compk ( Topk or Randk ) ( Stich et al. , 2018 ; Alistarh et al. , 2018 ) , that is EC [ ‖x− Compk ( x ) ‖2 ] ≤ ( 1− k/d ) ‖x‖2 , ∀x ∈ Rd , ( 3 ) where EC is the expectation taking on the compressor and ‖ · ‖ is the ` 2-norm . For any x ∈ Rd , Topk ( x ) ∈ Rd selects the top k largest elements ( in terms of the absolute value ) of x with corresponding indices and sets other d− k elements to zeros ; while Randk ( x ) ∈ Rd randomly ( in a uniform distribution ) selects k elements from xwith corresponding indices and other d−k elements are zeros . It is obvious that ‖x− Topk ( x ) ‖2 ≤ ‖x− Randk ( x ) ‖2 and ER [ ‖x− Randk ( x ) ‖2 ] = ( 1− k/d ) ‖x‖2 . ( 4 ) Existing studies use the same error estimate for both Topk and Randk in distributed SGD by exploiting the properties of ( 4 ) , which can not differentiate the convergence behavior of two operators . In practice , however , TopK-SGD has a much faster convergence speed ( in term of iterations ) than SGD with Randk ( RandK-SGD ) as empirically shown in ( Stich et al. , 2018 ) . We also compare the convergence performance between TopK-SGD and RandK-SGD on a 16-worker distributed setting with three popular convolutional neural networks ( VGG-16 ( Simonyan & Zisserman , 2014 ) , ResNet-20 and ResNet-50 ( He et al. , 2016 ) ) . Our results are shown in Fig . 1 . We observe that TopK-SGD achieves very similar performance to the original distributed SGD ( Dense-SGD ) , while RandK-SGD has much slower convergence than TopK-SGD . RandK-SGD even can not converge on ImageNet . Therefore , though existing studies show that TopK-SGD and RandK-SGD have the same convergence bound , their theoretical results can not explain the performance gap between TopK-SGD and RandK-SGD . Even some work ( Karimireddy et al. , 2019 ; Tang et al. , 2019 ) exploits δ ≤ 1 to replace k/d in ( 3 ) , they also fail to identify exact δ to distinguish Topk and Randk . In this paper , we dive into the details of the Topk operator in distributed SGD when training DNNs and provide a tighter bound than inequality ( 3 ) to explain the good convergence performance of TopK-SGD . The observation of gradients with Topk sparsification further enables us to propose a new computational-efficient selection algorithm for gradient which preserves the convergence property . Our contributions are summarized as follows . Contributions . ( 1 ) We empirically study the details of local stochastic gradients and observe that the coordinates of gradient follow bell shaped distributions through extensive experiments . ( 2 ) The bell shaped distribution enables us to intuitively explain that Topk should have a much tighter bound than Randk , and we exploit the distribution property to formulate how Topk outperforms Randk . ( 3 ) We design and implement an approximate top-k selection algorithm2 , which is much more efficient than existing top-k selection algorithms on GPUs . As compared with the existing sampling-based approximate top-k selection algorithm , we improve the scaling efficiency by 12-50 % on our 16-GPU cluster . 2Our system implementation will be made open-source after the review process . 2 RELATED WORK . Gradient Quantization . In distributed training of neural networks , the communicated gradients can be quantized to low-bit precision ( e.g. , 16-bit ( Micikevicius et al. , 2018 ; Jia et al. , 2018 ) , 3-bit ( Wen et al. , 2017 ) , 2.8-bit ( Alistarh et al. , 2017 ; Karimireddy et al. , 2019 ) and even 1-bit ( Seide et al. , 2014 ; Strom , 2015 ) ) while preserving nearly consistent convergence performance with the full precision ( 32-bit ) counterpart . Recently general frameworks of gradient quantization with error compensation are proposed to generalize the theoretical results of low-bit communication ( Wu et al. , 2018 ; Jiang & Agrawal , 2018 ; Karimireddy et al. , 2019 ; Tang et al. , 2019 ; Haddadpour et al. , 2019 ) . However , the quantization method can only reduce the communication traffic in 32× ( i.e. , 1-bit vs. 32-bit ) , and it could not be enough for large-scale models or low-bandwidth network connections . Gradient Sparsification . Compared to gradient quantization , gradient sparsification is a much more promising communication traffic reduction technique as it can sparsify up to three orders of magnitude gradients be zero with little impact on the model convergence ( Strom , 2015 ; Dryden et al. , 2016 ; Aji & Heafield , 2017 ; Chen et al. , 2018 ; Lin et al. , 2018 ; Shi et al. , 2019a ) . Due to the much success of gradient sparsification ( e.g. , Top-k sparsification ) in significantly reducing the communication traffic ( Lin et al. , 2018 ; Sun et al. , 2019 ) , much recent work tries to build theoretical guarantees for SGD with gradient specification ( Wangni et al. , 2018 ; Stich et al. , 2018 ; Alistarh et al. , 2018 ; Jiang & Agrawal , 2018 ; Shi et al. , 2019b ; Karimireddy et al. , 2019 ; Tang et al. , 2019 ) . These theoretical frameworks try to generalize the sparsification operator with the bound of inequality ( 3 ) to derive the convergence results for SGD with gradient sparsification . However , the existing analysis fails to go insight into the details of gradient sparsification of Topk which could have better convergence than other compression operators ( e.g. , Randk ) . Gradient Distribution3 . Glorot & Bengio ( 2010 ) study the distribution of activation values of DNNs and also their corresponding gradients . They empirically showed that back-propagated gradients have Gaussian-like distributions , which helps understand the difficulty of training deep neural networks . A similar plot is shown in ( Micikevicius et al. , 2018 ) , where the distribution of gradients helps analyze if the 16-bit representation of gradients would be overflow or underflow . These work has demonstrated that the gradients during training are likely located near zeros . We extend the similar studies on the gradient distribution for TopK-SGD . 3 STUDY ON STOCHASTIC GRADIENTS . 3.1 GRADIENT DISTRIBUTION . In previous gradient sparsification studies ( Strom , 2015 ; Dryden et al. , 2016 ; Aji & Heafield , 2017 ; Chen et al. , 2018 ; Lin et al. , 2018 ) , the basic rule of sparsification is to select “ significant ” elements of the gradients because they contribute more to the updates . The Topk operator selects the exact local top-k elements of gradients so that it achieves nearly consistent convergence performance with Dense-SGD . Therefore , we would like to understand what is the difference between “ significant ” elements of the gradients and randomly selected ones . We conduct extensive experiments to study the gradient distributions on three areas of deep learning applications , including image classification , language modeling , and speech recognition . The selected models are : 1 ) Feed-forward Neural Networks ( FNNs ) . An FNN with three hidden fully connected layers ( FNN-3 ) on the MNIST ( LeCun , 1998 ) data set . 2 ) Convolutional Neural Networks ( CNNs ) . LeNet-5 ( LeCun et al. , 2015 ) on MNIST , ResNet-20 ( He et al. , 2016 ) and VGG-16 ( Simonyan & Zisserman , 2014 ) on CIFAR10 ( Krizhevsky et al. , 2010 ) . And 3 ) Recurrent Neural Networks ( RNNs ) . Long Short Term Memory networks ( LSTMs ) on the Penn Treebank ( PTB ) ( Marcus et al. , 1993 ) and the AN4 ( Acero , 1990 ) data sets . For PTB , we adopt a 2-layer LSTM model ( LSTM-PTB ) with 1500 hidden units per layer , and for AN4 , we use a 5-layer LSTM model ( LSTM-AN4 ) with 800 hidden units per layer . The details of the experimental settings are shown in Table 1 . As the compression operator is applied on the gradients , we first measure the distributions of the gradient ’ s elements ( histograms ) on Dense-SGD . The results demonstrate the similar shapes as ( Glorot & Bengio , 2010 ) , while ours covers various applications ( refer to Appendix A.2 ) . Our interest is on TopK-SGD to check if 3The distribution we discussed in this paper is over coordinates on a particular vector ( e.g. , activation outputs , full gradients ) . gradients distributions perverse the same properties as Dense-SGD . During the training process of TopK-SGD ( k = 0.001d for a d-dimension model ) , we measure the histograms of local gradients accumulated with the residuals ( i.e. , upt = g p t + p t ) . The histograms of u 1 t with different t on different models are shown in Fig . 2 , where we only show the gradients from the first worker as different workers have very close gradient distributions . The corresponding cumulative distributions are presented in Appendix A.1 . It is seen that different models have different shapes on the accumu- lated gradients , but one common feature is that most coordinates of ut are close to zero . Compared to the full gradient SGD ( Appendix A.2 ) , TopK-SGD shows wider distributions , which could be mainly caused by residual accumulation . When selecting top-k largest values ( in terms of absolute values ) from ut , the selected values should be located at the left and right sides on the histograms . Therefore , performing Topk on ut should generate a vector whose ` 2-norm is very close to that of ut , that is ‖Topk ( ut ) ‖2 / ‖ut‖2 . The intuitive result inspires us to formulate how much close of ‖Topk ( ut ) ‖2 to ‖ut‖2 . Specifically , we would like to derive a variable γ ≤ ( 1 − k/d ) such that ‖ut − Topk ( ut ) ‖2 ≤ γ‖ut‖2 holds .

Summary: Top-k algorithm is a gradient sparsification method which gains its popularity to train deep networks due to its high compression ratio. However due to its computation overhead, it is not efficient on GPUs. This paper performs empirical study on the distribution of the gradients when training various deep networks, and provides a Gaussian approximation analysis to improve the convergence of the top-k algorithm. Further, the paper proposes to use a Gaussian-k algorithm to perform similar gradient sparsification with a much lower computational cost without losing much convergence accuracy compared to Dense-SGD.