LOW-PRECISION STOCHASTIC GRADIENT LANGEVIN DYNAMICS
Year: 2021

Abstract:
Low-precision optimization is widely used to accelerate large-scale deep learning. Despite providing better uncertainty estimation and generalization, sampling methods remain mostly unexplored in this space. In this paper, we provide the first study of low-precision Stochastic Gradient Langevin Dynamics (SGLD), arguing that it is particularly suited to low-bit arithmetic due to its intrinsic ability to handle system noise. We prove the convergence of low-precision SGLD on strongly log-concave distributions, showing that with full-precision gradient accumulators, SGLD is more robust to quantization error than SGD; however, with low-precision gradient accumulators, SGLD can diverge arbitrarily far from the target distribution with small stepsizes. To remedy this issue, we develop a new quantization function that preserves the correct variance in each update step. We demonstrate that the resulting low-precision SGLD algorithm is comparable to full-precision SGLD and outperforms low-precision SGD on deep learning tasks.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 PRELIMINARIES

heading: 3.1 STOCHASTIC GRADIENT LANGEVIN DYNAMICS
text: 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.

heading: 3.2 LOW-PRECISION TRAINING
text: 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).

heading: 4 LOW-PRECISION SGLD
text: 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.

heading: 4.1 FULL-PRECISION GRADIENT ACCUMULATORS
text: 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).

heading: 4.2 LOW-PRECISION GRADIENT ACCUMULATORS
text: To further reduce the costs, we apply low-precision gradient accumulators to SGLD which we denote SGLDLP-L. Mimicking the update of SGDLP-L, the update rule of SGLDLP-L is
θk+1 = Qw ( θk − αQg ( ∇Ũ(θk) ) + √ 2αξk+1 ) . (2)
Surprisingly, while we can prove a convergence result for SGLDLP-L, our theory suggests that it can diverge arbitrarily from the target distribution with small stepsizes. Theorem 2. We run SGLDLP-L under the same assumptions as in Theorem 1. Then
W2(µK , π) ≤ (1− αm)KW2(µ0, π) + 1.65(M/m)(αd)1/2 + min
( Ψ∆2wd
4m , M∆w
√ d
2m
)
+
√ (α∆2g + α −1∆2w)d+ 4ακ 2
4m + ( (1− αm)K + 1 ) ∆w√d 2 .
This theorem suggests that as the stepsize α decreases, W2 distance between the stainary distribution of SGLDLP-L and the target distribution increases. When α decreases, either our bound becomes loose or SGLDLP-L diverges from the target distribution. We empirically test SGLDLP-L on a standard Gaussian distribution in Figure 1, showing that the reason is the later. We use 8-bit fixed point and assign 2 of them to represent the integer part. We can see that SGLDLP-F always converges to the target distribution with different stepsizes whereas SGLDLP-L diverges from the target distribution and the divergence increases when the stepsize decreases.
One may choose a stepsize that minimizes the above W2 distance to avoid divergence, however, getting this optimal stepsize is generally difficult as the constants are unknown in practice. Moreover, enabling a small stepsize in SGLD is desirable, since SGLD needs a small stepsize to reduce asymptotic bias of the posterior approximation (Welling & Teh, 2011).

