SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Christopher Subia-Waud
School of Electronics & Computer Science
Southampton, UK
cc2u18@soton.ac.uk
Srinandan Dasmahapatra
sd@ecs.soton.ac.uk
Weight-sharing quantization has emerged as a technique to reduce energy expenditure during inference in large neural networks by constraining their weights to a limited set of values. However, existing methods often assume weights are treated solely based on value, neglecting the unique role of weight position. This paper proposes a probabilistic framework based on Bayesian neural networks (BNNs) and a variational relaxation to identify which weights can be moved to which cluster center and to what degree based on their individual position-specific learned uncertainty distributions. We introduce a new initialization setting and a regularization term, enabling the training of BNNs with complex dataset-model combinations. Leveraging the flexibility of weight values from probability distributions, we enhance noise resilience and compressibility. Our iterative clustering procedure demonstrates superior compressibility and higher accuracy compared to state-of-the-art methods on both ResNet models and the more complex transformer-based architectures. In particular, our method outperforms the state-of-the-art quantization method top-1 accuracy by  $1.6\%$  on ImageNet using DeiT-Tiny, with its 5 million+ weights now represented by only 296 unique values. Code available at https://github.com/subiawaud/PFWN.
Weight-sharing quantization is a technique developed to lower the energy costs associated with inference in deep neural networks. By constraining the network weights to take on a limited set of values, the technique can significantly reduce the data-movement costs within hardware accelerators, which represent the primary source of energy expenditure (DRAM read costs can be as much as 200 times higher than multiplication costs [Horowitz, 2014, Jouppi et al., 2017, Sze et al., 2017]). Storing the weight values close to computation and reusing them multiple times also becomes more feasible, thanks to the limited range of possible values. Various studies have explored the effectiveness of weight-sharing quantization, including [Ullrich et al., 2017, Subia-Waud and Dasmahapatra, 2022, Wu et al., 2018b, Han et al., 2015, 2016].
Weight-sharing quantization approaches face a common challenge in determining how much a single weight can be moved to a cluster center. Traditional methods rely on the magnitude or relative movement distances between weight and cluster to decide which weights can be moved to which clusters [Han et al., 2015, Subia-Waud and Dasmahapatra, 2022], irrespective of which filter or layer in the network the weight is located. However, this assumption neglects the fact that weights with the same value may have different roles in the network, and their placement within the architecture may affect their likelihood of being moved to a particular cluster. We posit that context-dependent movement of weights to clusters can, instead, better preserve the representational capacity of the
Figure 1: An overview of the PwFN process.
network while reducing its complexity. We build upon previous attempts [Wu et al., 2018b, Achterhold et al., 2018] to use a probabilistic framework to do so. The idea is to use a probability distribution to capture the flexibility in weight values, informing clustering decisions which reduce the entropy of the network and lower the unique parameter count without performance degradation.
We look to make progress on this problem with the use of Bayesian neural networks and a variational relaxation to identify optimal cluster configurations for the compression of neural networks with a method we call probabilistic weight fixing networks (PwFN). By incorporating the insights from the previous works on minimizing the relative weight-to-cluster distance [Subia-Waud and Dasmahapatra, 2022] and favoring additive powers-of-two cluster centroid values [Li et al., 2019b], we propose a novel initialization setting. Further, we discover that a simple regularization term which encourages maximal noise resilience is sufficient to avoid the variance of the weights' distribution from collapsing to zero. This regularization term facilitates the training of model-dataset pairings that were previously considered intractable when employing the variational Bayes-by-backdrop approach [Blundell et al., 2015]. These changes, when combined with a novel iterative clustering procedure, enable us to achieve superior compressibility. We can represent ResNet family models trained on the ImageNet dataset with fewer unique parameters and a reduced weight-space entropy while maintaining higher accuracy than the current state-of-the-art. Furthermore, we demonstrate the effectiveness of applying PwFN to transformer-based architectures, again achieving state-of-the-art results.
Weight-sharing quantization. Consider a neural network that consists of  $N$  weights  $\boldsymbol{w} = \{w_{1}, w_{2}, \ldots, w_{N}\}$ . Each weight is typically unique and can take up to  $2^{b}$  distinct values, where  $b$  is the bit-width of the number system used. However, this poses a challenge since each time a weight is used, at least one memory read is required. Since memory reads dominate the computational cost of inference, it is desirable to reduce them. Weight-sharing quantization addresses this challenge using a reduced pool of cluster centers  $\boldsymbol{c} = \{c_{1}, c_{2}, \ldots, c_{k}\}$ , with  $k \ll N$  and defining a map  $\boldsymbol{w} \rightarrow \boldsymbol{c}$ , where each weight  $\boldsymbol{w}_{i} \in \boldsymbol{w}$  is mapped to a cluster center:  $w_{i} \mapsto c_{k} \in \boldsymbol{c}$ .
Two considerations are pivotal for formulating this process: determining which values should be included in the reduced cluster pool  $c$ , and identifying the mapping of weights in  $w$  to the clusters in  $c$ .
Which values should be in the cluster pool? Insights from the literature highlight that pre-trained weight distributions often adhere to a heavy-tailed distribution [Barsbey et al., 2021, Hodgkinson and Mahoney, 2021, Gurbuzbalaban et al., 2021]. Moreover, hardware implementations show a preference for powers-of-two values for weights due to their cost-effective multiplication properties
[Vogel et al., 2019, Przewlocka-Rus and Kryjak, 2023, Lee et al., 2017]. While some methods advocate for exclusive use of powers-of-two as cluster centers [Zhou et al., 2017], others propose a more flexible approach, suggesting additive powers-of-two [Li et al., 2019b, Subia-Waud and Dasmahapatra, 2022] – but only if it doesn't compromise performance. In our study, we align with this perspective, populating the cluster pool with powers-of-two values, but making exceptions for additive powers-of-two when they enhance performance.
How to identify which of the weights should be moved to which cluster? Previous studies have shown that distance metrics can be utilized to determine which fixed clusters can accommodate certain weights without negatively affecting the performance. Weight-sharing clustering techniques can rely on these metrics to assign weights to clusters. The commonly used distance metrics include Euclidean distance [Wu et al., 2018b] and Relative movement distance [Subia-Waud and Dasmahapatra, 2022]. However, these metrics implicitly assume that all weights with the same value should be treated equally, irrespective of their location in the network. This may not be valid as moving a small (large) weight by a small (large) distance affects classification outcome depending on where the weight is in the network. To overcome this limiting assumption, we apply in this paper a Bayesian neural network (BNN) training approach to obtain a metric to determine the allowed movement of weights to clusters.
Bayesian Neural Networks. BNNs attempt to use the methods of Bayesian inference in modeling predictive problems. Rather than the weights in a network coming from point estimates (i.e., a single value for each weight), a BNN attempts to model many (ideally all) configurations of weight values throughout a network and make predictions, weighting each configuration by its probability. Exact Bayesian inference on the weights would require the computation of the integral  $P(\pmb{y}|\pmb{x}, D) = \int P(\pmb{y}|\pmb{x}, \pmb{w}) P(\pmb{w}|\pmb{D}) d\pmb{w}$  where predictions for each allowed  $\pmb{w}$  are averaged over. Unfortunately, the marginalization over  $P(\pmb{w}|\pmb{D})$  is intractable for even simple networks, so approximations are needed. Approaches to this include Laplace approximation [MacKay, 1992], gradient MCMC [Welling and Teh, 2011], expectation propagation updates [Hernández-Lobato and Adams, 2015], and variational methods such as Bayes-by-backprop (BBP) [Blundell et al., 2015]. Using dropout at inference time has also been shown to be a variational approximation [Gal and Ghahramani, 2016] which is much more amenable to scale – albeit with reduced expressive power [Jospin et al., 2022]. BBP, of which our approach uses a variation of, models each of the network weights as coming from a Gaussian distribution and uses the re-parameterization trick for gradient descent to learn the parameters of the weight distributions. We give the derivations to do so in the Appendix. BBP has been demonstrated to work well in more complex model-dataset combinations than other BNN approaches (aside from dropout) but is still not able to scale to modern architectures and problems such as the ImageNet dataset and Resnet family of networks [Jospin et al., 2022].
One reason for the limited applicability is that the weight prior, which is a scaled mixture of two zero-mean Gaussians [Blundell et al., 2015], is not a well-motivated uninformative prior [Fortuin, 2022]. Instead, we initialize our training from a pre-trained network, in the spirit of empirical Bayes, and propose a prior on weight variances that encourages noise-resilience. A noise-resilient network can be thought of as one with a posterior distribution with large variances at its modes for a given choice of weight coordinates. Thus, small changes in weight space do not meaningfully alter the loss. The corresponding characterization of the flat minima in the loss landscape is well supported in the optimization literature [Foret et al., 2020, Kaddour et al., 2022].
Quantization. Quantization is a technique used to reduce the number of bits used to represent components of a neural network in order to decrease energy costs associated with multiplication of 32-bit floating-point numbers. There are different approaches to quantization, such as quantizing weights, gradients, and activations. Clipping+scaling quantization maps the weight  $w_{i}$  to  $w_{i}^{\prime} = \text{round}\left(\frac{w_{i}}{s}\right)$ , where round() is a predetermined rounding function and  $s$  is a scaling factor learned channel-wise, layerwise or with even further granularity. Quantization can be performed post-training or with quantization-aware training [Jacob et al., 2018].
Weight sharing quantization uses clustering techniques to group weights and fix the weight values to their assigned cluster centroid. These weights are stored as codebook indices, which enables compressed representation methods like Huffman encoding to compress the network further. Unlike
Figure 2: For DeiT small, we show a box plot of the entropies and unique counts per input channel for each Q,K,V by layer and with the mean of each layer (calculated across all attention heads) shown in the black lines.
clipping+scaling quantization, weight sharing techniques share the pool of weights across the entire network. Several studies have proposed different methods for weight sharing quantization. For example, Zhou et al. [Zhou et al., 2017] restrict layerwise rounding of weights to powers-of-two for energy-cheap bit-shift multiplication, while [Li et al., 2019a] suggest additive powers-of-two (APoT) to capture the pre-quantized distribution of weights better. [Wu et al., 2018a] use a spectrally relaxed k-means regularization term to encourage the network weights to be more amenable to clustering and focus on a filter-row codebook for convolution. Similarly, [Stock et al., 2020] and [Fan et al., 2020] focus on quantizing groups of weights into single codewords rather than the individual weights themselves. [Ullrich et al., 2017] use the distance from an evolving Gaussian mixture as a regularization term to prepare the clustering weights. However, their approach is computationally expensive. In contrast, our formulation reduces computation by adding cluster centers iteratively when needed and only considering weights that are not already fixed for a regularization prior. A work closely related to ours is that of [Achterhold et al., 2018], where the authors formulate post-training quantization and pruning as a variational inference problem with a 'quantizing prior' but due to its optimization complexity and difficulties with weights not being drawn tightly enough into clusters the method was only demonstrated to work on simple dataset-model combinations and for the case of ternary quantization. WFN [Subia-Waud and Dasmahapatra, 2022] is a recent weight-sharing approach that uses the minimizing of the relative movement distance to determine which weights are to be clustered and has demonstrated that a low weight-space entropy and few unique weights with a single codebook can maintain accuracy. In our work, we go much further in reducing weight-space entropy and unique weight count by using a context-aware distance metric and probabilistic framework in determining which weights can be moved where. Additionally, almost all of the works reviewed do not quantize the first and last layers [Li et al., 2019a, Jung et al., 2019, Zhou et al., 2016, Yamamoto, 2021, Oh et al., 2021, Li et al., 2022] in order to maintain performance and in some cases don't quantize the bias terms [Li et al., 2022], - we challenge this view and attempt a full network quantization, leaving only the batch-norm and layer-norm parameters at full precision.
3 PwFN
In PwFN, we follow  $T$  fixing iterations each of which combines a training and a clustering stage in order to reach an outcome of a highly compressed/quantized network with a single whole-network codebook. We model each individual weight  $w_{i}$  as a draw from a distribution with learnable parameters and use these parameters as guidance for the clustering stage. We model each  $w_{i} \in \mathbf{w}$  as coming from a Gaussian distribution  $\mathcal{N}(\mu_i,\sigma_i)$  and during the training stage we use a form of BBP to train the weight distribution parameters  $\pmb{\mu} = (\mu_1,\dots ,\mu_N)$  and  $\pmb{\sigma} = (\sigma_{1},\dots ,\sigma_{N})$  to minimize both the task performance loss and to encourage the weight distributions to tell us exactly how much noise can be added to each weight without affecting performance. Both  $\pmb{\mu}$  and  $\pmb{\sigma}$  are trained with an additional regularization term that encourages larger values of  $\pmb{\sigma}$  to counter the model reverting to the point estimates with  $\sigma_{i} = 0$  to minimize the classification loss. During the clustering stage, we look
Figure 3: The regularization term acts to stop the  $\sigma$  uncertainty values from collapsing to zero. This experiment is run using the Cifar10 dataset with ResNet-18, stopping after 30 epochs.
to use this information to move the  $\mu_{i}$  values to one of a handful of cluster centers. We favour the cluster centers to be hardware multiplication-friendly powers-of-two, or additive-powers-of-two. After  $T$  iterations of training and clustering, each of the weights' distributions in the networks will have their  $\pmb{\mu}$  values centered on one of the  $k$  clusters  $c$  in the codebook.
After the  $T$  fixing iterations there are two options depending on the downstream usage of the network: either the network can be converted into point estimates and the weights set to the exact  $\mu$  values giving us a quantized network. Or, we can use the extra information given to us by modelling each weight as a distribution as a way of quantifying uncertainty of a particular prediction. If, after multiple samples of  $\boldsymbol{w}$ , a model changes its prediction for a fixed input, this tells us that there is uncertainty in these predictions – with this information being useful for practical settings.
PWFN Training. Consider a network parameterized by  $N$  weights  $\boldsymbol{w} = \{w_1, \dots, w_N\}$ . In PWFN, each weight  $w_i$  is not a single value but is instead drawn from a distribution  $w_i \sim \mathcal{N}(\mu_i, \sigma_i)$ , and instead of learning the  $w_i$  directly, the learning process optimizes each  $\mu_i$  and  $\sigma_i$ . In a forward pass, during training, we sample weight values  $w_i$  according to its distribution:
$$
w _ {i} = \mu_ {i} + \sigma_ {i} \epsilon , \epsilon \sim \mathcal {N} (0, 1). \tag {1}
$$
The forward pass is stochastic under fixed  $\mu, \sigma$ . If trained correctly, the  $\sigma_{i}$  values give us information about the amount of noise a particular weight  $w_{i}$  can handle without affecting performance. Said another way, if we can find a configuration  $\boldsymbol{w} = (\boldsymbol{\mu}, \boldsymbol{\sigma})$  which maintains task performance despite the randomness introduced by the  $\sigma_{i}$  parameters, then we will know which of the corresponding weights can be moved and to what degree. In PWFN, we train  $\mu, \sigma$  following the BBP optimization process [Blundell et al., 2015] with some changes both in terms of initialization and the priors on  $\mu$  and  $\sigma$ .
Large  $\sigma$  constraint for  $w$ . Given the usual cross-entropy or other performance loss, there is a clear direction of travel during gradient descent towards having small  $\sigma_{i}$  values and less uncertainty in the network parameters. A prior on the distribution of weights is therefore needed to prevent the  $\sigma = 0$  point estimate solution being found which would leave us with no weight movement information. In the original BBP set-up, the authors looked to prevent vanishing variance by regularising the distribution of weights according to a prior distribution of a mixture of zero-mean Gaussian densities with different variances (the parameters of the prior they find through an exhaustive search). The motivation for doing so was because the empirical Bayes approach didn't perform well due to the network favoring updating these parameters over the posterior (since there are fewer) and the link to the successful spike-and-slab prior [Mitchell and Beauchamp, 1988] — where values are concentrated around 0 (the slab) or another value known as the spike – favoring sparsity. Instead, we hypothesize that a good network can handle the most noise injection whilst maintaining performance. These networks are likely more compressible since they have been trained to accept changes to their weight values without performance degradation during training.
Figure 4:  $p_t$  follows the same schedule as [Subia-Waud and Dasmahapatra, 2022] (left). In the middle and right plots, we see that PwFN achieves very small entropy values by majority of weights to only a very small (4 or 5) cluster values and the rest are assigned as outliers, most of which are powers-of-two.
We attempt this by making our  $\sigma$  values to be large. Networks with large  $\sigma$  have, probabilistically, more noise added to the  $\mu$  values during training and so have to learn to have robust performance under such circumstances. We note that this acts as a push-pull relationship with the performance loss, which favours low  $\sigma$  values. The motivation is that, much like  $L_{1}$  norms enforcing sparsity, this formulation will train the network to produce a large  $\sigma_{i}$  for noise-resilient parameter  $w_{i}$ , whilst maintaining a noise-sensitive weight  $w_{j}$  to have a small  $\sigma_{j}$  despite the prior pull. The regularised loss function for training the training phases of the algorithm is:
$$
- \log P (\mathcal {D} | \boldsymbol {\mu}, \boldsymbol {\sigma}) + \alpha \mathcal {L} _ {\mathrm {R E G}} (\boldsymbol {\sigma}), \tag {2}
$$
where the regularization term is:
$$
\mathcal {L} _ {\mathrm {R E G}} (\boldsymbol {\sigma}) = \sum_ {i = 1} ^ {N} \mathcal {L} \left(\sigma_ {i}\right) = - \sum_ {i = 1} ^ {N} \left(\sigma_ {i} - S\right) \Theta (S - \sigma_ {i}), \tag {3}
$$
with  $\Theta(x) = 1$  for  $x \geq 0$  and 0 otherwise. The  $\Theta$  function prevents the optimization from finding a network with a subset of  $\sigma$  with infinitely large values and dominating the cross entropy term.  $S$  is thus a cutoff on how large the values in  $\sigma$  can be.  $\alpha$  is a hyperparameter controlling the formation of a noise-resilient network where the majority of the weights can receive noise injection without hurting performance, not just a few. In Figure 3 we illustrate the effect on the distribution of  $\sigma$  under different  $\alpha$  values for a ResNet-18 trained on the Cifar-10 dataset. As we increase  $\alpha$  the  $\sigma$  values no longer collapse to zero giving us information for downstream clustering.
Initialization using Relative Distance from Powers-of-two. For each weight  $w_{i}$  we need to specify its prior distribution so as to derive the posterior using Bayesian updating. We assume that the posterior distribution is Gaussian with a diagonal covariance matrix:  $P(w_{i};\mu_{i},\sigma_{i})$  whose parameters  $\mu_i,\sigma_i$  are trained using BBP. To initialize the prior distributions for  $\mu_{i}$  and  $\sigma_{i}$  we set  $P^0 (\mu) = \prod_iP^0 (\mu_i)$  where  $P^0 (\mu_i)\propto \delta_{\mu_i,w_i}$  for the pre-trained weight value  $w_{i}$ . For a Gaussian posterior we would typically require an unknown  $\sigma$  to be drawn from a Gamma conjugate prior distribution. Instead, we set  $\sigma_{i}$  to be a known function of the  $\mu_{i}$  at initialization. In [Subia-Waud and Dasmahapatra, 2022] relative distances to the preferred powers of two values for the weight was used to determine weight movement. To favour anchoring weights at powers of two, we set the standard deviations to be smallest  $(2^{-30})$  at either edge of each interval between the nearest integer powers of two,  $(2^{x_i}\leq \mu_i\leq 2^{x_i + 1})$  for integer  $x_{i}$ , and largest at the midpoint of the interval. We introduce a parabolic function  $\sigma_{i}(\mu_{i})$  as a product of relative distances of each pre-trained weight value  $(\mu_{i})$  to the nearest lower and upper powers of two:
$$
\sigma_ {i} \left(\mu_ {i}\right) = (0. 0 5) ^ {2} \left(\frac {\left| 2 ^ {x _ {i}} - \mu_ {i} \right|}{\left| 2 ^ {x _ {i}} \right|}\right) \left(\frac {\left| \mu_ {i} - 2 ^ {x _ {i} + 1} \right|}{\left| 2 ^ {x _ {i} + 1} \right|}\right), \tag {4}
$$
(Full details and visualization is given in the appendix).
PwFN Clustering. In Figure 1 we show a schematic of the clustering stage in which we use the information garnered from the weights' distribution parameters to identify cluster centers and their
assignment. PwFN clustering is a two-step method running for  $t = 1, \dots, T$  iterations. At each step we set a fraction  $p_t$  of the weights to be fixed, so that  $|W_{\mathrm{fixed}}^t| = Np_t$ . The remaining weights at iteration stage  $t$  are trainable and called  $W_{\mathrm{free}}^t$ . We follow the scheme first proposed in [Subia-Waud and Dasmahapatra, 2022] in setting  $p_t$  (Figure 4, left). All of the weights  $w_i$  that are assigned to  $W_{\mathrm{fixed}}^t$  will have their  $\mu_i$  values fixed to one of the set of cluster centers. At the last iteration,  $|W_{\mathrm{free}}^T| = 0$  and  $p_T = 1$ , as all weights have been fixed to their allocated cluster centroids.