Low-precision sampling

When running simulations for a dynamic system on a computer, the impact of round-off errors on number representations (e.g., floating-point representations) can not be naïvely presumed to be negligible. It has been observed that a large deviation can occur between numerical behavior and the theoretical behaviour Guihéneuf, 2015. In Bayesian literature, the same question arises: how does the rounding error from computer simulation affect the convergence property of a Markov chain? Roberts et al. (1998) shows a rather alarming example that a well-behaved Markov chain (with Feller continuity and geometrical ergodicity) becomes transient (i.e., will not converge to any target distribution) after adding an arbitrarily small roundoff error. Hoffman (2021) also warned that roundoff errors can severely degrade the performance of the Metropolis-Hastings sampling algorithm. To sum up, a common wisdom for conventional MCMC practice is to utilize high-precision numerical representations to limit the negative impact of roundoff errors. However, in this chapter, we will embrace the low-precision representations in the era of SG-MCMC samplers.

Low-Precision Deep Neural Network¶

The rise of large-scale deep neural networks (DNNs) and large language models (LLMs) has led to significant computational and memory costs of model training and inference. To improve efficiency, low-precision deep learning Hubara et al., 2016Jacob et al., 2018Rastegari et al., 2016Zhang et al., 2018Zhou et al., 2016Li et al., 2017 — which uses fewer bits to represent values of model parameters, activations, and gradients, and thus can substantially lower resource demands — has emerged as a promising direction. Precision formats such as 16-bit (FP16), 8-bit (INT8), 4-bit, and even binary (1-bit) offer memory reduction, energy savings, and speedup in inference, enabling deployment on edge devices and platforms, and specialized accelerators. On the other hand, well-designed low-precision optimization algorithms can maintain the model utility, such that the performance degradation from full-precision models to low-precision models is tolerable.

Modern foundations of low-precision deep learning include Gupta et al. (2015), which leverages 16-bit wide fixed-point arithmetic and stochastic rounding to accelerate the training speed of DNNs, and Vanhoucke et al. (2011), which discusses the possibility of running DNN inference on x86 CPUs rather than GPUs using 8-bit fixed-point arithmetic. In general, there are two directions for achieving lower precision representations of DNN models: quantization-aware training (QAT) and post-training quantization (PTQ). The former simulates quantization during the training process, allowing the model to learn a robust solution that accounts for quantization errors induced by rounding; the latter, on the other hand, applies quantization to off-the-shelf pre-trained models. It is arguably true that QAT outperforms PTQ in terms of better accuracy, but incurs greater computational costs of model training. Other worth mentioning techniques of low-precision models are, but not limited to, non-uniform quantization and mixed precision training. Opposite to conventional quantization methods that map floating point values to an evenly spaced quantized target space, non-uniform quantization chooses an unevenly spaced quantized target space that can allocate more quantization levels where the values are dense. This unevenly spaced quantized target space can be of logarithmic scale Miyashita et al., 2016, or adaptively learned from the training process Zhang et al., 2018. Mixed-precision methods Dong et al., 2019Micikevicius et al., 2017 assign different bit-widths to different layers of the neural networks or different computation steps of the training, leveraging the fact that different components of DNNs have varying sensitivities to quantization error. Theoretical investigations related to low-precision training are available in Li et al. (2017)Sa et al. (2018)Sakr et al. (2019)Markov et al. (2023).

Inevitably, low-precision optimization achieves better efficiency at the expense of performance sacrifice. One natural way to mitigate performance degradation caused by quantization error is low-precision Bayesian sampling. We argue that posterior sampling is particularly suited for low-precision arithmetic because of its inherent robustness to system noise. In particular: (1) Bayesian posterior explores plausible weight space instead of converging to a single point, thus it should not require precise weights and be tolerant to non-full-precision models; and (2) it allows us to perform model averaging during inference using an ensemble of models sampled from the posterior, which compensates for the performance drop caused by coarse representations of individual models Zhu et al., 2019.

Low-Precision SG-MCMC¶

Preliminary¶

The Bayesian counterpart of QAT directly samples low-precision models (i.e., low-precision weights of the neural network) from a target distribution (typically, a Bayesian posterior distribution). A straightforward idea is to adapt existing Monte Carlo sampling algorithms for low-precision sampling tasks. In the regime of deep learning, where the training data sets are usually massive, the SG-MCMC sampling algorithm (see Chapter Stochastic Gradient MCMC) is the most suitable choice since it can avoid a full scan of the data set when calculating the likelihood terms.

