Low-Precision Stochastic Gradient Langevin Dynamics

1 INTRODUCTION . Low-precision optimization has become increasingly popular in reducing compute and memory costs of training deep neural networks ( DNNs ) . It uses fewer bits to represent numbers in model parameters , activations and gradients and thus can drastically lower resource demands ( Gupta et al. , 2015 ; Zhou et al. , 2016 ; De Sa et al. , 2017 ; Li et al. , 2017 ) . Prior work has shown that using 8-bit numbers in training DNNs achieves about 4× latency speed ups and memory reduction compared to 32-bit numbers on a variety of deep learning tasks ( Sun et al. , 2019 ; Yang et al. , 2019 ; Wang et al. , 2018b ; Banner et al. , 2018 ) . As datasets and architectures grow rapidly , performing lowprecision optimization enables training large-scale DNNs efficiently and enables many applications on different hardware and platforms . Despite the impressive progress in low-precision optimization , low-precision sampling remains largely unexplored . However , we believe stochastic gradient Markov chain Monte Carlo ( SGMCMC ) methods ( Welling & Teh , 2011 ; Chen et al. , 2014 ; Ma et al. , 2015 ) are particularly suited for low-precision arithmetic because of their intrinsic robustness to system noise . In particular : ( 1 ) SGMCMC explores weight space instead of converging to a single point , thus it should not require precise weights or gradients ; ( 2 ) SGMCMC even adds noise to the system to encourage exploration and so is naturally more tolerant to quantization noise ; ( 3 ) SGMCMC performs Bayesian model averaging during testing using an ensemble of models , which enables coarse representations of individual models to be compensated by the overall model average ( Zhu et al. , 2019 ) . SGMCMC is particularly compelling in Bayesian deep learning due to its ability to characterize complex and multimodal DNN posteriors , providing state-of-the-art generalization accuracy and calibration ( BNNs ) ( Zhang et al. , 2020 ; Li et al. , 2016 ; Gan et al. , 2016 ) . Moreover , low-precision approaches are especially appealing in this setting , where at test time we must store samples from a posterior over millions of parameters , and perform multiple forward passes through the corresponding models , which incurs significant memory and computational expenses . In this paper , we provide the first comprehensive study of low-precision Stochastic Gradient Langevin Dynamics ( SGLD ) ( Welling & Teh , 2011 ) . We start by analyzing the theoretical convergence of low-precision SGLD on strongly log-concave distributions , proving that SGLD with full-precision gradient accumulators ( SGLDLP-F ) converges to the target distribution within a distance that is asymptotically smaller than the distance between the SGD estimation and the optimum . Surprisingly , we find that SGLD with low-precision gradient accumulators ( SGLDLP-L ) can diverge arbitrarily far away from the target distribution with small stepsizes . We solve this issue by developing a new quantization function that preserves the correct variance in each update step . We prove that when using this new quantization function , SGLD converges to the target distribution within a bounded level of accuracy even with small stepsizes . We illustrate the behavior of different low-precision SGLD variants in Figure 1 . Empirically we demonstrate low-precision SGLD on deep learning tasks , including CIFAR-10 , CIFAR-100 , and ImageNet with a ResNet-18 , and the IMDB dataset with a LSTM . We show that our variance-corrected quantizer significantly improves the performance of low-precision SGLD . Moreover , the improvement of SGLD over SGD is larger in low-precision than in full-precision , demonstrating the promise of low-precision stochastic sampling . 2 RELATED WORK . Most work on accelerating SGLD training has focused on distributed learning with synchronous or asynchronous communication ( Ahn et al. , 2014 ; Chen et al. , 2016 ; Li et al. , 2019 ) . Another direction to reduce training costs is to shorten training time by accelerating the convergence using either variance reduction techniques ( Dubey et al. , 2016 ; Baker et al. , 2019 ) or a cyclical learning rate schedule ( Zhang et al. , 2020 ) . To speed up testing , distillation techniques are often used to save both compute and memory by transferring the knowledge of an ensemble of models to a single model ( Korattikara et al. , 2015 ; Wang et al. , 2018a ) . Low-precision computation has become one of the most common approaches to reduce latency and memory consumption in deep learning and is widely supported on new emerging chips including CPUs , GPUs and TPUs ( Micikevicius et al. , 2017 ; Krishnamoorthi , 2018 ; Esser et al. , 2019 ) . Two main directions to improve low-precision training include developing new number formats ( Sun et al. , 2019 ; 2020 ) and studying mixed precision schemes ( Courbariaux et al. , 2015 ; Zhou et al. , 2016 ; Banner et al. , 2018 ) . Recently , one line of work studies applying Bayesian framework to learn a deterministic quantized neural network ( Soudry et al. , 2014 ; Cheng et al. , 2015 ; Achterhold et al. , 2018 ; van Baalen et al. , 2020 ; Meng et al. , 2020 ) . Despite impressive progress of low-precision deep learning , few work takes advantage of it for Bayesian neural networks ( BNNs ) . Su et al . ( 2019 ) proposes a method to train binarized variational BNNs and Cai et al . ( 2018 ) develops an efficient hardware for training low-precision variational BNNs . The only work on low-precision MCMC known to us is Ferianc et al . ( 2021 ) , which directly applies post-training quantization in optimization ( Jacob et al. , 2018 ) to quantize BNNs trained by Stochastic Gradient Hamiltonian Monte Carlo ( Chen et al. , 2014 ) . In contrast , we study training low-precision models by SGLD from scratch , which can reduce both training and testing costs . 3 PRELIMINARIES . 3.1 STOCHASTIC GRADIENT LANGEVIN DYNAMICS . In the Bayesian setting , given some dataset D , a model with parameters θ , and a prior p ( θ ) , we are interested in sampling from the posterior p ( θ|D ) ∝ exp ( −U ( θ ) ) , where the energy function is U ( θ ) = − ∑ x∈D log p ( x|θ ) − log p ( θ ) . When the dataset is large , the cost of computing a sum over the entire dataset is expensive . Stochastic Gradient Langevin Dynamics ( SGLD ) reduces the cost by using a stochastic gradient estimation ∇Ũ , an unbiased estimator of ∇U usually based on a subset of the dataset D. Specifically , SGLD updates the parameter θ in the ( k + 1 ) -th step following the rule θk+1 = θk − α∇Ũ ( θk ) + √ 2αξk+1 , where α is the stepsize and ξ is a standard Gaussian noise . Compared to the SGD update , the only difference is the additional Gaussian noise . This close connection makes SGLD convenient to implement on existing deep learning tasks for which SGD is the typical learning algorithm . 3.2 LOW-PRECISION TRAINING . We study training a low-precision model by SGLD from scratch , to reduce both training and testing costs . We use fixed point in our theoretical analysis and empirical demonstration , which is a common number type for reducing DNNs costs ( Gupta et al. , 2015 ; Lin et al. , 2016 ; Li et al. , 2017 ; Yang et al. , 2019 ) . We additionally use block floating point , another common number representation , for empirical evaluations on deep learning tasks ( Song et al. , 2018 ) . In order to train a network using low-precision numbers , we need a quantization function Q to convert a real-valued number into a rounded number . Such functions include deterministic rounding and stochastic rounding . Suppose that we use W bits to represent numbers with F of those W bits to represent the fractional part . Then the quantization gap ∆ = 2−F is the distance between consecutive representable numbers . The lower and upper bounds of the representable numbers are l = −2W−F−1 and u = 2W−F−1 − 2−F , respectively . The deterministic rounding quantizes a number to its nearest representable neighbour : Qd ( θ ) = sign ( θ ) · clip ( ∆ ⌊ |θ| ∆ + 1 2 ⌋ , l , u ) , where clip ( x , l , u ) = max ( min ( x , u ) , l ) . Stochastic rounding quantizes a number with probability : Qs ( θ ) = { clip ( ∆ ⌊ θ ∆ ⌋ , l , u ) , w.p . ⌈ θ ∆ ⌉ − θ∆ clip ( ∆ ⌈ θ ∆ ⌉ , l , u ) , w.p . 1− ( ⌈ θ ∆ ⌉ − θ∆ ) . Qs is unbiased , which means E [ Qs ( θ ) ] = θ . It is more favorable than Qd in practice since it can preserve gradient information when the gradient update is smaller than the quantization gap ( Gupta et al. , 2015 ; Courbariaux et al. , 2016 ) . In what follows , we use Qw and ∆w to denote the weights ’ quantizer and quantization gap , Qg and ∆g to denote gradients ’ quantizer and quantization gap . Please note that Qg means quantizing the error and gradient in each layer in backward propagation ( Wu et al. , 2018 ; Yang et al. , 2019 ) . Besides the quantization function , we also need to decide what and where to quantize . There are two common choices depending on whether we store an additional copy of full-precision weights . Fullprecision gradient accumulators use a full-precision weight buffer to accumulate gradient updates and only quantize weights before computing gradients . SGD with full-precision gradient accumulators ( SGDLP-F ) updates the weights as θk+1 = θk − αQg ( ∇Ũ ( Qw ( θk ) ) ) . Gradient accumulators are frequently updated during training , therefore it will be ideal to also represent it in low-precision to further reduce the costs . The update of SGD with low-precision gradient accumulators ( SGDLP-L ) is θk+1 = Qw ( θk − αQg ( ∇Ũ ( θk ) ) ) . Both full- and low-precision gradient accumulators have been widely used in low-precision training . Low-precision gradient accumulators are cheaper and full-precision gradient accumulators generally have better performance because of more precisely accumulating gradient updates ( Courbariaux et al. , 2015 ; Li et al. , 2017 ) . 4 LOW-PRECISION SGLD . In this section , we first study the convergence bound of low-precision SGLD with full-precision gradient accumulators ( SGLDLP-F ) on strongly log-concave distributions and show that SGLDLP-F converges to the target distribution within a distance that is asymptotically smaller than that between the SGD estimation and the optimum . Next we analyze low-precision SGLD with low-precision gradient accumulators ( SGLDLP-L ) under the same setup and prove that SGLDLP-L can diverge arbitrarily far away from the target distribution with a small stepsize , which is typically required by SGLD to reduce asymptotic bias . Finally , we solve this problem by developing a variance-corrected quantization function and further prove that our SGLDLP-L converges with small stepsizes . 4.1 FULL-PRECISION GRADIENT ACCUMULATORS . Due to the similarity between SGLD and SGD , applying low-precision to SGLD is straightforward . Similar to SGDLP-F , we can do low-precision SGLD with full-precision gradient accumulators ( SGLDLP-F ) as the following . θk+1 = θk − αQg ( ∇Ũ ( Qw ( θk ) ) ) + √ 2αξk+1 . ( 1 ) We now prove that SGLDLP-F will converge to the target distribution given small stepsizes and large number of iterations . Our convergence analysis of low-precision SGLD is based on ( Dalalyan & Karagulyan , 2019 ) , where the target distribution is assumed smooth and strongly log-concave . We additionally assume the energy function has Lipschitz Hessian following recent work in low-precision optimization ( Yang et al. , 2019 ) . Specifically , the energy function U satisfies the following , U ( θ ) − U ( θ′ ) −∇U ( θ′ ) ᵀ ( θ − θ′ ) ≥ ( m/2 ) ‖θ − θ′‖22 , ‖∇U ( θ ) −∇U ( θ′ ) ‖2 ≤M ‖θ − θ′‖2 , ‖∇2U ( θ ) −∇2U ( θ′ ) ‖2 ≤ Ψ‖θ − θ′‖2 , ∀θ , θ′ ∈ Rd for some positive constants m , M and Ψ . We further assume the variance of the stochastic gradient is bounded by E [ ‖∇Ũ ( θ ) − ∇U ( θ ) ‖22 ] ≤ κ2 for some constant κ . For simplicity , we consider SGLD with a constant stepsize α . We measure the convergence of SGLD in terms of 2-Wasserstein distance . We use stochastic rounding for both weight and gradient quantizers as it is generally better than deterministic rounding and has also been used in previous low-precision theoretical analysis ( Li et al. , 2017 ; Yang et al. , 2019 ) . Theorem 1 . We run SGLDLP-F under the above assumptions and with a constant stepsize α ≤ 2/ ( m+M ) . Let π be the target distribution , µK be the distribution obtained after K iterations and µ0 be the initial distribution , then W2 ( µK , π ) ≤ ( 1− αm ) KW2 ( µ0 , π ) + 1.65 ( M/m ) ( αd ) 1/2 + min ( Ψ∆2wd 4m , M∆w √ d 2m ) + √ ( ∆2g +M 2∆2w ) αd+ 4ακ 2 4m . There are three main takeways of this theorem : First , it shows that SGLDLP-F converges to the accuracy floor min ( Ψ∆2wd 4m , M∆w √ d 2m ) which depends on the quantization gap ∆w given large enough K and small enough α . Second , our convergence bound is O ( ∆2w ) which is better than the distance between the SGD estimation and the optimum O ( ∆w ) ( Yang et al. , 2019 ) ( see Appendix D for details ) . Third , if we further assume the energy function is quadratic , that is Ψ = 0 , then SGLDLP-F converges to the target distribution asymptotically . This is similar to SGDLP-F on quadratic function which converges to the optimum asymptotically ( Li et al. , 2017 ) .

Summary: This paper discusses using the stochastic gradient Langevin dynamics with low-precision implementation. The authors present convergence results in Wasserstein distance norm for three implementation cases, which distinguish specific low-precision thresholding schemes and different quantization techniques. The authors also report simulation outcome of proposed low-precision SGLD algorithms and show their convergence in actual application.