L OW -P RECISION S TOCHASTIC G RADIENT L ANGEVIN D YNAMICS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ﬁrst 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 han- dle 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
I NTRODUCTION
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 efﬁciently 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 signiﬁcant memory and computational expenses.
In this paper, we provide the ﬁrst 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.
[Figure]
Surprisingly, ﬁnd 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 signiﬁcantly 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.
R ELATED W ORK
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 efﬁcient 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.
P RELIMINARIES 3. S TOCHASTIC G RADIENT L ANGEVIN D YNAMICS
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
[Equation]
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 .
Speciﬁcally, SGLD updates the parameter θ in the ( k + 1) -th step following the rule
[Equation]
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. L OW -P RECISION T RAINING
We study training a low-precision model by SGLD from scratch, to reduce both training and testing costs.
We use ﬁxed 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 ﬂoating 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 = − 2 W − F − 1 and u = 2 W − F − 1 − 2 − F , respectively.
The deterministic rounding quantizes a number to its nearest representable neighbour:
[Equation]
where clip ( x, l, u ) =
max(min
( x, u ) , l ) .
Stochastic rounding quantizes a number with probability:
[Equation]
Q s is unbiased, which means E [ Q s ( θ )]
= θ .
It is more favorable than Q d 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 Q w and ∆ w to denote the weights’ quantizer and quantization gap, Q g and ∆ g to denote gradients’ quantizer and quantization gap.
Please note that Q g 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
[Equation]
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
[Equation]
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).
L OW -P RECISION SGLD
In this section, we ﬁrst 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. F ULL -P RECISION G RADIENT A CCUMULATORS
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.
[Equation]
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).
Speciﬁcally, the energy function U satisﬁes the following,
[Equation]
for some positive constants m , M and Ψ .
We further assume the variance of the stochastic gradient is bounded by E
[ (cid:107)∇ ˜ U ( θ ) − ∇ U ( θ ) (cid:107) 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
[Equation]
There are three main takeways of this theorem: First, it shows that SGLDLP-F converges to the accuracy ﬂoor min (cid:16)
[Equation]
2 w d 4 m , M
[Equation]
m (cid:17) which depends on the quantization gap ∆ w given large enough
K and small enough α .
Second, our convergence bound is O (∆ 2 w ) 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).
4. L OW -P RECISION G RADIENT A CCUMULATORS
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
[Equation]
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
[Equation]
This theorem suggests that as the stepsize α decreases, W 2 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 ﬁxed 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 W 2 distance to avoid divergence, however, getting this optimal stepsize is generally difﬁcult 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).
4. V ARIANCE -C ORRECTED Q UANTIZATION
To approach correcting the problems with na ¨ ıve SGLDLP-L, we ﬁrst 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.
Speciﬁcally, given the stochastic gradient ∇ ˜ U , the update of full-precision SGLD can be viewed as sampling from a Gaussian distribution for each dimension i
[Equation]
Using stochastic rounding as the weight quantizer and the gradient quantizer in SGLDLP-L gives us
[Equation]
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 Q g 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
[Equation]
Algorithm given: Stepsize α , number of training iterations K , gradient quantizer Q g , deterministic rounding
Q d , stochastic rounding Q s and quantization gap of weights ∆ w . for k = 1 : K do update θ k +1 ← Q vc (cid:16) θ k − αQ g (cid:16) ∇ ˜ U ( θ k ) (cid:17) , 2 α, ∆ w , Q d , Q s (cid:17) end for output 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 Q vc in Algorithm 2.
Q vc 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 < v s .
However that case rarely happens in practice, because the stepsize has to be extremely small.
The main idea of Q vc 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 , Q vc ﬁrst adds a portion of the Gaussian noise and uses a sample from the weight grid to make up the remaining; in the other situation, Q vc directly samples from the weight grid to achieve the target variance.
Although Q vc 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 Q vc , 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 (cid:104)(cid:13)(cid:13)(cid:13)
Q g ( ∇ ˜ U
( θ k ))
(cid:13)(cid:13)(cid:13) 1 (cid:105) ≤ G .
Let v 0 =
∆ 2 w / 4 .
Then
[Equation]
where A = (cid:26) 5 v 0 d, if 2 α > vmax (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.
E XPERIMENTS
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 input : ( µ , v , ∆ , Q d , Q s ) //
Q vc returns a variable with mean µ and variance v v 0 ← ∆ 2 /// ∆ 2 / 4 is the largest possible variance that stochastic rounding can cause if v > v 0 then x ← µ + √ v − v 0 ξ , where ξ ∼ N (0 , I d ) r ← x − Q d ( x ) for all i do c i ←   ∆ , w.p. v 0 + r 2 i + | r i | ∆ 2∆− ∆ , w.p. v 0 + r 2 i −| r i | ∆ 2∆// When v > v 0 , we add a portion of Gaussian noise and 0 , otherwise // sample from the weight grid to make up the remaining end for θ ← Q d ( x ) + sign ( r ) (cid:12) c else r ← µ − Q s ( µ ) for all i do v s ← (cid:16) 1 − | r i | ∆ (cid:17) · r 2 i + | r i | ∆ · ( − r i + sign ( r i )∆)if v > v s then c i ←   ∆ , w.p. v − v s 2∆− ∆ , w.p. v − v s 2∆// When v ≤ v 0 , we sample from the weight grid 0 , otherwise // to achieve the target variance θ i ← Q s ( µ ) i + c i else θ i ← Q s ( µ ) i end if end for end if clip θ if outside representable range return θ
5. L OGISTIC R EGRESSION AND M ULTILAYER P ERCEPTRON
We ﬁrst 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 satisﬁes the assumptions in Section 4.
We use ﬁxed 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 veriﬁes 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 signiﬁcantly 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.
[Figure]
Bits
T e s t
E rr o r
[Figure]
[Figure]
SGDLP
[Figure]
L Naive SGLDLP-L
VC SGLDLP-L SGDLP-F
SGLDLP-F SGDFP (7.68%)
SGLDFP (7.63%)        10 Number of Fractional Bits
T e s t E rr
[Figure]
[Figure]
SGDLP-L Naive SGLDLP-L
VC SGLDLP-L SGDLP-
[Figure]
SGLDLP-F SGDFP (2.03%)
SGLDFP (1.9%)
[Figure]
5. R ES N ET AND LSTM
We consider image and sentiment classiﬁcation 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 ﬁxed point for weights and gradients but full-precision for activations since we ﬁnd low-precision activations signiﬁcantly harm the performance.
Similar to the results in previous sections, SGLDLP-F is better than SGDLP-F and VC SGLDLP-L signiﬁcantly 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 ﬂoating point (BFP) which is another common number type and is often preferred over ﬁxed point on deep models due to less quantization error caused by overﬂow and underﬂow (Song et al., 2018).
Following the block design in Yang et al.
, we use small-block for ResNet and big-block for LSTM.
The Q vc function naturally generalizes to
BFP and only needs a small modiﬁcation (see Appendix E for the algorithm of Q vc with BFP).
By using BFP, the results of all low-precision methods improve over ﬁxed 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 Q vc 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.
5. I MAGE N ET
Finally, we test low-precision SGLD on a large scale image classiﬁcation 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-
[Table] [Table]
error) than in full-precision (0.17% top-1 error), showing the advantages of low-precision SGLD on large-scale deep learning tasks.
C ONCLUSION
We provide the ﬁrst 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 ﬁnd 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 classiﬁcation tasks.
In the future, it will be interesting to extend low-precision computation to other stochastic MCMC methods and improve theoretical bounds to better reﬂect empirical performance.
We hope this work shed lights on accelerating MCMC methods for Bayesian deep learning.
A P ROOF OF T HEOREM
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
[Equation] The noise in the gradient ζ k has the following three assumptions: [Equation] 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
[Equation]
Proof.
We write the SGLDLP-F update as the following
[Equation]
where
[Equation]
˜ U ( Q w ( θ k ))
− ∇ U ( Q w ( θ k ))
+ ∇ U
( Q w ( θ k )) − ∇ U ( θ k ) .
Since E [ ∇ ˜ U
( x )]
= ∇ U ( x ) and E
[ Q ( x )] = x , we have
E [
[Equation]
By the assumption, we know that
[Equation]
It follows that
[Equation]
Let f : R → R d denote the function f ( a ) =
∇ U ( θ k + a ( Q w ( θ k ) − θ k )) .
By the mean value theorem, there will exist an a ∈ [0 , 1] (a function of the weight quantization randomness) such that
[Equation]
So,
[Equation]
Now, by the assumption (cid:107)∇ 2 U ( x ) − ∇ 2 U ( y ) (cid:107) 2 ≤ Ψ (cid:107) x − y (cid:107) , we get (cid:107) E [ ζ k ] (cid:107) = (cid:13)(cid:13) [Equation] ∇ 2 U
[Equation]
θ k +
[Equation]
w
[Equation]
θ k ) − θ k ))
− ∇ 2 U ( θ
[Equation]
) (cid:1) ( Q w ( θ k ) − θ k ) (cid:3)(cid:13)(cid:13) ≤ E
[Equation]
∇
[Equation]
This combined with the previous result gives us
[Equation]
Now considering the variance of ζ k
[Equation]
E [Equation] We set δ and σ in Theorem 4 in Dalalyan & Karagulyan (2019) as the following [Equation]
[Equation]
The assumptions in Theorem 4 in Dalalyan & Karagulyan (2019) are satisﬁed, thus we can apply the result in Theorem 4 and get
[Equation]
B P ROOF OF T HEOREM
Proof.
The update of SGLDLP-L is
[Equation]
Let ψ k
+
=
θ k − [Equation] g ( ∇ ˜ U
( θ k ))
+ √ 2 αξ k
+
so that θ k = Q w ( ψ k ) .
Then,
[Equation]
where
[Equation]
Similar to the previous proof, we know that
[Equation]
and it sufﬁces to set δ as above.
On the other hand, the variance will be bounded by
[Equation]
, where F ( θ ) = 12 α (cid:107) θ (cid:107) 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
[Equation]
This is essentially replacing the M 2 in the previous analysis with α − 2 , so, supposing the distribution of ψ K +1 is ν K , it will give us the rate of
[Equation]
We also have
[Equation]
Combining these two results
[Equation]
we get
[Equation]
C P ROOF OF T HEOREM
Proof.
The update of VC SGLDLP-L is
[Equation]
We ignore the variance of Q g since it is relatively small compared to weight quantization variance in practice.
Q vc is deﬁned as in Algorithm 1 and we have E [ θ k +1 ]
= θ k − α ∇ U ( θ k ) .
Let ψ k
+
=
θ k − αQ g ( ∇
˜ U
( θ k ))
+ √ 2 αξ k
+1 then it follows that ψ
[Equation]
and ψ
[Equation]
+ where ζ k = ψ k − θ k [Equation] = ψ [Equation]
α
[Equation]
g
[Equation]
∇ U ( ψ k ) .
Note that E [ ψ k − θ k ] = 0 .
Similar to previous proof, we know that (cid:107) [Equation] k
[Equation]
22 (cid:105) .
When 2 α > v 0 = ∆ 2 w 4 , we have that
E [ (cid:107) ψ k − θ k (cid:107) 2 ] = E (cid:20)(cid:13)(cid:13)(cid:13)(cid:16) θ k −
[Table]
Let
[Equation]
then | b | ≤ ∆ w 2 and
[Equation]
[-( θ k − 1 )) (cid:17) + √ 2 αξ k − Q d (cid:16) θ k − 1 − αQ g ( ∇ ˜ U ( θ k − 1 )) + √ 
 α − v 0 ξ k (cid:17) − sign ( r ) c (cid:13)(cid:13)(cid:13) 2 (cid:21) Let b = Q d (cid:16) θ k − 1 − αQ g (-]≤