Now we briefly recall two seminal SG-MCMC algorithms, the SGLD and SGHMC. Given a dataset $\mathcal D$, a statistical model with parameters $\theta\in \mathbb{R}^P$, and a prior $p(\theta)$, we aim to sample from the posterior $p(\theta| \mathcal D)\propto \exp( -U(\theta ))$, where the energy function is defined as

Stochastic Gradient Langevin Dynamics (SGLD) updates the parameter in the $(k+1)$-th step following the rule

where $\alpha$ is the stepsize, $\xi_{k+1}$ is standard Gaussian noise, and the stochastic gradient $\nabla\tilde U$ is an unbiased estimator of $\nabla U$ based on a subset of the dataset $\mathcal D$. The use of stochastic gradients helps reduce the computational cost. Compared to the Stochastic Gradient Descent (SGD) update, the only difference is that SGLD adds an additional Gaussian perturbation at each step, enabling SGLD to characterize the full distribution rather than converge to a single point. Due to this close connection, it is convenient to implement SGLD on existing deep learning tasks for which SGD is the default learning algorithm.

Stochastic Gradient HMC Chen et al., 2014 is derived by replacing $\nabla U$ with a noisy gradient $\nabla \tilde U+\xi$ in the Hamiltonian Monte Carlo, and boils down to the unadjusted underdamped Langevin algorithm with stochastic gradient, where the underdamped Langevin dynamics can be defined as

where $u$, $\gamma$ denote the hyperparameters of the inverse mass and friction, respectively. Furthermore, Cheng et al. (2018) proposed the following discretization of underdamped Langevin dynamics (3) with stochastic gradient:

where $\mathbf{\xi}_{k}^{v}$, $\mathbf{\xi}_{k}^{\theta}$ are normal distributed in $\mathbb{R}^P$ satisfying that :

Before formally introducing the low-precision SG-MCMC algorithms, we review the basic concepts of low-precision training. To represent numbers, e.g., the weights of a DNN, in low-precision, one simple way is to use fixed-point representation, which has been utilized in both theory and practice Gupta et al., 2015Li et al., 2017Yang et al., 2019. Specifically, suppose that we use $W$ bits to represent a number, where $F$ of those $W$ bits are used to represent the fractional part. Then, there is a gap between consecutive representable numbers, $\Delta=2^{-F}$, which is called* quantization gap*. In addition, the representable numbers also have a lower bound $L= -2^{W-F-1}$ and an upper bound $U = 2^{W-F-1}-2^{-F}$. Therefore, when the number of bits (i.e., $W$ and $F$) decreases, the accuracy and range of number representation decreases. This chapter will use this type of number representation in our theoretical analysis and empirical demonstration, following existing literature Li et al., 2017Yang et al., 2019.

Another popular number representation is floating point, where each number is decomposed into its own sign, exponent, and significand (or mantissa). Between fixed point and floating point,* block floating point* allows numbers within a block to share a common exponent Song et al., 2018. This chapter will also present some deep learning experiments using block floating point for deep learning experiments since it has been shown to be more favourable for deep models Yang et al., 2019.

Given the low-precision number representations, we also need a quantization function (or a quantizer) $Q$ which converts a real-valued number into a low-precision number. Popular choices include deterministic rounding and* stochastic rounding*. More specifically, the deterministic rounding function $Q^d$ quantizes a number to its nearest representable neighbour, i.e.,

where $\text{clip}(x,L,U) = \max[\min(x,U),L]$ and the above operator is applied to the whole vector elementwisely. On the other hand, stochastic rounding $Q^s$ quantizes a number to one of its representable neighbours based on a probabilistic rule:

A key design property of $Q^s$ is that $\mathbb{E}\left[Q^s(\theta)\right]=\theta$, meaning the quantized number is unbiased. $Q^s$ is generally preferred over $Q^d$ in practice since it can preserve gradient information, especially when the scale of the gradient update is smaller than the quantization gap (which is always quantized to 0 if $Q_d$ is implemented) Gupta et al., 2015Wang et al., 2018. In what follows, we utilize stochastic rounding as our quantizer, and when necessary, $Q_W$ and $\Delta_W$ denote the weights’ quantizer and quantization gap, $Q_G$ and $\Delta_G$ denote gradients’ quantizer and quantization gap.

