Neural gradients are near-lognormal: improved quantized  and sparse training

1 INTRODUCTION . Neural gradients are used in the training process of deep networks to backpropagate the error-gradient throughout the model , thus allowing to compute the required weight updates . As these neural gradients are needed for a substantial ratio of the underlying computations ( about 23 ) , compressing them can alleviate data-throughput requirements and accelerate the training process . Compression of neural gradients reduce the memory footprint for the intermediate calculation and the bandwidth of data transfer inside the HW accelerator . Moreover , in term of distributed training in model parallelism the neural gradients are one of the main bottlenecks that need to be transferred between devices ( Rong et al. , 2020 ; Gupta et al. , 2020 ) . Many previous works ( Banner et al. , 2019 ; Fang et al. , 2020 ) compress tensors such as weights and activations by approximating their distributions using an analytically tractable density . These works often assume a bell-shaped distribution such as Gaussian or Laplace distributions , which have been reported to fail for neural gradients ( Ye et al. , 2019 ) . One key observation in this paper is that neural gradient distributions are heavy-tailed , fundamentally different from the light-tailed distributions of weights and activations . Further statistical and distributional tests reveal gradient magnitudes follow a lognormal distribution . Adopting this lognormal observation , our paper suggests two main applications — quantization and pruning , used to reduce the computational and memory burden of neural gradients . To tackle these challenges , we first formalize the problems and find closed-form expressions that enable us to predict the optimal quantization and pruning policies . These measures are easy to use and depend only on the estimated lognormal parameters . ∗Equal contribution . In Figure 1 we summarize these applications and their derivation . The first application uses the lognormal prior to enabling low-precision floating-point ( FP ) quantization of the gradients . Here we optimize two tasks . The first task is to find a partition between mantissa and exponent bit-widths that minimizes quantization noise for a given n bit FP gradient representation . The second task is to scale these gradients so that they would be properly represented within a limited dynamic range ( distance between the maximum and minimum that FP format can represent ) . We provide useful insights that make empirically-based heuristics such as loss scaling ( Micikevicius et al. , 2018 ) a more grounded approach with a theoretical basis . Optimizing both tasks we obtain state-of-the-art results for FP quantization of the neural gradients . The second application performs accurate and predictable stochastic pruning of gradients on the fly , which results in two state-of-the-art pruning schemes . The first translates the desired sparsity level into an accurate threshold , and the other enables combined use of different sparsity levels at different layers ( heterogeneous sparsity ) . 2 RELATED WORK . Quantization and pruning of neural networks have seen a tremendous amount of works ( e.g. , Nahshan et al . ( 2019 ) ; Choi et al . ( 2018 ) ; Baskin et al . ( 2018 ) ; Louizos et al . ( 2018 ) ; Frankle & Carbin ( 2018 ) ; Cambier et al . ( 2020 ) ) aiming to reduce both bandwidth and memory footprint , as well as computation time . Most of these methods focus on the quantization or pruning of the weights / activations in the forward path ( Banner et al. , 2019 ; Nahshan et al. , 2019 ) or the weight gradients ( Bernstein et al. , 2018 ; Alistarh et al. , 2016 ) in the backward path . So far , neural gradients got less attention . Some of these methods ( Banner et al. , 2019 ; Ye et al. , 2019 ; Fang et al. , 2020 ) use a systemic and rigorous statistical approach to optimize various distortion measures . For example , Banner et al . ( 2019 ) used the normal distributional assumption ( of weights and activations ) to analytically minimize the mean-squared quantization error . Our work follows a similar line to rigorously optimize similar performance measures for quantization and pruning of gradient distributions , which are different from that of the weights and activations . Gradient Quantization . While a lot of research focused on the quantization of weights and activations for inference ( Krishnamoorthi , 2018 ; Choi et al. , 2018 ; Jain et al. , 2020 ) , there were also major advances in quantization during training , many of them faced difficulty trying to represent the high dynamic range of the gradients ( Banner et al. , 2018 ; Wu et al. , 2018 ) . Cambier et al . ( 2020 ) suggests keeping for each tensor shift and scale in full precision numbers to make them fit in FP8 format dynamic range . Mantissa vs exponent allocation of the available bits has proven crucial in deep learning workloads , where for example , BF16 ( 1-8-7 : sign-exponent-mantissa ) has shown greater success compared to traditional FP16 ( 1-5-10 ) format due to wider dynamic range ( Henry et al. , 2019 ; Kalamkar et al. , 2019 ) . Research over required format and trade-offs in exponent versus mantissa is on-going with growing interest over lower precision representations such as FP8 . Some works have explored using different FP8 formats : Wang et al . ( 2018 ) has used ( 1-5-2 ) , while ( Sun et al. , 2019 ) suggests using one type of FP8 format for the forward ( 1-4-3 ) , and a different FP8 format ( 1-5-2 ) for the backward pass , after empirically assessing the different possible formats . Additionally , with the growing usage of FP16 mixed-precision training , researchers and practitioners faced the need to use loss scaling ( Micikevicius et al. , 2018 ) , adjusting the tensor distribution to avoid over/under-flow . This procedure required , in some cases , an intensive parameter search and was not guaranteed to succeed in all cases . Gradient pruning . Focusing on the computationally-intensive back-propagation , ” meProp ” ( Sun et al. , 2017 ) prunes the K smallest absolute-valued entries of the neural gradients on the fly , using the top-k algorithm . Works following it , replaced pruning with quantization to induce sparsity ( Wiedemann et al. , 2020 ) , or used top-k pruning as well on the copies of weights and activations used in back-prop ( Aamir Raihan & Aamodt , 2020 ) . Ye et al . ( 2019 ) , inspired by conventional unbiased estimators like stochastic rounding , suggested ” stochastic pruning ” , reaching higher sparsity levels . Yet , the authors assumed the gradients are normally distributed , leading to incorrect estimation of the threshold and a large difference between the required sparsity and the one obtained . As we shall see later , using the correct statistical distribution model is essential to determine the proper threshold . 3 NEURAL GRADIENTS DISTRIBUTION . Many prior works ( Banner et al. , 2019 ; Bernstein et al. , 2018 ) take the assumption that tensor data e.g. , weights ( W ) and activations ( A ) is sampled from a Gaussian distribution . Recently , Ye et al . ( 2019 ) ; Wiedemann et al . ( 2020 ) used the same assumption for the distribution of neural gradients∇A . In the section , we discuss this assumption . We show that neural gradients are better approximated by lognormal distributions , i.e. , the gradient logarithm values are normally distributed , as opposed to the gradients themselves . In Fig . 2 , we plot the histogram of neural gradient magnitudes at linear and log scales in one layer of ResNet18 - ImageNet dataset . At a linear scale , the distribution has a few gradients of huge magnitude and a great many gradients of small magnitudes ( Fig 2a ) . Plotting the histogram on a logarithmic scale reveals a distribution close to a symmetric normal ( Gaussian ) distribution ( Fig 2b ) . This is the hallmark of the lognormal distribution . Finally , when plotting the theoretical quantiles of the normal distribution against the quantiles of the gradient distribution ( Q-Q plot ) , we see that the points follow a strongly nonlinear pattern in Fig 2c , suggesting that the data is not distributed as a standard normal distribution . Note that in the Q-Q plot for lognormal distribution ( Fig 2d ) , almost all points lie on a straight line . We further estimate the goodness of fit of the neural gradients to normal and lognormal distributions . To that end , we measure the static distance ( largest vertical line ) between the cumulative distribution function ( CDF ) of the empirically observed distribution and the CDF of the reference distribution ( also known as Kolmogorov-Smirnov test ( Smirnov , 1948 ) ) . For each model and dataset in Table 1 , we calculate the average ( across all layers ) of the static distance to normal and lognormal distributions . The analysis is performed on the absolute value of the gradients , excluding the zero-valued entries . Note that lognormal distribution gets a better fit . Additional statistical distributions in Appendix A.1 . 4 APPLICATION I - OPTIMAL FLOATING-POINT QUANTIZATION . Floating-point ( FP ) representations can cover a very wide dynamic range with a relatively small number of digits . The mantissa-exponent tradeoff control the balance between the dynamic range of numbers that can be represented ( exponent ) to the precision of these numbers ( mantissa ) . This dynamic range is especially important for the heavy-tailed distributions that characterize neural gradients . In this section , we study the optimal characteristics of the FP quantizer . 4.1 PROBLEM FORMULATION . We can decompose any positive real value x ∈ R+ as follows : x = 2ln x = M∈ [ 1,2 ) ︷ ︸︸ ︷ 2ln x−bln xc ·2 E∈Z︷ ︸︸ ︷ blnxc , ( 1 ) where M ∈ [ 1 , 2 ) is the mantissa and E ∈ Z the exponent . Given N bits , we allocate 1 bit for the sign and seek the optimal allocation of n1 bits to M and n2 bits to E , such that n1 + n2 = N − 1 . Accordingly , we define the quantized xq as : xq = { 2Emax E ≥ Emax Mq · 2E −Emax ≤ E ≤ Emax 0 E ≤ −Emax ( 2 ) where Emax = 2n2−1 and Mq is the quantized mantissa with the range [ 1 , 2 ) divided into 2n1 quantization levels with a spacing of ∆ = 12n1 . Finally , we measure the relative error between the FP number xq and the real number x , which is simply the difference between the two numbers divided by the real number ( Widrow & Kollár , 2008 ) : η ( n1 , n2 ) = ∣∣∣∣xq − xx ∣∣∣∣ ( 3 ) 4.2 ANALYTICAL DERIVATION OF THE RELATIVE ERROR . We assume that x ∼ Lognormal ( µ , σ2 ) . Note thatE = blnxc ≈ lnx ∼ N ( µ , σ2 ) . In Appendix A.4 we split the range into three parts according to E : ( i ) −Emax ≤ E ≤ Emax ; ( ii ) E ≥ Emax ; ( iii ) E ≤ −Emax , and calculate the expected contribution for each term . A closed-form formula for the expected relative error could be obtained as follows : E [ η ( n1 , n2 ) ] = 2Φ ( Emax σ ) − 1 8 · ln ( 2 ) · ( 2n1 − 1 ) + 2Emax−1e σ2 ln2 ( 2 ) 2 ( erf ( σ ln 2√ 2 + Emax√ 2σ ) − 1 ) − 1 2 erf ( Emax√ 2σ ) + 3 2 − Φ ( Emax σ ) ( 4 ) where Φ ( x ) is the CDF of N ( 0 , 1 ) . In Fig . 3a we show that analytical results stated by Eq . ( 4 ) are in good agreement with simulations for FP8 with various number of exponent bits . Simulations were obtained by quantizing 10,000 values , generated from a lognormal distribution with σ = 1 , 3 , 5 .

Summary: This work proposed a very interesting idea that the back-propagated errors have log-normal distributions. The authors could extend this intriguing observation into computation efficient algorithms; reduced-precision floating-point quantization or the pruning of the back-prop error, which are very interesting. The authors also provide details of derivations as well as the data analysis along with several small ideas (in Appendix) to strongly support their ideas.