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 investigate posterior sampling of weight-quantized neural networks with SGLD. The paper gives the convergence of SGLD under Wasserstein distance with quantized weights and gradients. It also considers the situation when the accumulator is also in low precision, and proposes a variance correction strategy to deal with the additional quantization noise.