To perform a gradient update in low-precision presentations, there are two common strategies depending on whether we keep an additional copy of full-precision weights. The SGD with Full-precision gradient accumulators (SGDLP-F) algorithm uses a full-precision weight buffer to accumulate gradient updates and only quantizes weights before computing gradients. Particularly, it updates the weights as follows,

where $\theta_{k+1}$ and $\theta_k$ are full-precision during the update, and the gradient computation is quantized for forward and backward propagation Courbariaux et al., 2015Li et al., 2017.

Gradient accumulators need to be frequently updated during training. Hence, to further reduce the computational costs, low-precision gradient accumulators represent them in low-precision format, and SGD with low-precision gradient accumulators (SGDLP-L) performs the update as follows,

where $\theta$ is always represented in low precision. In comparison, low-precision gradient accumulators are cheaper and faster due to having all numbers in low precision, whereas full-precision gradient accumulators better preserve small gradient updates, leading to better performance in general Courbariaux et al., 2015Li et al., 2017.

Low-Precision Stochastic Gradient Langevin Dynamics¶

In this section, we combine the low-precision rounding and the SGLD update (2). As shown in Equation (2), the update of a full-precision SGLD is simply a full-precision SGD update plus a Gaussian noise. Therefore, the low-precision counterpart of SGLD is naturally the low-precision SGLD training with an additional Gaussian noise in each step. Given the SGDLP-F, we can do low-precision SGLD with full-precision gradient accumulators (SGLDLP-F) as follows:

Despite this small change, the Gaussian noise turns out to help counteract the rounding noise introduced by quantization, which makes SGLDLP-F more robust to inaccurate number representation and converges better than SGDLP-F as shown in Theorem 1.

To facilitate the theoretical analysis, we assume that the target distribution is smooth and strongly log-concave, and the energy function has a Lipschitz Hessian, i.e., $\forall~ \theta,\theta'\in \mathbb{R}^P$,

These assumptions are commonly used in the literature Dalalyan & Karagulyan, 2019Yang et al., 2019Zhang et al., 2022Wang et al., 2024.

This theorem shows that SGLDLP-F converges to the accuracy floor at

given a large iteration number $K$ and small stepsize $\alpha$. Furthermore, if $\Psi=0$ (i.e., the energy function is quadratic), then SGLDLP-F indeed converges to the target distribution asymptotically. This is comparable to the optimization result that SGDLP-F converges to the optimum asymptotically on a quadratic lossLi et al., 2017. This theorem also recovers the bound of full-precision SGLD (i.e., $\Delta_W=0$) presented by Dalalyan & Karagulyan (2019). In general, when the energy function is not exactly quadratic, the convergence of SGLDLP-F to the target distribution has a $\mathcal{O}(\Delta_W^2)$ rate, but SGDLP-F to the global optimum has a $\mathcal{O}(\Delta_W)$ rate Yang et al., 2019. Because $\Delta_W = 2^{-F}$ where $F$ is the number of fractional bits, it implies that to achieve the same convergence accuracy, SGLD only needs half the number of fractional bits as SGD needs. This advantage of SGLD over SGD fits into the literature of comparing sampling and optimization convergence bounds Ma et al., 2019Talwar, 2019.

To further reduce the computational costs, we can utilize low-precision gradient accumulators and mimic the update of SGDLP-L in Equation (9). Naturally, this leads to the following update rule of SGLD with low-precision gradient accumulators (SGLDLP-L),

Unfortunately, unlike the convergence behaviour of SGLDLP-F, SGLDLP-L may diverge arbitrarily far away from the target distribution with a small stepsize.

Note that an $\alpha^{-1}$ factor appears in one term on the right-hand side upper bound; this theorem implies that as the stepsize $\alpha$ of SGLDLP-L decreases, $W_2$ distance between the actual sampling distribution and the target distribution may increase. To empirically verify it, we run SGLDLP-L on a standard Gaussian distribution in Figure Figure 1 with the stepsize $\alpha=0.001$ and 0.0001. The experiment uses 8-bit fixed-point, and 3 out of 8 bits are used to represent the fractional part. The simulations clearly demonstrate that SGLDLP-L diverges farther away from the target distribution under a smaller stepsize, whereas SGLDLP-F samples always successfully resemble the target distribution, aligning with the result in Theorem 1.