heading: 4.3 VARIANCE-CORRECTED QUANTIZATION
text: To approach correcting the problems with naı̈ve SGLDLP-L, we first need to identify the source of the issue. We start by showing that the reason is the variance of each dimension of θk+1 becomes larger due to using stochastic rounding as the low-precision gradient accumulators. Specifically, given the stochastic gradient ∇Ũ , the update of full-precision SGLD can be viewed as sampling from a Gaussian distribution for each dimension i
θk+1,i ∼ N ( θk,i − α∇Ũ(θk)i, 2α ) , for i = 1, · · · , d.
Using stochastic rounding as the weight quantizer and the gradient quantizer in SGLDLP-L gives us E [θk+1,i] = E [ Qs ( θk,i − αQs ( ∇Ũ(θk) ) i + √ 2αξk+1,i )] = θk,i − α∇Ũ(θk)i which keeps the same mean of θk+1 as in full-precision. However the variance of θk+1,i is larger than needed. If we ignore the variance from Qg and the stochastic gradient, since they are present and have been shown to work well in SGLDLP-F, the variance of θk+1,i is
Var [θk+1,i] = E [ Var [ Qs ( θk,i − α∇U(θk)i + √ 2αξk+1,i )∣∣∣ξk+1,i]] + Var [ E [ Qs ( θk,i − α∇U(θk)i + √ 2αξk+1,i
)∣∣∣ξk+1,i]] =
∆2w 4 χk+1,i + 2α.
Algorithm 1 Variance-Corrected Low-Precision SGLD (VC SGLDLP-L). given: Stepsize α, number of training iterationsK, gradient quantizerQg , deterministic rounding Qd, stochastic rounding Qs and quantization gap of weights ∆w. for k = 1 : K do
update θk+1 ← Qvc ( θk − αQg ( ∇Ũ(θk) ) , 2α,∆w, Q d, Qs )
end for output: samples {θk}
where χk+1,i ∈ [0, 1] depends on the distance of θk+1,i to the discrete grid. Empirically, from Figure 1, we can see that naı̈ve SGLDLP-L gives the correct mean estimate but the wrong variance estimate, while both are estimated correctly in SGLD which adds the proper amount of noise in each update. This validates our intuition that using stochastic rounding in low-precision gradient accumulators adds more variance than needed leading to an inaccurate variance estimate.
To address this issue, we propose a new quantization function Qvc in Algorithm 2. Qvc always guarantees the correct mean, E [θk+1,i] = θk,i − α∇Ũ(θk)i, and guarantees the correct variance Var [θk+1,i] = 2α most of the time except when v < vs. However that case rarely happens in practice, because the stepsize has to be extremely small. The main idea of Qvc is to directly sample from the discrete weight space instead of quantizing a real-valued Gaussian sample. It does so considering two cases: when the Gaussian noise is larger than the largest possible stochastic rounding variance ∆2/4,Qvc first adds a portion of the Gaussian noise and uses a sample from the weight grid to make up the remaining; in the other situation, Qvc directly samples from the weight grid to achieve the target variance. Although Qvc does not provide a sample following a Gaussian distribution (as it is discrete), we show that this does not affect the performance in both theory and practice.
We now prove that SGLDLP-L using Qvc, denoting VC SGLDLP-L, converges to the target distribution up to a certain accuracy level with small stepsizes. Theorem 3. We run VC SGLDLP-L as in Algorithm 1. Besides the same assumptions in Theorem 1, we further assume the gradient is bounded E [∥∥∥Qg(∇Ũ(θk))∥∥∥ 1 ] ≤ G. Let v0 = ∆2w/4. Then
W2(µK , π) ≤ (1− αm)KW2(µ0, π) + 1.65(M/m)(αd)1/2 + min
( ΨA
m , M √ A m
)
+
√ α∆2gd+ 4ακ 2
4m +
A αm + ( (1− αm)K + 1 )√ A.
where A = {
5v0d, if 2α > v0 max (2∆wαG, 4αd) , otherwise
This theorem shows that when the stepsize α → 0, VC SGLDLP-L converges to the target distribution up to an error instead of diverging. Our bound is O( √ ∆w) which is equivalent to the distance between SGD with low-precision gradient accumulators and the optimum (Li et al., 2017; Yang et al., 2019). However, we show empirically that VC SGLDLP-L has better dependency on the quantization gap than SGD. We leave the improvement of the theoretical bound for future work.
We empirically demonstrate VC SGLDLP-L on the standard Gaussian distribution under the same setting as in the previous section in Figure 1. Regardless of the stepsize, VC SGLDLP-L converges to the target distribution and approximates the target distribution as well as SGLDLP-F, showing that preserving the correct variance is the key to ensure correct convergence.

heading: 5 EXPERIMENTS
text: We demonstrate low-precision SGLD with full-precision gradient accumulators (SGLDLP-F) and with variance-corrected low-precision gradient accumulators (VC SGLDLP-L) on a logistic regression and multilayer perceptron on MNIST dataset (Section 5.1), ResNet-18 on CIFAR datasets and LSTM on IMDB dataset (Section 5.2), and ResNet-18 on Imagenet dataset (Section 5.3). We use qtorch (Zhang et al., 2019) to simulate low-precision training on these experiments.
Algorithm 2 Variance-Corrected Quantization Function Qvc. input: (µ, v, ∆, Qd, Qs) // Qvc returns a variable with mean µ and variance v v0 ← ∆2/4 // ∆2/4 is the largest possible variance that stochastic rounding can cause if v > v0 then x← µ+ √ v − v0ξ, where ξ ∼ N (0, Id)
r ← x−Qd(x) for all i do
ci ←  ∆, w.p. v0+r 2 i +|ri|∆ 2∆2 −∆, w.p.v0+r 2 i−|ri|∆ 2∆2 // When v > v0, we add a portion of Gaussian noise and
0, otherwise // sample from the weight grid to make up the remaining end for θ ← Qd(x) + sign(r) c
else r ← µ−Qs(µ) for all i do vs ← ( 1− |ri|∆ ) · r2i + |ri| ∆ · (−ri + sign(ri)∆) 2
if v > vs then
ci ←  ∆, w.p.v−vs2∆2 −∆, w.p.v−vs2∆2 // When v ≤ v0, we sample from the weight grid 0, otherwise // to achieve the target variance
θi ← Qs(µ)i + ci else θi ← Qs(µ)i
end if end for
end if clip θ if outside representable range return θ

heading: 5.1 LOGISTIC REGRESSION AND MULTILAYER PERCEPTRON
text: We first empirically verify our theorems on a logistic regression on MNIST dataset. We use N (0, 1/6) as the prior distribution following (Yang et al., 2019) and the resulting posterior distribution satisfies the assumptions in Section 4. We use fixed point numbers with 2 integer bits and vary the number of fractional bits which is corresponding to varying the quantization gap ∆.
From Figure 2a, SGLDLP-F is more robust to the decay of bits than SGDLP-F since SGLDLPF outperforms SGDLP-F on all number of bits and recovers the full-precision result with 6 bits whereas SGDLP-F needs 10 bits. This verifies Theorem 1 that SGLDLP-F converges to the target distribution up to an error and is more robust to the quantiazation gap than SGDLP-F. With lowprecision gradient accumulators, we can see that VC SGLDLP-L is significantly better than naı̈ve SGLDLP-L, indicating the effectiveness of using the variance-corrected quantization function for quantizing gradient accumulators. These results verify Theorem 2 and Theorem 3. Besides, VC SGLDLP-L outperforms SGDLP on all number of bits and even outperforms SGDLP-F when using 2 fractional bits. VC SGLDLP-L matches full-precision SGLD results with 6 bits whereas SGDLPL needs 8 or 10 bits. These observations demonstrate that SGLD is still more favorable than SGD when using low-precision gradient accumulators.
To test whether these results apply to non-log-concave distributions, we replace the logistic regression model by a two-layer multilayer perceptron (MLP). The MLP has 100 hidden units and RELU nonlinearilities. From Figure 2b, we observe similar results as on the logistic regression, suggesting that empirically our analysis still stands on non-log-concave distributions. We also provide negative log-likelihood (NLL) comparisons in terms of number of bits on the logistic regression and MLP in Figure 3 in Appendix.

heading: 5.2 RESNET AND LSTM
text: We consider image and sentiment classification tasks: CIFAR datasets (Krizhevsky et al., 2009) on ResNet-18 (He et al., 2016), and IMDB dataset (Maas et al., 2011) on LSTM (Hochreiter & Schmidhuber, 1997). We use 8-bit number representation since it becomes increasingly popular in training deep models and is powered by new generation of chips (Sun et al., 2019; Banner et al., 2018; Wang et al., 2018b). We report test errors averaged over 3 runs with the standard error in Table 1.
Fixed Point We use 8-bit fixed point for weights and gradients but full-precision for activations since we find low-precision activations significantly harm the performance. Similar to the results in previous sections, SGLDLP-F is better than SGDLP-F and VC SGLDLP-L significantly outperforms naı̈ve SGLDLP-L and SGDLP-L across datasets and architectures. Notably, the improvement of SGLD over SGD becomes larger when using more low-precision arithmetic. For example, on CIFAR-100, VC SGLDLP-L outperforms SGDLP-L by 2.24%, SGLDLP-F outperforms SGDLP-F by 0.54% and SGLD outperforms SGDLP by 0.06%. This demonstrates that SGLD is particularly compatible with low-precision deep learning because of its natural ability to handle system noise.
Block Floating Point We also consider block floating point (BFP) which is another common number type and is often preferred over fixed point on deep models due to less quantization error caused by overflow and underflow (Song et al., 2018). Following the block design in Yang et al. (2019), we use small-block for ResNet and big-block for LSTM. The Qvc function naturally generalizes to BFP and only needs a small modification (see Appendix E for the algorithm of Qvc with BFP). By using BFP, the results of all low-precision methods improve over fixed point. SGLDLP-F can match the performance of SGLDFP with all numbers quantized to 8-bit except gradient accumulators. VC SGLDLP-L still outperforms naive SGLDLP-L indicating the effectiveness of Qvc with BFP. Again, SGLDFP-F and VC SGLDLP-L outperform their SGD counterparts on all tasks, suggesting the general applicability of low-precision SGLD with different number types.
Cyclical SGLD We further apply low-precision to a recent variant of SGLD, cSGLD, which utilizes a cyclical learning rate schedule to speed up convergence (Zhang et al., 2020). We observe that the results of cSGLD-F are very close to those of cSGLDFP, and VC cSGLDLP-L can match or even outperforms full-precision SGD with all numbers quantized to 8 bits! These results indicate that diverse samples from different modes can countereffect the quantization error by providing complementary predictions.

heading: 5.3 IMAGENET
text: Finally, we test low-precision SGLD on a large scale image classification dataset, ImageNet, with ResNet-18. We train SGD for 90 epochs and train SGLD for 10 epochs using the trained SGD model as the initialization. The improvement of SGLD over SGD is larger in low-precision (0.76% top-1
error) than in full-precision (0.17% top-1 error), showing the advantages of low-precision SGLD on large-scale deep learning tasks.

heading: 6 CONCLUSION
text: We provide the first comprehensive study of low-precision SGLD. In theory, we prove that with fullprecision gradient accumulators, SGLD can converge to the target distribution within a distance that is asymptotically smaller than the distance between the SGD estimation and the optimum; with lowprecision gradient accumulators, SGLD can diverge arbitrarily far away from the target distribution with small stepsizes. We find that the issue is caused by the wrong variance in each update and thus develop a new varaince-corrected quantization function that preserves the correct variance. We prove that SGLD with this quantization function converges to the target distribution up to a certain level depending on the quantization gap. In practice, we verify our theoretical results and demonstrate SGLDLP-F and VC SGLDLP-L are comparable to full-precision SGLD with 8-bit on image and sentiment classification tasks.
In the future, it will be interesting to extend low-precision computation to other stochastic MCMC methods and improve theoretical bounds to better reflect empirical performance. We hope this work shed lights on accelerating MCMC methods for Bayesian deep learning.

heading: A PROOF OF THEOREM 1
text: Our proof follows Theorem 4 in Dalalyan & Karagulyan (2019), which provides a convergence bound of Langevin dynamics with noisy gradient. We state the result of Theorem 4 in Dalalyan & Karagulyan (2019) below.
We consider Langevin dynamics whose update rule is
θk+1 = θk − α (∇U(θk) + ζk) + √ 2αξk+1.
The noise in the gradient ζk has the following three assumptions: E [ ‖E [ζk|θk]‖22 ] ≤ δ2d, E [ ‖ζk −E [ζk|θk]‖22 ] ≤ σ2d, ξk+1 is independent of (ζ0, · · · , ζk)
where δ > 0 and σ > 0 are some constants. Under the same assumptions in Section 4, we have the convergence bound for the above Langevin dynamics. Theorem 4. We run the above Langevin dynamics with α ≤ 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 + δ √ d
m +
σ2(αd)1/2
1.65M + σ √ m .
Proof. We write the SGLDLP-F update as the following
θk+1 = θk − αQg(∇Ũ(Qw(θk))) + √ 2αξk+1
= θk − α (∇U(θk) + ζk) + √ 2αξk+1
where
ζk = Qg(∇Ũ(Qw(θk)))−∇U(θk) = Qg(∇Ũ(Qw(θk)))−∇Ũ(Qw(θk))
+∇Ũ(Qw(θk))−∇U(Qw(θk)) +∇U(Qw(θk))−∇U(θk).
Since E[∇Ũ(x)] = ∇U(x) and E[Q(x)] = x, we have E[ζk] = E [ Qg(∇Ũ(Qw(θk)))−∇Ũ(Qw(θk)) ] + E [ ∇Ũ(Qw(θk))−∇U(Qw(θk)) ] + E [∇U(Qw(θk))−∇U(θk)]
= E [∇U(Qw(θk))−∇U(θk)] .
By the assumption, we know that E [ ‖∇U(Qw(θk))−∇U(θk)‖22 ] ≤M2E [ ‖Qw(θk)− θk‖22 ] ≤M2 · ∆ 2 wd
4 .
It follows that
‖E[ζk]‖22 = ‖E [∇U(Qw(θk))−∇U(θk)]‖ 2 2 ≤ E [ ‖∇U(Qw(θk))−∇U(θk)‖22 ] ≤M2 · ∆ 2 wd
4 .
Let f : R→ Rd denote the function
f(a) = ∇U(θk + a(Qw(θk)− θk)).
By the mean value theorem, there will exist an a ∈ [0, 1] (a function of the weight quantization randomness) such that
f(1)− f(0) = f ′(a).
So,
E[ζk] = E [∇U(Qw(θk))−∇U(θk)] = E [ ∇2U(θk + a(Qw(θk)− θk))(Qw(θk)− θk) ] = E [ ∇2U(θk)(Qw(θk)− θk)
] + E [( ∇2U(θk + a(Qw(θk)− θk))−∇2U(θk) ) (Qw(θk)− θk)
] = E [( ∇2U(θk + a(Qw(θk)− θk))−∇2U(θk) ) (Qw(θk)− θk) ] .
Now, by the assumption ‖∇2U(x)−∇2U(y)‖2 ≤ Ψ‖x− y‖, we get
‖E[ζk]‖ = ∥∥E [(∇2U(θk + a(Qw(θk)− θk))−∇2U(θk)) (Qw(θk)− θk)]∥∥
≤ E [∥∥(∇2U(θk + a(Qw(θk)− θk))−∇2U(θk)) (Qw(θk)− θk)∥∥]
≤ E [∥∥∇2U(θk + a(Qw(θk)− θk))−∇2U(θk)∥∥2 ‖Qw(θk)− θk)‖] ≤ E [Ψ ‖a(Qw(θk)− θk)‖ ‖Qw(θk)− θk)‖]
≤ ΨE [ ‖Qw(θk)− θk‖2 ] ≤ Ψ∆ 2 wd
4
This combined with the previous result gives us
‖E[ζk]‖ = min
( Ψ∆2w √ d
4 , M∆w 2
) .
Now considering the variance of ζk,
E [ ‖ζk −E[ζk]‖22 ] ≤ E [ ‖ζk‖22
] ≤ E [∥∥∥Qg(∇Ũ(Qw(θk)))−∇Ũ(Qw(θk))∥∥∥2 2
] + E [∥∥∥∇Ũ(Qw(θk))−∇U(Qw(θk))∥∥∥2 2 ] + E [ ‖∇U(Qw(θk))−∇U(θk)‖22
] ≤ ∆2gd
4 + κ2 +M2 · ∆
2 wd 4 .
We set δ and σ in Theorem 4 in Dalalyan & Karagulyan (2019) as the following
δ = min
( Ψ∆2w √ d
4 , M∆w 2
) , σ2d = ∆2gd
4 + κ2 +M2 · ∆
2 wd 4 = (∆2g +M 2∆2w)d+ 4κ 2 4 .
The assumptions in Theorem 4 in Dalalyan & Karagulyan (2019) are satisfied, thus we can apply the result in Theorem 4 and get
W2(µK , π) ≤ (1− αm)KW2(µ0, π) + 1.65(M/m)(αd)1/2 + δ √ d
m +
σ2(αd)1/2
1.65M + σ √ m
≤ (1− αm)KW2(µ0, π) + 1.65(M/m)(αd)1/2 + δ √ d
m +
√ σ2αd
m
= (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 .

heading: B PROOF OF THEOREM 2
text: Proof. The update of SGLDLP-L is
θk+1 = Qw ( θk − αQg(∇Ũ(θk)) + √ 2αξk+1 ) .
Let ψk+1 = θk − αQg(∇Ũ(θk)) + √ 2αξk+1 so that θk = Qw(ψk). Then,
ψk+1 = θk − αQg(∇Ũ(θk)) + √ 2αξk+1
= Qw(ψk)− αQg(∇Ũ(Qw(ψk))) + √ 2αξk+1 = ψk − α(∇U(ψk) + ζk) + √ 2αξk+1
where
ζk = ψk − θk
α +Qg(∇Ũ(θk))−∇U(ψk)
= ψk − θk
α +Qg(∇Ũ(θk))−∇Ũ(θk) +∇Ũ(θk)−∇U(θk) +∇U(θk)−∇U(ψk)
Similar to the previous proof, we know that
E[ζk] = E [∇U(θk)−∇U(ψk)] = E [∇U(Qw(ψk))−∇U(ψk)] , so
‖E[ζk]‖ ≤ min
( Ψ∆2wd
4 , M∆w
√ d
2
) ,
and it suffices to set δ as above. On the other hand, the variance will be bounded by E [ ‖ζk −E[ζk]‖22 ] ≤ E [∥∥∥Qg(∇Ũ(θk))−∇Ũ(θk)∥∥∥2]+ E [∥∥∥∇Ũ(θk)−∇U(θk)∥∥∥2] + E [∥∥∥∥ψk − θkα +∇U(θk)−∇U(ψk) ∥∥∥∥2 ]
≤ ∆2gd
4 + κ2 + E
[ ‖∇F (ψk)−∇F (θk)‖2 ] ,
where F (θ) = 12α ‖θ‖ 2 − U(θ). Observe that since U is m-strongly convex and M -smooth, and α−1 ≥M/2, F must be h−1-smooth, and so
E [ ‖ζk −E[ζk]‖22 ] ≤ ∆2gd
4 + κ2 +
1 α2 E [ ‖ψk − θk‖2 ] ≤ ∆2gd
4 + κ2 +
∆2wd
4α2 .
This is essentially replacing theM2 in the previous analysis with α−2, so, supposing the distribution of ψK+1 is νK , it will give us the rate of
W2(νK , π) ≤ (1− αm)KW2(ν0, π) + 1.65(M/m)(αd)1/2 + min
( Ψ∆2wd
4m , M∆w
√ d
2m
)
+
√ (α∆2g + α −1∆2w)d+ 4ακ 2
4m .
We also have
W2(µk, νk) =
( inf
J∈J (x,y)
∫ ‖x− y‖2 dJ(x, y) )1/2 ≤ E [ ‖θk+1 − ψk+1‖2 ] 1 2 ≤ ∆w √ d
2 .
Combining these two results, we get
W2(µK , π) ≤W2(µK , νK) +W2(νK , π)
≤ (1− αm)KW2(ν0, π) + 1.65(M/m)(αd)1/2 + min
( Ψ∆2wd
4m , M∆w
√ d
2m
)
+
√ (α∆2g + α −1∆2w)d+ 4ακ 2
4m +
∆w √ d
2
≤ (1− αm)KW2(µ0, π) + 1.65(M/m)(αd)1/2 + min
( Ψ∆2wd
4m , M∆w
√ d
2m
)
+
√ (α∆2g + α −1∆2w)d+ 4ακ 2
4m + ( (1− αm)K + 1 ) ∆w√d 2 .

heading: C PROOF OF THEOREM 3
text: Proof. The update of VC SGLDLP-L is θk+1 = Q vc ( θk − αQg(∇Ũ(θk)), 2α,∆w, Qd, Qs ) We ignore the variance of Qg since it is relatively small compared to weight quantization variance in practice. Qvc is defined as in Algorithm 1 and we have E[θk+1] = θk − α∇U(θk). Let ψk+1 = θk − αQg(∇Ũ(θk)) + √ 2αξk+1 then it follows that
ψk+1 − θk+1 = θk − αQg(∇Ũ(θk)) + √
2αξk+1 − θk+1, and
ψk+1 = ψk − α(∇U(ψk) + ζk) + √ 2αξk+1
where
ζk = ψk − θk
α +Qg(∇Ũ(θk))−∇U(ψk)
= ψk − θk
α +Qg(∇Ũ(θk))−∇Ũ(θk) +∇Ũ(θk)−∇U(θk) +∇U(θk)−∇U(ψk).
Note that E[ψk − θk] = 0. Similar to previous proof, we know that ‖E[ζk]‖2 = ‖E [∇U(θk)−∇U(ψk)]‖2 ≤M2E [ ‖ψk − θk‖22 ] .
When 2α > v0 = ∆2w 4 , we have that
E[‖ψk − θk‖2]
= E [∥∥∥(θk−1 − αQg(∇Ũ(θk−1)))+√2αξk −Qd (θk−1 − αQg(∇Ũ(θk−1)) +√2α− v0ξk)− sign(r)c∥∥∥2]
Let
b = Qd ( θk−1 − αQg(∇Ũ(θk−1)) + √ 2α− v0ξk ) − ( θk−1 − αQg(∇Ũ(θk−1)) + √ 2α− v0ξk ) ,
then |b| ≤ ∆w2 and E [ ‖ψk − θk‖2 ] = E
[∥∥∥(θk−1 − αQg(∇Ũ(θk−1)))+√2αξk − (θk−1 − αQg(∇Ũ(θk−1)) +√2α− v0ξk)− b− sign(r)c∥∥∥2] = E
[∥∥∥√2αξk −√2α− v0ξk − b− sign(r)c∥∥∥2] ≤ E
[∥∥∥√2αξk −√2α− v0ξk − b∥∥∥2]+ E [‖sign(r)c‖2] ≤ E [∥∥∥∥∣∣∣√2αξk −√2α− v0ξk∣∣∣+ ∆w2 ∥∥∥∥2 ] + v0d ≤ ( √ 2α− √ 2α− v0)2E[‖ξk‖2] + ( √ 2α− √ 2α− v0)∆wE[|ξk|] + 2v0d
≤ ( ( √ 2α− √ 2α− v0)2 + ( √ 2α− √ 2α− v0)∆w + 2v0 ) d
Since 2xy ≤ x2 + y2, we get
( √ 2α− √ 2α− v0)∆w ≤ ( √ 2α− √ 2α− v0)2 + ∆2w 4 = ( √ 2α− √
2α− v0)2 + v0 It follows
E[‖ψk − θk‖2] ≤ ( 2( √ 2α− √ 2α− v0)2 + 3v0 ) d.
Note that √
2α− √ 2α− v0 = 2α− (2α− v0)√ 2α+ √ 2α− v0 = v0√ 2α+ √ 2α− v0 ≤ v0√ 2α .
The expectation becomes
E[‖ψk − θk‖2] ≤ ( v20 α + 3v0 ) d.
Since 2α > v0, it follows E[‖ψk − θk‖2] ≤ (2v0 + 3v0) d = 5v0d.
Let A = 5v0d. Then ‖E[ζk]‖2 ≤M2 ·A,
and also ‖E[ζk]‖ ≤ Ψ ·A.
Therefore, δ = min ( ΨA,M √ A ) .
The variance will be bounded by E [ ‖ζk −E[ζk]‖22 ] ≤ E [∥∥∥Qg(∇Ũ(θk))−∇Ũ(θk)∥∥∥2]+ E [∥∥∥∇Ũ(θk)−∇U(θk)∥∥∥2] + E [∥∥∥∥ψk − θkα +∇U(θk)−∇U(ψk) ∥∥∥∥2 ]
≤ ∆2gd
4 + κ2 +
1 α2 E [ ‖ψk − θk‖2 ] ≤ ∆2gd
4 + κ2 +
A α2 .
Supposing the distribution of ψK+1 is νK , it will give us the rate of
W2(νK , π) ≤ (1− αm)KW2(ν0, π) + 1.65(M/m)(αd)1/2 + min ( Ψ ·A m , M √ A m )
+
√ α∆2gd+ 4ακ 2
4m +
A
αm .
We also have
W2(µK , νK) =
( inf
J∈J (x,y)
∫ ‖x− y‖2 dJ(x, y) )1/2 ≤ E [ ‖θK+1 − ψK+1‖2 ] 1 2 ≤ √ A.
Combining these two results, we get
W2(µK , π) ≤W2(µK , νK) +W2(νK , π)
≤ (1− αm)KW2(ν0, π) + 1.65(M/m)(αd)1/2 + min ( Ψ ·A m , M √ A m )
+
√ α∆2gd+ 4ακ 2
4m +
A αm + √ A
≤ (1− αm)KW2(µ0, π) + 1.65(M/m)(αd)1/2 + min ( Ψ ·A m , M √ A m )
+
√ α∆2gd+ 4ακ 2
4m +
A αm + ( (1− αm)K + 1 )√ A.
When 2α < ∆ 2 w 4 , since we assume that the gradient is bounded by E [∥∥∥Qg(∇Ũ(θk))∥∥∥
1
] ≤ G
E[‖ψk − θk‖2] = E [∥∥∥(θk−1 − αQg(∇Ũ(θk−1)))− θk +√2αξk∥∥∥2]
= E [∥∥∥(θk−1 − αQg(∇Ũ(θk−1)))− θk∥∥∥2]+ E [∥∥∥√2αξk∥∥∥2] ≤ max ( 2E
[∥∥∥(θk−1 − αQg(∇Ũ(θk−1)))−Qs (θk−1 − αQg(∇Ũ(θk−1)))∥∥∥2] , 4αd) Using the bound equation (6) in Li & De Sa (2019),
E [∥∥∥(θk−1 − αQg(∇Ũ(θk−1)))−Qs (θk−1 − αQg(∇Ũ(θk−1)))∥∥∥2] ≤ ∆wαE [∥∥∥Qg(∇Ũ(θk−1))∥∥∥ 1
] ≤ ∆wαG
It follows
E[‖ψk − θk‖2] ≤ max (2∆wαG, 4αd) Let A = max (2∆wαG, 4αd). The rest is that same as the case 2α > v0.

heading: D COMPARISON TO SGD BOUNDS
text: Following previous work in comparing sampling and optimization methods (Ma et al., 2019; Talwar, 2019), we also compare our 2-Wasserstein distance bound with previosu SGD bounds. Previous low-precision SGD convergence bounds are shown in terms of the squared distance to the optimum∥∥θ̄K − θ∗∥∥22 (Yang et al., 2019). In order to compare our bounds with theirs, we consider a 2- Wasserstein distance between two point distributions. Let µK be the point distribution assigns zero
probability everywhere except θ̄K and π be the point distribution assigns zero probability everywhere except θ∗. Then we get
W2(µK , π) =
( inf
J∈J (x,y)
∫ ‖x− y‖2 dJ(x, y) )1/2 ≤ E [∥∥θ̄K − θ∗∥∥2] 12 . From (Yang et al., 2019), we know that E
[∥∥θ̄K − θ∗∥∥2] 12 is proportional to ∆w. Therefore, our 2-Wasserstein distance is O(∆2w) whereas SGD’s 2-Wasserstein distance is O(∆w), which shows SGLD is more robust to the quantization error.

heading: E ALGORITHMS WITH (BLOCK) FLOATING POINT NUMBERS
text: The qunatization gap is floating point and block floating point is computed as ∆w(µ)← {
2E[µ]−W+2 where E[µ] = clip(blog2(max |µ|)c , l, u) block floating point 2E[µ]−W where E[µ] = clip(blog2(|µ|)c , l, u) floating point
(3)
Then deterministic rounding and stochastic rounding are defined with ∆w. VC SGLDLP-L with (block) floating point is outlined in Algorithm 3
Algorithm 3 VC SGLDLP-L with (Block) Floating Point. given: Stepsize α, number of training iterationsK, gradient quantizerQg , deterministic rounding with (block) floating point Qd, stochastic rounding with (block) floating point Qs, F bits to represent the shared exponent (block floating point) or the exponent (floating point), W bits to represent each number in the block (block floating point) or the mantissa (floating point). let l← −2F−1, u← 2F−1 − 1 for k = 1 : K do
compute µ← θk − αQg ( ∇Ũ(θk−1) ) compute ∆w(µ) following Equation (3) update θk+1 ← Qvc ( µ, 2αT,∆w(µ), Q d, Qs )
end for output: samples {θk}
Algorithm 4 Variance-Corrected Quantization Function Qvc with (Block) Floating Point. input: (µ, v, ∆, Qd, Qs) v0 ← ∆2/4 if v > v0 then x← µ+ √ v − v0ξ, where ξ ∼ N (0, Id)
r ← x−Qd(x) recompute ∆← ∆w(x) following Equation (3)
c←  ∆, w.p.v0+r 2+|r|∆ 2∆2 −∆, w.p.v0+r 2−|r|∆ 2∆2 0, otherwise return Qd(x) + sign(r) · c
else the same as in fixed point numbers end if
The Qvc function with (block) floating point in Algorithm 4 is the same as Algorithm 1 except the line in red. We recompute the quantization gap ∆ after adding Gaussian noise to make sure it aligns with the quantization gap of x.

heading: F EXPERIMENTAL DETAILS AND ADDITIONAL RESULTS
text: F.1 SAMPLING METHODS
For both SGLD and low-precision SGLD, we collected samples {θk}JK=1 from the posterior of the model’s weight, and obtained the prediction on test data {x∗, y∗} by Bayesian model averaging
p(y∗|x∗,D) ≈ 1 J J∑ j=1 p(y∗|x∗,D, θj).
F.2 MNIST
We train all methods on logistic regression and MLP for 20 epochs with learning rate 0.1 and batch size 64. We additionally report negative log-likelihood (NLL) comparisons in Figure 3.
F.3 CIFAR AND IMDB
For CIFAR datasets, we use batch size 128, learning rate 0.5 and weight decay 5e − 4. We train the model for 245 epochs and used the same decay learning rate schedule as in Yang et al. (2019). We collect 50 samples for SGLD. For cyclical learning rate schedule, we use 7 cycles and collect 5 models per cycle (35 models in total).
For IMDB dataset, we use batch size 80, learning rate 0.3 and weight decay 5e− 4. We use a twolayer LSTM. The embedding dimension and the hidden dimension are 100 and 256 respectively. We train the model for 50 epochs and used the same decay learning rate schedule as on CIFAR datasets. We collect 20 samples for SGLD. For cyclical learning rate schedule, we use 1 cycles and collect 20 models.
F.4 IMAGENET
We use batch size 256, learning rate 0.2 and weight decay 1e − 4. We use the same decay learning rate schedule as in He et al. (2016) and collect 20 models for SGLD.

Review summaries:
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. 

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.

The author proposed a comprehensive study of low-precision Stochastic Gradient Langevin Dynamics (SGLD). They proposed the convergence analysis of low-precision Stochastic Gradient Langevin Dynamics on strongly log-concave distributions.  Motivated by the diverging fact of SGLDLP-L with small stepsizes, they proposed a variance-corrected quantization function to ensure a bounded bias. Empirical experiments based on large-scale DNN examples are evaluated.