[-∇ ˜ U ( θ k − 1 )) + √ 
 α − v 0 ξ k (cid:17) − (cid:16) θ k − 1 − αQ g ( ∇ ˜ U ( θ k − 1 )) + √ 
-][Equation]
Since 2 xy ≤ x 2 + y 2 , we get
[-˜ U ( θ k − 1 )) + √ 
 α − v 0 ξ k (cid:17) − b − sign ( r ) c (cid:13)(cid:13)(cid:13) 2 (cid:21) = E (cid:20)(cid:13)(cid:13)(cid:13) √ 2 αξ k − √ 2 α − v 0 ξ k − b − sign ( r ) c (cid:13)(cid:13)(cid:13) 2 (cid:21) ≤ E (cid:20)(cid:13)(cid:13)(cid:13) √ 2 αξ k − √ 2 α − v 0 ξ k − b (cid:13)(cid:13)(cid:13) 2 (cid:21) + E (cid:104) (cid:107) sign ( r ) c (cid:107) 2 (cid:105) ≤ E (cid:34)(cid:13)(cid:13)(cid:13)(cid:13)(cid:12)(cid:12)(cid:12) √ 2 αξ k − √ 2 α − v 0 ξ k (cid:12)(cid:12)(cid:12) + ∆ w 2 (cid:13)(cid:13)(cid:13)(cid:13) 2 (cid:35) + v 0 d ≤ ( √ 2 α − √ 
-][Equation]
√[-2 α − v 0 ) 
 + ( √ 2 α − √ 2 α − v 0 ) ∆ w + 2 v 0 (cid:17) d-]