One possible remedy is to choose a stepsize that minimizes the right handed side of (15). However, this is not feasible in practice since the constants (such as $M$ and $m$) are general unknown. Moreover, enabling a small stepsize in SGLD is usually desirable, since a smaller stepsize helps reduce the asymptotic bias of the posterior approximation Welling & Teh, 2011.

One key reason causing the divergence of SGLDLP-L is that the variance of each dimension of $\theta_{k+1}$ becomes larger due to using low-precision gradient accumulators. More precisely, given the stochastic gradient $\nabla\tilde{U}$, the update of full-precision SGLD is equivalent to sampling from a Gaussian distribution

Note that both the weight quantizer $Q_W$ and the gradient quantizer $Q_G$ are unbiased stochastic rounding, hence SGLDLP-L satisfies

which shares the same mean as $\theta_{k+1}$ in full-precision. But on the other hand, the variance of $\theta_{k+1,i}$ is now larger than needed. Let’s ignore the variance caused by $Q_G$ and the stochastic gradient (since they are present and have been shown to work well in SGLDLP-F), then the variance of $\theta_{k+1,i}$ is

for some $\chi_{k+1,i}\in [0,1]$. This result shows that the variance of the SGLDLP-L update is larger than the ideal variance value $2\alpha$. Figure Figure 1 also reflects this, as the naïve SGLDLP-L estimates the mean correctly but the variance wrongly.

To correct the inflated variance and enable SGLD with low-precision gradient accumulators, Zhang et al. (2022) introduced a new variance-corrected quantization function $Q^{\text{vc}}$. The main idea of $Q^{\text{vc}}$ is to directly sample from the discrete low-precision space instead of quantizing a real-valued Gaussian sample. To begin with, if we want a sample over $\{\Delta_W,-\Delta_W,0\}$ with mean $\mu\ge 0$ and variance $v\le\Delta^2_W /4$, we could sample it with the following categorical distribution,

providing that the probability values are valid (i.e., non-negative and smaller than 1). To use this categorical distribution to preserve the correct mean and variance for quantized $\theta_{k+1}$, we consider two cases: (i) when the ideal variance $2\alpha$ is larger than the largest possible stochastic rounding variance, i.e., $\Delta_W^2/4$, $Q^{\text{vc}}$ first adds a small Gaussian noise and uses the sampler in Equation (19) to make up the remaining variance; (ii) Otherwise, $Q^{\text{vc}}$ directly samples from Equation (19) to achieve the target variance. The full description of $Q^{\text{vc}}$ is described in Algorithm 1. It is not difficult to see that the use of $\text{Cat}(\cdot,\cdot)$ sampler is valid in Algorithm 1, with the sampling probabilities being nonnegative. This variance-corrected quantizer leads to the following SGLDLP-L algorithm, named VC SGLDLP-L:

The variance-corrected quantizer $Q^{\text{vc}}$ always guarantees the correct mean, $\mathbb{E}(\theta_{k+1,i})= \theta_{k,i} - \alpha\nabla\tilde U(\theta_k)_i$, and further ensures the variance $\mbox{Var}(\theta_{k+1,i})$ matches the ideal value $2\alpha$ most of the time except when $v=2\alpha<v_s$, where $v_s$ is defined in Algorithm 1. We argue that $v=2\alpha<v_s$ rarely happens in practice since the stepsize is often very small. Although $Q^{\text{vc}}$ only preserves the correctness of the first two moments (i.e., mean and variance), it is sufficient to achieve satisfactory performance in both theory and practice, as shown by Theorem 3 and later numerical evidence.

As shown by this theorem, when the stepsize $\alpha \rightarrow 0$, the distance between the sampling distribution of VC SGLDLP-L and the target distribution no longer diverges to infinity. Instead, under proper choices of $\alpha$ and $K$, VC SGLDLP-L converges to the target distribution in $\mathcal{O}(\sqrt{\Delta_W})$, which matches the convergence rate of SGD with low-precision gradient accumulators to the optimum Li et al., 2017Yang et al., 2019.

Numerical Experiments¶

Now we present some numerical results for these Langevin-based low-precision sampling algorithms. We evaluate the generalization accuracy and uncertainty estimation of SGLDLP-F and variance-corrected VC SGLDLP-L on the ResNet-18 model on the CIFAR datasets Krizhevsky et al., 2009Krizhevsky & Hinton, 2009. We use the same quantization for weights, gradients, activations, and backpropagation errors unless otherwise stated. By default, SGLD algorithms collect samples from the approximated posterior of the model’s weights and predict labels for test data via Bayesian model averaging. The comparison baseline is the performance of low-precision SGD, which is commonly used for frequentist low-precision deep learning Sun et al., 2019Sun et al., 2020. We use SGDFP, SGDLP-F, and SGDLP-L to denote full-precision SGD, low-precision SGD with full-precision gradient accumulators, and low-precision SGD with low-precision gradient accumulators, respectively. Both fixed-point and block-floating-point (BFP) low-precision formats are implemented. The numerical results are provided in Table 1, which reports average performance over 3 runs with the standard error.

1. Fixed Point: We use 8-bit fixed point for weights and gradients but keep the activations full-precision since low-precision activations significantly harm the empirical performance. The results show that SGLDLP-F is better than SGDLP-F, and VC SGLDLP-L significantly outperforms SGDLP-L. Moreover, the improvement of SGLD over SGD gets larger under more low-precision arithmetic: on CIFAR-100, VC SGLDLP-L outperforms SGDLP-L by 2.24% and SGLDLP-F outperforms SGDLP-F by 0.54%, but SGLDFP outperforms SGDFP merely by 0.06%. This suggests that SGLD is particularly suitable for low-precision deep learning because of its natural robustness against system noise.

2. Block Floating Point: BFP is often preferred over fixed point on deep models since it causes less quantization error induced by overflow and underflow Song et al., 2018. Following the block design in Yang et al. (2019), we use small-block for ResNet experiments. Also, the $Q^{\text{vc}}$ function can naturally generalize to BFP with only a minor modification (see Zhang et al. (2022) for the detailed algorithm of $Q^{\text{vc}}$ with BFP). With BFP, the results of all low-precision methods improve over fixed-point counterparts. Specifically, SGLDLP-F performs as well as SGLDFP with all numbers quantized to 8-bit except gradient accumulators. Similar to the fixed-point case, SGLDFP-F and VC SGLDLP-L outperform low-precision SGD results. Overall, these results suggest the universal applicability of low-precision SGLD with different computer number formats.

3. Expected Calibration Error: Besides accuracies, Table 1 also reports the results of expected calibration error (ECE, which measures how well a model’s predicted probabilities match its actual empirical accuracy) Guo et al., 2017 to demonstrate the uncertainty estimation of low-precision SGLD. The summarized results indicate that SGLDLP-F and VC SGLDLP-L achieve almost the same or even lower ECE than full-precision SGLD. It proves that SGLD retains the capability to provide well-calibrated predictions while using low-precision representations. VC SGLDLP-L sometimes gives lower ECE than SGLDLP-F, which may be due to the regularization effect of low-precision arithmetic. On the other hand, the ECE of low-precision SGD increases significantly compared to its full-precision counterpart, implying that the quantization error makes the standard DNNs more overconfident, which might lead to wrong decisions for corner cases in real-world applications.

Low-Precision Stochastic Gradient Hamiltonian Monte Carlo¶

In this section, we combine low-precision arithmetic with SGHMC algorithm (4). Recall that SGHMC is equivalent to a discretized underdamped Langevin algorithm with stochastic gradient. The stationary distribution of underdamped Langevin dynamics is defined on the augmented probability space: $\pi(\theta,v)\propto \exp( - U(\theta)-\|v\|_2^2/2u)$. Therefore, our theoretical investigation will investigate the 2-Wasserstein distance^[1] between the distribution of $(\theta_K,v_K)$ and $\pi(\theta,v)$.

Adopting the update rule in equations (4), we introduce low-precision SGHMC with full gradient accumulators (SGHMCLP-F) as follows:

where the distributions of $\xi^v$ and $\xi^\theta$ follow (5). SGHMCLP-F keeps full-precision parameters $v_k$, $\theta_k$ at each iteration and quantizes them to low-precision representations before taking the gradient. Following the previous section, we still assume strong convexity of $U(\theta)$. In addition, we denote $\theta^*=\arg\min_\theta U(\theta)$ and $\rho= M/m$, and assume $\Delta_G=\Delta_W=\Delta$ for the simplicity of analysis presented in the following sections.