[Equation]
It follows
[Equation]
Note that
[Equation]
The expectation becomes
[Equation]
Since 2 α > v 0 , it follows
[Equation]
Let A = 5 v 0 d .
Then
[Equation]
, and also
[Equation]
E [ ζ
[Equation]
] (cid:107) ≤ Ψ · A. Therefore,
[Equation]
The variance will be bounded by
[Equation]
Supposing the distribution of ψ K +1 is ν K , it will give us the rate of
[Equation]
We also have
[Equation]
Combining these two results
[Equation]
we get
[Equation]
When 2 α < ∆ 2 w
[Equation]
since we assume that the gradient is bounded by E (cid:104)(cid:13)(cid:13)(cid:13)
Q g
[Equation]
Using the bound equation (6) in Li & De Sa (2019),
[Equation]
It follows
[Equation]
Let A = max (2∆ w αG, 4 αd ) .
The rest is that same as the case 2 α > v 0 .
D C OMPARISON TO SGD B OUNDS
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 (cid:13)(cid:13)
¯ θ K − θ ∗ (cid:13)(cid:13) 2
(Yang et al., 2019).
In order to compare our bounds with theirs, we consider a 2Wasserstein 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
[Equation]
From (Yang et al., 2019), we know that E (cid:104)(cid:13)(cid:13) ¯ θ K − θ ∗ (cid:13)(cid:13) 2 (cid:105) 12 is proportional to ∆ w .
Therefore, our 2-Wasserstein distance is O (∆ 2 w ) whereas SGD’s 2-Wasserstein distance is O (∆ w ) , which shows
SGLD is more robust to the quantization error.
E A LGORITHMS WITH (B LOCK ) F LOATING POINT N UMBERS
The qunatization gap is ﬂoating point and block ﬂoating point is computed as ∆ w ( µ ) ← (cid:26) 2 E [ µ ] − W +2 where E [ µ ] = clip ( (cid:98) log 2 (max | µ | ) (cid:99) , l, u ) block ﬂoating point 2 E [ µ ] − W where E [ µ ] = clip ( (cid:98) log 2 ( | µ | ) (cid:99) , l, u ) ﬂoating point (3)
Then deterministic rounding and stochastic rounding are deﬁned with ∆ w .
VC SGLDLP-L with (block) ﬂoating point is outlined in Algorithm
Algorithm 3 VC SGLDLP-L with (Block) Floating Point.
given: Stepsize α , number of training iterations K , gradient quantizer Q g , deterministic rounding with (block) ﬂoating point Q d , stochastic rounding with (block) ﬂoating point Q s , F bits to represent the shared exponent (block ﬂoating point) or the exponent (ﬂoating point), W bits to represent each number in the block (block ﬂoating point) or the mantissa (ﬂoating point).
let l ← − 2 F − 1 , u ← 2 F − 1 −for k = 1 : K do compute µ ← θ k − αQ g (cid:16) ∇
˜ U
( θ k − 1 ) (cid:17) compute ∆ w ( µ ) following Equation (3) update θ k +1 ← Q vc (cid:0) µ, 2 αT, ∆ w ( µ ) , Q d , Q s (cid:1) end for output : samples { θ k } Algorithm input : ( µ , v , ∆ , Q d , Q s ) v 0 ← ∆ 2 /if v > v 0 then x ← µ + √ v − v 0 ξ , where ξ ∼ N (0 , I d ) r ← x − Q d ( x ) recompute ∆ ← ∆ w ( x ) following Equation (3) c ←   ∆ , w.p. v 0 + r 2 + | r | ∆ 2∆− ∆ , w.p. v 0 + r 2 −| r | ∆ 2∆0 , otherwise return Q d ( x ) + sign ( r ) · c else the same as in ﬁxed point numbers end if
The Q vc function with (block) ﬂoating 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 . [Figure]
F E XPERIMENTAL D ETAILS AND A DDITIONAL R ESULTS F. S AMPLING 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
[Equation]
F. 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. CIFAR AND IMDB
For CIFAR datasets, we use batch size 128, learning rate 0 .
5 and weight decay 5 e − 4 .
We train the model for 245 epochs and used the same decay learning rate schedule as in Yang et al.
We collect 50 samples for SGLD.
For cyclical learning rate schedule, we use 7 cycles and collectmodels per cycle (35 models in total).
IMDB dataset, we use batch size 80, learning rate 0 .
3 and weight decay 5 e − 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 collectmodels.
F. I MAGE N ET
We use batch size 256, learning rate 0 .
2 and weight decay 1 e − 4 .
We use the same decay learning rate schedule as in He et al.
and collect 20 models for SGLD.