Theorem 1 implies that under strongly log-concave target distributions, the low-precision SGLDLP-F also achieves $\epsilon$ accuracy within $\tilde{\mathcal{O}}\left(\epsilon^{-2}\log(\epsilon^{-1})(\Delta^2+1)\right)$ iterations with learning rate of order $\epsilon^2/(\Delta^2+1)$. Thus, under strong convexity, SGHMCLP-F achieves no better convergence speed than SGLDLP-F. This is not surprising, since the quantization applied to the gradients in the full-precision gradient accumulator algorithm is equivalent to adding extra noise to the stochastic gradients. Theoretically, as shown by Cheng et al. (2018), for a strongly log-concave target distribution, SGHMC doesn’t exhibit any advantage over the overdamped Langevin algorithm, especially when stochastic gradients are used.

To reduce the computation cost, the SGHMC with low-precision gradient accumulators (SGHMCLP-L) can be trivially implemented as

Similar to SGLDLP-L, SGHMCLP-L has intrinsic flaws. It directly quantizes the weights after each update, hence a small stepsize update is often quantized to zero update of $\theta$ and $v$, resulting in the sample distribution converging to a Dirac distribution at the initial point. In such cases, ensuring convergence becomes challenging. Wang et al. (2024) investigated the convergence of SGHMCLP-L under 2-Wasserstein distance, exhibiting that the naïve SGHMCLP-L has a much worse convergence upper bound than SGHMCLP-F. To remedy this, we adapt the variance correction quantizer (Algorithm 1) to SGHMCLP-L. Let $\mathbb{V}_{v}^{hmc} = u(1-e^{-2\gamma\eta})$ and $\mathbb{V}_{\theta}^{hmc} = u\gamma^{-2}(2\gamma\eta+4e^{-\gamma\eta}-e^{-2\gamma\eta}-3)$, which are the desired variances induced from the underdamped Langevin dynamics (4). The VC SGHMCLP-L can be done as follows:

Next, we present the convergence result for VC SGHMCLP-L under a strongly log-concave target distribution.

Theorem 5 shows that the variance-corrected quantization function resolves the overdispersion problem for the naïve SGHMCLP-L algorithm. The $W_2$ distance between the sample distribution and target distribution can be arbitrarily close to $\tilde{\mathcal{O}}(\sqrt{\Delta})$. This convergence rate of VC SGHMCLP-L is the same as that of VC SGLDLP-L (see the discussion after Theorem 3).

Advantage of LP Stochastic Gradient HMC under Non-Convexity¶

So far, we have introduced the low-precision SGLD and SGHMC algorithms and explored their theoretical convergence behaviour. The theorems in previous sections don’t showcase any advantage of low-precision SGHMC over low-precision SGLDLP for strongly log-concave distributions. This doesn’t align with the common understanding of SGHMC performance. Cheng et al. (2018)Gao et al. (2022) showed that for both strongly log-concave and non-log-concave target distributions, SGHMC converges to the target faster than the best-known convergence rate of SGLD in the 2-Wasserstein distance. Beyond this, SGHMC is analogous to stochastic gradient methods augmented with momentum, which is shown to have more robust updates w.r.t. gradient estimation noise Liu et al., 2020. Since the quantization-induced stochastic error in low-precision updates acts as extra gradient noise, we expect that low-precision SGHMC can perform better than low-precision SGLD under certain scenarios beyond strong convexity of $U(\theta)$.

Wang et al. (2024) explored the convergence of low-precision SGHMC and SGLD under non-log-concave distributions. Instead of assuming $U(\theta) - U(\theta') - \nabla U(\theta')^\intercal (\theta - \theta') \ge (m/2) \left\lVert \theta - \theta' \right\rVert_2^2$, the dissipative condition is considered:

Such an assumption is commonly used Raginsky et al., 2017Zou et al., 2019Gao et al., 2022 in the analysis of sampling from non-log-concave distributions. The theorem below summarizes the theoretical findings of Wang et al. (2024), which emphasize how the convergence rate depends on the quantization gap $\Delta$.

Note that $\Delta<1$ in general, thus Theorem 6 suggests that HMC-based low-precision sampling algorithms have better dependence on the quantization error than SGLD-based algorithms for non-log-concave distributions.

Numerical Experiments¶

Similarly to the experiment settings for low-precision SGLD, we consider image tasks CIFAR-10 and CIFAR-100 on the ResNet-18. The simulation results are summarized in Table 1. Together with Table 1, it shows that across all evaluation metrics, the low-precision SGHMC has either a better or comparable performance with its SGLD counterpart.

Post-Training Low-Precision Sampling¶

The previous section focuses on training low-precision Bayesian DNN models directly by SG-MCMC from scratch. In contrast, it is feasible first to train a full-precision (Bayesian) DNN and then obtain a low-precision Bayesian DNN model via additional processing, such as fine-tuning, over the pre-trained model. More relevant content about model compression with quantization can be found in Chapter Bayesian compression: with quantization

Quantized Pre-Trained Bayesian models¶

To convert a pre-trained Bayesian DNN to its low-precision counterpart, one common approach is to quantize the sampled weights yielded by the full-precision posterior sampling (e.g., from SG-MCMC algorithms). In simple words, let $\theta^{(1)},\dots,\theta^{(T)},\dots$ be the posterior samples in $\mathbb{R}^P$, then a post-training quantization technique is applied to each sample.

This process seems trivial, but there are multiple factors that need to be considered.

Sampling Thinning¶

Sample thinning refers to the selection of a subset from the original posterior samples. That is, only a proportion of samples will be quantized and used for low-precision Bayesian model average-based inferences, for the sake of computational efficiency. One would also avoid samples from the burn-in period of posterior sampling to reduce the sampling bias. One common way is to include one sample every other $f$ samples (where $f$ denotes the thinning frequency), skipping the first $b$ samples (where $b$ denotes the burn-in period length).

To determine the best thinned sample set, one needs to assess the quality of the samples, say how well the empirical distribution determined by this sample set resembles the true posterior distribution. To this end, we employ the Kernel Stein Discrepancy Gorham & Mackey, 2017. KSD is a distance measure between two distributions, where the target distribution is continuously differentiable with an intractable normalization constant. Given a kernel function $k(\cdot,\cdot)$, we define the Stein Kernel as

Then, the KSD between (the empirical distribution of) a sample $\theta^{(1)},\dots, \theta^{(T)}$ and the target distribution is

KSD can be used to search for the optimal choices of $b$ and $f$ or the whole subset Riabiz et al., 2022. However, it is worth mentioning that the computation of the Stein Kernel scales up quickly with the dimension $d$, thus it could be very expensive for huge models.

Calibration of Quantization Range and Gap¶

If training low-precision (Bayesian) DNNs from scratch (e.g., Section Low-Precision SG-MCMC), one usually needs to determine the quantization scheme (e.g., the quantization gap $\Delta$ in the quantization functions $Q^d$ or $Q^s$) beforehand. In contrast, when a pre-trained model is available (e.g., a full-precision posterior sample $\theta^{(t)}$), we have direct access to the distribution of its weight values; we can use a representative dataset to approximate the distribution of activations along the forward propagation. This information can help us choose an appropriate quantization scheme to quantize this $\theta^{(t)}$ that fits the computational / memory budgets and has a small impact on its predictive performance. Note that in Section Low-Precision SG-MCMC, we only consider a rounding mechanism that maps a high-precision (e.g., full floating-point value) to its near low-bit fixed-point values. But in general, any mapping from real values to a discrete space can be used. For example, an uniform quantization Jacob et al., 2018 with $W$-bit quantization is defined as

where $\lfloor\cdot\rceil$ denotes rounding to the nearest integer. Then $a$ and $b$, which denote the quantization range, can be set as the minimum and maximum values among entries of the pre-trained model $\theta^{(t)}$.

Parameter Fine-tuning¶

Although a quantized network $Q(\theta^{(t)})$ for some quantization function $Q$ is ready for inference, it may suffer significant performance drops. In such a case, one can first fine-tune the model $\theta^{(t)}$ for a small number of gradient descent steps to compensate for the performance drops caused by quantization, then apply quantization to the fine-tuned model.

Let $f(x;\theta)$ be the neural network model parameterized by the weight vector $\theta$, then the empirical loss is defined as

for some loss function $l(\cdot,\cdot)$. Instead of running fine-tuning gradient descent on $\mathcal {L}(\theta)$, we design a new neural network model $f^q(\cdot;\theta)$ which approximates $f(\cdot;\theta)$ with quantization applied on weights and activations. For example, Jacob et al. (2018)Ferianc et al. (2021) insert simulated quantization modules after weights input and every tensor product operation. Therefore, running gradient descent on $\mathcal L^q = \sum l(y, f^q(x;\theta))$ initialized from $\theta^{(t)}$, leads to a nearby model whose after-quantization performance is acceptable.

Quantize Pre-trained Variational Bayesian DNNs¶

If a pre-trained variational Bayesian neural network model (refer to Book Part %s for details) is given, one may additionally quantize the variational parameters Ferianc et al., 2021Subedar et al., 2021. Suppose that the variational posterior distribution of one weight entry $\theta_{(i)}$ is $\mbox{N}(\mu,\sigma^2)$, then given a quantization function $Q$, the quantized posterior sample of $\theta_{(i)}$ is generated by $Q(Q(\mu)+Q(\sigma)\xi)$ where $\xi$ is a standard Gaussian variable. To accommodate performance degradation induced by quantization, one can also incorporate additional fine-tuning for the variational parameters Ferianc et al., 2021, similarly to the procedure discussed in the above paragraph.

Low-Bit Sub-Distribution¶

As discussed earlier, a pre-trained full-precision model can provide valuable insights into the distribution of weight values and inform the design of a quantization scheme. The SG-MCMC algorithms in Section Low-Precision SG-MCMC yield posterior samples whose entries belong to an equally spaced discrete space $\mathcal S=\{0, \pm\Delta, \pm2\Delta,\dots\}$. When a pre-trained model is available, one can train a flexible discrete space $\mathcal S$ from which the posterior samples are drawn. The elements in $\mathcal S$ can be full floating-point values; thus, model weights of a neural network achieve low-bit representation by sharing the same full precision values in $\mathcal S$. This can be viewed as a generalization of sparse DNN, which encourages network weights to share a common value of 0.

Given a pre-trained DNN model $\theta$, a straightforward method is to implement a clustering algorithm over the entries of $\theta$ and obtain the cluster centers $\mathcal S=\{q_1, \dots, q_M\}$. Then we would like to restrict the posterior samples to $\mathcal S$. To accomplish this, we can adapt the mean-field variational Bayes (refer to Chapter Introduction to variational inference for details) and define the discrete variational distribution as

where variational parameters satisfy $\sum_m p_{im}=1$. To facilitate the training of variational inference, we can introduce parametric modelling for $p_{im}$ (e.g., $p_{im}$ relates to $|\theta_{(i)}-q_m|$, where $\theta_{(i)}$ is the pre-trained weight). To enhance the model performance, we can also allow $q_1,\dots, q_M$ to update in the variational training procedure Dong et al., 2022.

It is worth mentioning that no actual low-precision sampling happens in the variational training process since variational inference boils down to an optimization problem for the variational parameters. Once these parameters are learned, the (approximate) posterior samples are easily drawn from (37).

Conclusions and Discussions¶

This chapter introduced low-precision sampling algorithms that are useful for modern Bayesian deep learning applications, including SG-MCMC-based direct low-precision samplers and post-training low-precision sampling techniques. These low-precision methods are crucial for deploying large models on resource-constrained hardware and improving Bayesian inference speed.

As discussed in the introduction section, one main reason for pursuing low-precision arithmetics is to reduce the memory and computation costs for DNN training or DNN inference. Beyond this, the low-precision representation is also helpful to reduce communication cost for scenarios where frequent information transfer among several parties is required, but only with limited bandwidth. One such emerging research field is federated learning (FL) Kairouz & McMahan, 2021. Federated learning is a learning paradigm where multiple parties (e.g., hospitals or edge devices) collaboratively train a global model while keeping the local training data private and decentralized. In a federated SGLD algorithm, each party computes the gradient based on its local data, and all local gradients are shared and aggregated via communication. The aggregated gradient is then used in the SGLD update (2) of the global model. Low-precision representations can be used to compress the gradient communications. Vono et al. (2022) utilized a stochastic rounding mechanism and studied how quantization error impacts the convergence of the global Bayesian model.

The theorems and algorithms in this chapter are mostly depicted regarding the theme that quantization of the weight parameter $\theta$ suffices. However, neural networks have a multilayer structure, and quantization can be employed on each layer along the forward propagation process (i.e., online conversion of activation values to low-precision representations). In addition, normalization layers, such as Batch Normalization and Layer Normalization, are often sensitive to quantization error and may require different techniques (e.g., BatchNorm folding) to mitigate the impact of quantization. These issues are less explored and discussed in the field of Bayesian low-precision DNNs. Interested readers may refer to Gholami et al. (2022)Kim et al. (2023).