Probabilistic Weight Fixing: Large-scale training of neural network weight uncertainties for quantisation.
Authors: Chris Subia-Waud
Srinandan Dasmahapatra
Venue: NeurIPS 2023 poster (2023)
Keywords: Quantization
compression
bayesian neural networks
accelerators

Abstract: 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/PWFN.

Probabilistic Weight Fixing: Large-scale training of neural network weight uncertainties for quantization
Christopher Subia-Waud
Srinandan Dasmahapatra
School of Electronics & Computer Science University of Southampton
School of Electronics & Computer Science University of Southampton Southampton, UK sddecs soton.... ac uk
Southampton, UK cc2u180soton. ac.uk
Abstract
Weight-sharing quantization has emerged as a technique to reduce energy expen- diture 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 variational relaxation identify which weights be moved which cluster center and to what degree based on their individual position-specific learned uncertainty distributions. We introduce 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 distri- butions, we enhance noise resilience and compressibility. Our iterative clustering procedure demonstrates superior compressibility and higher accuracy compared state-of-the-aas art methods on both ResNet models and the more complex transformer- based architectures. particular, our method outperforms the oftte-of-hhe-at quantization method top-1 accuracy by 1.6% on ImageNet using DeiT-Tiny, with million+ weights now represented by only 296 unique values. Code available at https:::/// //g1 thub. com/
Introduction
Weight-sharing quantization is a technique developed to lower the energy costs associated with inference deep neural networks. By constraining the network weights take on limited set of values, the technique can significantly reduce the data-movement costs within hardware accelerators, which represent the primary source energy expenditure (DRAM read costs can as much as 200 times higher than multiplication costs [Horowitz] 2014 Douppi 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 efffcctiveness of weight-sharing quantization, including (Ullrich et al. 2017 Wu et al. 2018b Han al. 2015 2016).
Subia- Waud and Dasmahapatra 2022
Weight-sharing quantization approaches face a common challenge in determining how much single weight can be moved to 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-dess -deeendn movement of weights to clusters can, instead, better preserve the representational capacity of the
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

IOOEM TID AUDE POEE doretoe
IRIVIT BLEET ALETO
5.75
0.50
0.25
0.00
0.25
0.50
2.75
1.00
0.75
0.50
0.25
0.00
0.25
0.50
0.75
100
1
The
-0.75
-0.50
-0.25
0.25
0.50
5.35
1.00
-0.25
-0.50
0.25
0.75
1.50
60
Figure I: An overview of the PWFN process.
network while reducing its complexity. We build upon previous attempts WW al. 2018b Achterhold al., 2018 use probabilistic framework do The idea is 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 variational relaxation identify optimal cluster configurations for the compression of neural networks with method we call probabilistic weight fixing networks (PWFN). By incorporating the insights from the previous works minimizing the relative weight-to-cluster distance (Subia- Waud and Dasmahapatra 2022] and favoring additive powers-of two cluster centroid values [Li et al.) (019b), we propose a novel initialization setting. Further, maximal noise resilience is sufficient to avoid the variance the weights' distribution from collapsing zero. This regularization term facilitates the training of model-dataset pairings that were previously considered intractable when employing the variational 2015). These changes, when combined with novel iterative clustering procedure, enable achieve superior compressibility. We can represent ResNet family models trained on the ImageNet dataset with fewer unique parameters and reduced weight-space entropy while maintaining higher Furthermore, we demonstrate the effectiveness of applying PWFN transformer- -aased architectures, again achieving state-of-the- results.
discover that simple regularization term which encourages
beees-bb-backkrop approach [Blundell et al.
accuracy than the current
Weight-sharing quantization. Consider a neural network that consists of weights w {wi, W....., WN' }. Each weight is typically unique and can take up to 2 distinct values, where b is the bit-width of the number system used. However, this poses a challenge since each time weight is used, at least one memory read is required. Since memory reads dominate the computational cost of inference, is desirable to reduce them. Weight-sharing quantization addresses this challenge using a reduced pool of cluster centers e {c1, {c1,02,.... C2...., Ck), with k < N and defining a map w + c, where each weight w, E w is mapped to a cluster center: + E 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
Which values should be in the cluster pool? Insights from the literature highlight that pre-trained
weight distributions often adhere to heavy-tailed distribution (Barsbey et al. 2021 Hodgkinson and Mahoney] 2021 Gurbuzbalaban et al. 202T]. Moreover, hardware implementations show preference for pfwers-of-two values for weights due to their cost-effective multiplication properties

---

Przewlocka-Ree 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 more flexible approach, suggesting additíve powers-of-two [Li et al. 2019h (Subia Waud and Dasmahapatra [022] - but only if it doesn't t compromise performance. In our study, we align with this perspective, populating the cluster pool with powers-of-twe values, but making exceptions for additive powers-of-two when they enhance performance.
Vogel et al. 2019
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, [202]) 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 Bayesian neural network (BNN) training approach to obtain a metric to determine the allowed movement of weights to clusters.
2 Related Work
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., 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(y (x, D) P(ylı, w) P( [w| D)w where predictions for each allowed w are averaged over. Unfortunately, the marginalization over P(w) D) is intractable for even simple networks, so approximations are needed. Approaches to this include Laplace approximation [MacKay, 1992]. gradient MCMC (Welling and 201T). expectation propagation updates Bermández-Lobate and Adams 2015). and variational methods such as Bayes-by-bar bybckkprpp (BBP) (Blundell et al. [2015] Using dropout at inference time has also been shown to be a variational approximation Ga and Ghahramani 2016 which is much
more amenable 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- parameterizaaion 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 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 well-motivated uninformative prior [Fortuin 2022]. Instead, we initialize our training from a pre-trained network, in the spirit of empirical Bayes, and propose prior weight variances that encourages noise-resi 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 neural network in order to decrease energy costs associated with multiplication of 32- bit floating-port pointt numbers. There are different approaches to quantization, such as quantizing weights, gradients, and activations. Clipping+scaling quantization maps the weight w; to w round( ( ), where round() is a predetermined rounding function and s is a scaling factor learned woanell-wise. layerwise 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 Huffmman encoding compress the network further. Unlike
3

---

Grund
Layer
In PWFN, we follow T' fixing iterations each which combines a training and clustering stage in order to reach an outcome of highly compressed/quanized network with single whole-network codebook. We model each individual weight w; as a draw from a distribution with learnable parameters and use these parameters as guidance for the clustering stage. We model each coming from a Gaussian distribution A ((4,, o,) and during the training stage we use a form of BBP to train the weight distribution parameters f (Pi..... (NN) and o = (Ø1...., .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 μ and are trained with an Figure 2: For DeiT small, we show box plot of the entropies and unique counts per input channel for each Q.K.V 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., [017]] restrict layerwise rounding of weights to pfwers-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. 20188) use a spectrally relaxed k-means regularization term to encourage the network weights be more amenable to clustering and focus on a filter-row codebook for convolution. Similarly, (Stock et al. 2020 and Fan 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 regularization term 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 mateset-model combinations and for the case of ternary quantization. WFN (Subia-Waar 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 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 1 [i et al. 2022 in order to maintain performance and in some cases don't t quantize the bias terms [Li al.) 2022]
challenge this view and attempt full network quantization, leaving only the batch-norm and layer-norm parameters at full precision.
3 PWEN
E w as
additional regularization term that encourages larger values of to counter the model reverting the point estimates with = 0 to minimize the classification loss. During the clustering stage, we look

---

0.0
0.00781
0.03125
0.0030
0.0030
0.0030
0.0025
0.0025
0.0025
0.0015
0.0015
0.0015
0.0010
0.0010
0.0010
0.0005
0.0005
0.0005
0.0000
0,0000
0.1
0.0
0.1
0.2
0.1
0.0
0.1
0.2
0.0
0.1
0.2
Figure 3: The regularization term stop the a uncertainty values from collapsing This experiment is run using the Cifar 10 dataset with ResNet-18, 18, stopping after 30 epochs.
use this information move the 14; values cluster centers to be hardware multiplication-friendd After T iterations of training and clustering, each of the weights' distributions in the networks will have their la values centered on one of the k clusters in the codebook.
of handful of cluster centers. We favour the powers-of-two, or additive-powers-of-tye
After the fixing iterations there are two options depending on the downstream usage of the network: either the network be converted into point estimates and the weights set the exact fr values giving us a quantized network. Or, we can use the extra information given to us by modelling each weight distribution multiple samples of 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.
way of quantifying uncertainty of particular prediction. If, after
PWEN Training. Consider a network parameterized by N weights {wi, ... WN). J. In PWEN, each weight w; is not a single value but instead drawn from distribution w; ~ N N(u, o;), and instead of learning the w, directly, the learning process optimizes each μ4 and o;. In a forward pass, during training, we sample weight values w; according to its distribution:
W = Pi + o.C, 1).
(1)
The forward pass is stochastic under fixed μr, a. If trained correctly, the a; values give information about the amount of noise particular weight w; can handle without affecting performance, Said another way, if we can find configuration the randomness introduced by the parameters. then we will know which of the corresponding weights can be moved and what degree. In PWFN, we train μt, following the BBP optimization process [Blundell et al., 2015] with some changes both in terms of initialization and the priors on and a.
(μ, or) which maintains task performance despite
Large a 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 a, values and less uncertainty in the network parameters. A prior on the distribution of weights is therefore needed to prevent the point estimate solution being found which would leave with weight movement information. In the original BBP set-up, the authors looked to prevent vanishing variance by regularising the distribution of weights according to prior distribution of mixture zero-mean Gaussian densities with different variances (the parameters of the prior they find through an exhaustive search). The motivation for doing 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 the successful piike-and-slab prior [Mitchell and Beauchamp/ 1988 where values are concentrated around O (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.
5

---

schedule
Number Assi gned vs Cluster Value
Number of Unique Values vs Percent
1.0
00
140
large sercentage of dusters and cons steeed nutllers which enirapy
120
live
0.8
Number Assiones
100
0.7-
Number of Unique
80
0.5
60
40
00312
20
00000
0.021
0.2
0
2
-0.5
0.0
0.5
Cluster Value
Round
Figure 4: Pr follows the same schedule (Subia- Waud and Dasmahapatra 2022 (left). In the middle and right plots, we see that PWEN achieves very small entropy values by majority of weights to only very small (4 or 5) cluster values and the rest assigned as outliers, most of which are powers-of-two.
We attempt this by making our o values to be large. Networks with large o have, probabilistically; more noise added to the μe values during training and 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 values. The motivation is that, much like L1 norms enforcing sparsity, this formulation will train the network to produce a large o; for noise-resilient parameter w,, whilst maintaining noise-sensitiv weight w, to have a small o, despite the prior pull. The regularised loss function for training the training phases the algorithm is:
- log P(D)u, a) + aLREG(o),
(2)
where the regularization term is:
N Σ [(a,))=- o. S) - oz):
CREG
with ((z) = for x 2 0 and 0 otherwise. The θ function prevents the optimization from finding network with a subset of a with infinitely large values and dominating the cross entropy term. thus cutoff on how large the values in can be. a is 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 F we illustrate the effect on the distribution of a under different a values for a ResNet-18 trained on the Cifar-10 10 dataset. As we increase a the or values no longer collapse to zero giving us information for downstream clustering.
Initialization using Relative Distance from Powers-of-two. For each weight w; we need to specify its prior distribution so as to derive the posterior using Bayesian updating. We assume that the posterior distribution Gaussian with a diagonal covariance matrix: P(wa (4, whose parameters 14g. a; are trained using BBP. To initialize the prior distributions for μ4 and we set (n) = II: (P4)) where (14) xx Spy,e, for the pre-trained weight value wis. For a Gaussian posterior we would typically require
unknown a to be drawn from a Gamma conjugate prior distribution. Instead, we set d; to be a known function of the μı; at initialization. In [Subia- Waud and Dasmahapatra [022] 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 < 14 < 2= for integer Ty, and largest at the midpoint of the interval. We introduce a parabolic function (p:) product of relative distances of each pre-trained weight value (μ₂) to the nearest lower and upper powers of two:
1 Pal
- 122-11
(14))) (0.05)²
(4)
22
(Full details and visualization is given in the appendix).
PWFN Clustering. In Figure I we show 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 = 1...., T iterations. At each step we set fraction P: of the weights to be fixed, so that |W ixxeal Np.. The remaining weights at iteration stage / are trainable and called W free. We follow the scheme first proposed in [Subia- Waud and Dasmahapatra 2022] in setting Pt. (Figure left). All of the weights w; that are assigned to W lixed will have their μ4 values fixed to one of the set of cluster centers. At the last iteration, |Wire! and Pr = 1, as all weights have been fixed to their allocated cluster centroids.
We next introduce how a cluster center Cj is defined and how the mapping μ4: H C E is performed.
Let R = (- nr . 2/+T: 27. 0, T7 2 2 } be the set of all powers-of-two up to precision b. For weight to be desired additive power of two, a sum over most w elements of R is defined
to be a cluster center of order w. Formally, for P(R) the power set of R.
c² = i
(5)
I
PWEN begins with order w = 1, the powers-of-fe of-tso up to precision b as the proposal cluster set c², Next, for each weight w;; = (1z 7,) in the network, the value of o; used to determine how far away they are from each of the cluster centers using:
114 -gyl
Droub (wi, (3)
(6)
Oi
Interpret this Mahalanobis distance as: "how many sigmas (standard deviations) away is cluster Oj E e from weight w.". At iteration stage I, for each free weight we define d:' (i) minege Dro c)
as the cluster center that is the fewest sigmas away from wi E Wfe We denote by ng the number of weights with the smallest proh cluster 4 i.e., n = En Mick = =<(())] We then take the index k" of the cluster with the most number of weights nearest to a cluster: h = argmax, ng. Thus,
c- E c
(cf .c......C
o ) the cluster with the most number of weights nearest it.
We then order the weights in W by their distance to c... In detail, for Wi [W1, .u.,.... Wi;....................... ...,....,.................. W,). we reorder the weights by permuting the indices wn = Wa(s) where [1 II ,l....,77 n?] is a permutation, i + T r()). The ordered list [w], I..,..., w²,] satisfies
Dprob (u's c S Dpoob(u(4+1 c+
Next, we need to determine how many of these weights we should assign to cluster of. To do so, we define a threshold 8 and we take the first (()) weights from [uf...., w/)1 such that:
8(8)
1
Dprot ( S 8.
(8)
{(5)
As long as this is possible with E())= > O, we have identified both a cluster ef. and set of weights [w],..., ('(()) which can be moved from W to W We map the weights in [w],..... w/(0)] "'(()) =
(1) (M'RO) [(()))] to single weight wa- = (x-)) corresponding to cluster o Uk. c- and σe. sta([uá,.... (((o)) where std computes the standard deviation of its argument. This process is then repeated, finding the next most popular cluster until Np, weights are assigned a cluster. If ((d) O, before enough weights are assigned in iteration 4, then we have not been able to find any cluster centers Cj E d which are close enough to any weight, iee. Dpras (w:, <2) > 8 for all weights w; E Wt free and ej Ck- In this case, we set w +- w + 1 and + 25 giving us a higher-order additive powers-of-two of-two set and less restrictive ó value threshold. Since |c |e this increase in w makes more cluster centers available during the next clustering attempt.
Putting it All Together. Putting the training and clustering stages together, we have process for training a neural network whose weights are from Gaussian posterior distribution with diagonal covariance matrix by backpropagation (BPP) that favours configurations with long Gaussian tails, which the clustering stage can then use to identify which weights to move and to what extent. This process is repeated for 7' iterations, with the fraction p. of weights increasing with each 1 Pt+1 > pt until all of the weights are moved from W free to W Fixed at iteration T where pr 1. We give the full algorithm in the Appendix.
7

---

Table 1: Full comparison results. (w/o FL-Bias) refers to calculating the metrics without the first-last
layers and bias terms included. 'Params' refers to the unique parameter count in the quantized model, entropy is the full weight-space entropy. In-ch, layer, attn refer to whether the method uses separate codebook for each layer, in-channel and attention head respectively.
Separate Codebook
Model
Method
Layer
In-ch
Attn
Top-1 (Ensemble)
Entropy 23.3
Params
ResNet-18
Baseline
68.9
10756029
LSQ
/
APoT
X
69.9
5.7
1430
WFN
X
69.7
3.0
164
PWEN (no prior)
X
X
69.3 (69.6)
1.7
143
PWFN
X
X
70.0 (70.1)
2.5
155
ResNct-34
Baseline
73.3
24.1
19014310
LSQ
X
71.9
APoT
X
73,4
6.8
16748
X
WFN
73.0
3.8
233
PWEN (no prior)
(74.4)
1.2
147
PWEN
X
74.3 (74.6)
1.8
154
ResNet-50
Baseline
76.1
24.2
19915744
LSQ
X
75.8
WFN
X
76.0
4.1
261
X
PWEN (no prior)
X
77.2 (78.1)
3.5
334
PWFN
X
X
77.5 (78.3)
3.4
325
DeiT-Small
Baseline
79.9
16.7
19174713
LSQ+
77.8
Q-ViT
78.1
11.3
3066917
Q-VIT (wfo FL-Bias)
10.4
257149
PWEN (no prior)
78.0 (78.3)
2.7
352
PWFN
78.1 (78.5)
2.7
356
DeiT-Tiny
Baseline
72.9
15.5
5481081
LSQ+
68.1
Q-VTT
69.6
11.5
1117630
Q-VIT (w/o FL-Bias)
69.6
10.5
128793
PWEN (no prior)
71.4 (71.)
2.8
300
PWFN
71.2 (71.5)
2.8
296
DenseNet16l
Baseline
77.8
17.1
26423159
PWEN
X
X
X
77.6 (78.0)
1.1
125
4 Experiments
We conduct our experimentation on ImageNet dataset with wide range of models: ResNets- (18,34.50) [He et al. 2016), DenseNet-161 (Huang et al. [017] and the challenging DeiT' (small and tiny) [Touvron et 2021]. For each model, we convert all the parameters in the convolution and linear layers into Gaussian distributions where the mean value is set to be the weight value of the pre-trained model found in the Timm library. Thus, at test time with no further training, we retain the original accuracies. We set the variance parameters according to the setting described in Eq (12). We then apply nine rounds of the described weight fixing with three epochs of re- re-training each round, totalling to 27 epochs of training. We train using SGD with momentum 0.9 with a learning rate of 0.001. For all experiments, we fix ō = 1, a 2-11 which we found using grid-search on the Cifar- 10 dataset and works surprisingly well in all settings. For all our experiments train using 4x RT'X8000 GPUs and batch-size of 128. For the ensemble results, we sample 20 times different weights using the learned weights' distributions and report the mean accuracy.
5 Results
compare PWFN against range quantization approaches where the model weights have been made available so that make accurate measurements of entropy and unique parameter count.

---

Table 2: Comparison of the number of additional training epochs required by different fine-tuning
quantization methods.
Method
Num of additional epochs
ApoT
120
PWFN
27
WFN LSQ
27
90
QviT
300
For the ResNet family, we compare against the current sfate-of-the-art APoT [Li et al. 2019b and WFN [Subia- Waud and Dasmahapatra 2022 For the transformer models, there has only been one work released, Q- (Li et al. 2022 which has both the model saves and code released. For both APoT and Q-Vit, we compare the 3-bit models which are the closest in terms of weight-space entropy to that achieved by PWFN.
As presented in Table 2 PWFN requires substantially fewer additional training epochs than most methods, save for WFN. highlighting its training efficiency. We use straightforward regularization term that encourages an increase in a, and its computational cost is comparable to that of 11 regular- ization. While our approach does lead to greater memory demands due to the additional a parameters and their associated gradient updates, the overall simplicity of the method is more efficient than previous BNN training procedures making it feasible to tackle more complex model-dataset pairings. Additionally, we note that when using the quantized version for inference, there are no extra costs, and the BNN functions as a point-estimate network.
In Table T we can see the set of results. PWFN demonstrates superior entropy, unique parameter count and top-1 accuracy across the board. In addition the point-estimate accuracy using the mean of each of the weights' distributions (the cluster centers), we can additionally sample the weights from the learned distributions give us an ensemble of models the mean prediction of which gives us further accuracy gains which we show in brackets in the Table. The prior initialization gives slight but consistent accuracy improvement over using a uniform prior (PWFN (no prior)). We note that for both APoT and Q- Vit, different codebooks are used for different layers and for Q-Vit, different codebooks were additionally used for every attention head and input channel, and the bias terms were left unquantized, pushing up the parameter count and weight-space entropy substantially. We highlight this as a growing trend in the field, where relaxations such as leaving large parts of the network unquantized, or using different codebooks for ever granular parts of the network, are often used. Each relaxation comes at a cost in hardware, be that support for unquantized elements - such as the first and last layers or the use of different codebooks for various parts of the architecture. Figure 2 illustrates the variation in entropy and the count of unique parameters across different layers and attention components. A notable observation from our study is that the weights associated with the "value' component exhibit higher entropy in the final layer. This observation aligns with the notion that employing fixed quantization scheme for each layer necessitates a relaxation of the quantization constraints specifically for the last layer, as supported by prior studies (Li et al., 2019a| Jung et al.) 2019. Zhou et al.) 2016, Yamamoto, 2021 Oh et al., 2021 Li al. 2022] Moreover, this highlights an intriguing possibility that in the context of attention networks, such relaxation might be essential only for the 'value" weights, and not for the 'keys' and 'queries'.
In understanding how PWFN is able to compress network's ss representation to such a degree compared to WFN we look to how often the previously proposed relative distance threshold is maintained.
In Figure 5 it's evident that while the relative distance threshold established in WFN is, on average, maintained, there are edge-cases where it isn't. 't. This observation suggests that having context- specific noise tolerance benefits subsequent compression stages. Furthermore, the data indicates that these values are typically small (as seen in the left column), have a high frequency of occurrence (depicted in the middle), and are predominantly assigned during the middle (0.6, 0.7) and final rounds.
httpgs//gitchbb.ccca//hhh/AAA/ /atpa///aaaaaaa/aa///oaaaaanaaar Quantization
httpee//ggthhbb.ccmmmabbaaarrrrr Fix Fieterrerrtt..eeeeeerreeeee Networks

---

Max Roladra Oistance Assigned Value
Max Reladive Distance Assignod Number
Mox Relative Disiance Addigned Porcentage
0.7-
0.6
0.4
0.3
0.1-
0.0
Number
Acconetgge Cludwred
Mean Relative Distance Ansigned to Viluee
Mean Relstive Distance Ansigned to Number
Mean Relative Distance Assigned to Percentage
FAMAbN Ditance
0.3-
02-
Palabbe
0.1
0.0
10
1
107
10a
10
10a
10'
0.95
0.975
0.99
Number Auiggedd
Aeceetage Custere
Figure 5: The maximum (top left) and mean (bottom left) relative distance weight moves cluster by cluster value. The maximum relative distance is not well maintained with the number of weights assigned to that cluster (top middle), but the mean relative distance is (bottom middle). The maximum (top) relative distance for each cluster assignment and mean (bottom) relative distance by round are shown in the right-hand column. In all plots, we show in blue the threshold used in WFN.
Conclusion
This work proposed PWFN, training and clustering methodology that can both scale BNNs to complex model-dataset combinations and then use the weight uncertainties to inform quantization The result a set of networks with extremely low weight-space entropies that can be used downstream in accelerator designs to mitigate expensive data movement costs. Additionally, we have seen the potential of the probabilistic aspect of the learned networks with sampled ensemble giving noticeable accuracy gains. An exciting direction for future work is to explore how the uncertainty estimations and the duittf-disttibuttin performance of neural networks could be enhanced using PWFN to train Bayesian Neural Networks.
7 Acknowledgments
This work was supported by the UK Research and Innovation Centre for Doctoral Training in Machine Intelligence for Nano-electronic Devices and Systems [EP/S024298/1]
References
Jan Achterhold, Jan Mathias Koehler, Anke Schmeink, and Tim Genewein. Variational network quantization. In International Conference on Learning Representations, 2018.
Melih Barsbey, Milad Sefidgaran, Murat A Erdogdu, Gael Richard, and Umut Simsekli. Heavy tails in sgd and compressibility of overparametrized neural networks. Advances in Neural Information Processing Systems, 34:29364-2993 2021.
Charles Blundell, Julien Cornebise, Koray Kavukeuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International conference machine learning, pages 1613-1622, PMLR, 2015.
10

---

Angela Fan, Pierre Stock, Benjamin Graham, Edouard Grave, Rémi Gribonval, Herve Jegou, and Armand Joulin. Training with quantization noise for extreme model compression. arXiv preprint arXiv:2004.07320, 2020.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. arXiv preprint urXiv:2010.01412, 2020.
Vincent Fortuin. Priors in bayesian deep learning: A review. International Statistical Review, 90(3): 563-591, 2022.
Yarin Gal and Zoubin Ghahramani. Dropout bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pages 1050-1059.
PMLR, 2016.
Mert Gurbuzbalaban, Umut Simsekli, and Lingjiong Zhu. The heavy-tail phenomenon in sgd. In International Conference on Machine Learning, pages 3964-3975. PMLR, 2021.
Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint 1rXiv:1510.00149 2015.
Song Han, Xingyu Liu, Huizi Mao, Jing Pu, Ardavan Pedram, Mark A Horowitz, and William J Dally. Eie: Efficient inference engine on compressed deep neural network. ACM SIGARCH Computer Architecture News, 44(3)224--24, 2016.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 770-778, 2016.
José Miguel Mernández-Lobatt and Ryan Adams. Probabilistic backpropagation for scalable learning of bayesian neural networks. In International conference on machine learning, pages 1861-1869. PMLR, 2015.
Liam Hodgkinson and Michael Mahoney. Multiplicative noise and heavy tails in stochastic optimiza- tion. In International Conference on Machine Learning, pages 4262-4274. PMLR, 2021.
Mark Horowitz. 1.1 computing's s energy problem (and what we can do about it). In 2014 IEEE International Solid-State Circuits Conference Digest of Technical Papers (ISSCC), pages 10-14. IEEE, 2014.
Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Weinberger. Densely connected convolutional networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4700-4708, 2017.
Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient integer-arithhetti-onlly inference. In Proceedings of the IEEE conference on computer vision and pattem recognition, pages 2704-2713, 2018.
Laurent Valentin Jospin, Hamid Laga, Farid Boussaid, Wray Buntine, and Mohammed Bennamoun. Hands-on bayesian neural networks- a tutorial for deep learning users. IEEE Computational Intelligence Magazine, 17(2):29 48, 2022.
Norman P Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, et al. In-datacenter performance analysis of tensor processing unit. In Proceedings of the 44th annual international symposium on computer architecture, pages 1-12, 2017.
Sangil Jung, Changyong Son, Seohyung Lee, Jinwoo Son, Jae-Joon Han, Youngjun Kwak, Sung Ju Hwang, and Changkyu Choi. Learning to quantize deep networks by optimizing quantization intervals with task loss. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 4350 4359, 2019.
Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt Kusner. When do flat minima optimizers work? Advances in Neural Information Processing Systems, 35:16577-16595, 2022.
11

---

Diederik P Kingma and Max Welling. arXiv:1312.6114, 2013.
Auto-encoding variational bayes.
arXiv preprint
Edward H Lee, Daisuke Miyashita, Elaina Chai, Boris Murmann, and S Simon Wong. Lognet: Energy-efficient neural networks using logarithmic computation. In 2017 IEEE International IEEE, 2017.
Conference on Acoustics, Speech and Signal Processing (ICASSP), pages
Yuhang Li, Xin Dong. and Wei Wang. Additive powers-of-two quantization: An efficient non-uniform discretization for neural networks. In International Conference on Learning Representations, 2019a.
Yuhang Li, Xin Dong, and Wei Wang. Additive powers-of-two quantizztinnn An efficient non-uniform discretization for neural networks. arXiv preprini arXiv: 1rxiv:1909.1314444 2019b.
Zhexin Li, Tong Yang, Peisong Wang, and Jian Cheng. Q-vit: Fully differentiable quantization for vision transformer. arXiv preprint (eXiv:2201.077000 2022.
David JC MacKay. A practical bayesian framework for backpropagation networks. Neural compuia- tion, 433):4484472 1992.
Toby Mitchell and John Beauchamp. Bayesian variable selection in linear regression. Journal of the american statistical association, 83(404):102311032 1988.
Sangyun Oh, Hyeonuk Sim, Sugil Lee, and Jongeun Lee. Automated log-scale quantization for low-cost deep neural networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 742-751, 2021.
Dominika
Przewlocka- Rus and Tomasz Kryjak. Energy efficient hardware acceleration of neural networks with power-of-two quantisation. In Computer Vision and Graphics: Proceedings of the International Conference on Computer Vision and Graphics ICCVG 2022, pages 225-27 236. Springer, 2023.
Pierre Stock, Armand Joulin, Rémi Gribonval, Benjamin Graham, and Hervé Jégou. And the bit goes down: Revisiting the quantization of neural networks. In ICLR 2020-Eight International Conference on Learning Representations, pages 1-11, 2020.
Christopher Subia-War and Srinandan Dasmahapatra. Weight fixing networks, In Computer Vision- ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27 2022, Proceedings, Part XI, pages 415-431. Springer, 2022.
Vivienne Sze, Yu-Hsin Chen, Tien-Ju Yang, and Joel S Emer. Efficient processing of deep neural networks: A tutorial and survey. Proceedings of the IEEE, 105(12):225--2329, 2017.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers distillation through attention. In International conference on machine learning, pages 10347-10357. PMLR, 2021.
Karen Ullrich, Edward Meeds, and Max Welling. Soft weight-sharing for neural network compression. arXiv preprint arXiv: 07077000000000 2017.
Sebastian Vogel, Rajatha B Raghunath, Andre Guntoro, Kristof Van Laerhoven, and Gerd Ascheid. Bit-shift- based accelerator for cnns with selectable accuracy and throughput. 2019 22nd Euromicro Conference on Digital System Design (DSD), pages 663 667. IEEE, 2019.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pages 681-688.
2011.
Junru Wu, Yue Wang, Zhenyu Wu, Zhangyang Wang, Ashok Veeraraghavan, and Yingyan Lin. Deep k-means: Re-training and parameter sharing with harder cluster assignments for compressing deep convolutions. 35th Int. Conf. Mach. Learn. ICML 2018, 12:8523 2018a.
Junru Wa, Yue Wang, Zhenyu Wa, Zhangyang Wang, Ashok Veeraraghavan, and Yingyan Lin. Deep k-means: Re-training and parameter sharing with harder cluster assignments for compressing deep convolutions. In International Conference on Machine Learning, pages 5363-5372. PMLR, 2018b.
12

---

Kohei Yamamoto. Learnable companding quantization for accurate low-bit neural networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 5029-5038 2021.
Aojun Zhou, Anbang Yao, Yiwen Guo, Lin Xu, and Yurong Chen. Incremental network quantization: Towards lossless cnns with low-precision weights. arXiv preprint wwxXiv//02.030000 2017.
Shuchang Zhou, Yuxin Wu, Zekun Ni, Xinyu Zhou, He Wen, and Yuheng Zou. Dorefa-net: Train- ing low bitwidth convolutional neural networks with low bitwidth gradients. arXiv preprint 1606.06006.611000 2016.
8 Appendix
Bayes-by-backprop and PWFN The Bayesian posterior neural network distribution P(w) D) is approximated by distribution ((wr(0) whose parameters 0 are trained using back-propagation,, Bayes-by-backproy (BPP). The approximation
achieved by minimizing the Kullback-Leible (KI.) divergence D [] between /' and Q to find the optimal parameters 0*, These parameters 0* instantiate the means μt and variances of of the PWFN.
0r arg min K K[[Q(ur()) ((w) D)], where
8
P(w] D) EQ(w o) log (Q(w)))P(()
Q(wl0)
EQ( log
P(D[w))(er
The P(D) term does not contribute to the optimization and is dropped, leaving
H*
arg min (q(wlo) [log Q(w - log P(D loog P(w)],
0
arg min Σ log Q(w")(() log P(
log P(D|w" w
(10)
pricr dependent
data dependent
where the expectation value approximated by samples uhth (uml(() drawn from Q(w)(0) each of which instantiates neural network.
Gradient descent over each whth that instantiates a neural network is made possible by the re- parametrization trick Kingma and Welling 2013 Blundell et al., 2015). The idea is to regard each sample utd - or where cm p(c) random draw from some distribution that we take be an isotropic Gaussian: p(c) = (0, I) with the N V-dimensional identity matrix for the N weights of network W. These weights w"' th are used in the forward pass through the network while parameters and a are trainable. Then, for any function f'w)) we have EQ( [f(w]] (Pp(e)[((u so that
i ag Bo 0)] de Ip 0)] Ep(c)l
0f(w, 0) Ow @f(w, 0)
dw
00
00
The terms are all calculable, allowing to draw from distribution for each weight "h and backpropagate to the underlying distribution parameters == (14, or.. For wmth μ + omo, the derivatives are Oro" m dp I, and Bhr
80 em making the respective gradients
0f(w, O.) Of(w, o,o) and V. Öf(w. H. Ow
0f(w M.
cm
(11)
dw
0p
do
PWEN Clustering Algorithm
Prior Initialization
In addition to the prior initialization described main paper. we added an reweighting determined by the size of the values in the network. Using the definition of u (iei(u, pow2,, (u.)) re-weight by the third quartile 00.75 and re-write the initialization as:
Dra (14, ((w),(1))) x Dree( (142
f(4) = 0.0025 x
(12)
80.75
13

---

while W
fixed Σ Npr do
fixed new +- []
while fixednew empty do
Increase the order w w +1
Increase 0 + 26
e" + 1-ieer i P(R) |rl s
for each -...,,.,111 |
(i) t min Dprob(ws, c)
for each cluster center og E C
n + +i Hic = =c
k* f- arg maxe nE
Sort: [w/...., W'N + [WI,.... w.]..,,,,,,,,.............. w/ wa(e) permutation
where D Dprob (w', (ha) < Dprob (1 (1/+11,
i + 1, mean & prob (w1, cka)
while mean < do
fixednew t w!
mean + 74I mean + it1 x Pprob(w((34
i - i+1
end
8:=t t 3 x 81=t
end
Assign all the weights means H4, fixednew to cluster center c (1) and set each of the o, E fixednew to the variance of the weight means in fixed Finally, move them from W to W
end
Algorithm 1: Clustering Npt weights at the th iteration in PWFN.
and clamp the values be within the range [2 30 0.05] giving us our initial variance values.
o, m0x (finn(())))))) 2 2-30.
(13)
14

---

mOOA Wor Aicand
mocetleer Bxand
model wyd Leand
nocelteeri coma
model Ryeer Leavv
rodee ayerr Laarnc
mocktionre Brawd
mockiiae Buaedd
modellavert Oeznn
model prodelllyyert Oiaand
model
model leyert dovnanmple,
modee Luand
roodel syyo
and Duard
Figure 6: Here we compare the μ values for all weights given layer initialization (left) and after PWFN convergence and clustering (right).
15


Reviews:
[{"review_id": "vFSEUJYi7E", "reviewer": "Reviewer_TzYT", "summary": "This paper introduces a novel approach to weight-sharing quantization, a technique aimed at reducing energy costs associated with inference in deep neural networks. The authors propose a method that takes into account the context of weights in the network, arguing that this strategy can better preserve the network's representational capacity while reducing its complexity. They employ a probabilistic framework to capture the flexibility in weight values, which guides clustering decisions to reduce the network's entropy and lower the unique parameter count without compromising performance. The authors also suggest a novel initialization setting and a regularization term to prevent the collapse of the weights' distribution variance to zero. The paper demonstrates superior compressibility and accuracy in ResNet family models trained on the ImageNet dataset and transformer-based architectures.", "questions": "See Weaknesses", "limitations": "See Weaknesses", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "To the best of my knowledge, this paper is the first to utilize Bayesian Neural Networks for model quantization.\n\nThe authors conduct experiments on powers-of-two and additive powers-of-two quantization, which are more hardware-friendly than other types of quantization.", "weaknesses": "The paper's writing style is somewhat difficult to follow, particularly in the method presentation and the explanation of Table 1 (In-ch?).\n\nI recommend the authors use the term 'Quantization' instead of 'Quantisation' for consistency with most literature in the field.\n\nThe authors should address the training difficulty associated with Bayesian Neural Networks and discuss whether this approach is simpler for model quantization compared to other methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "First and foremost, we genuinely appreciate your constructive feedback on our paper. It is our aim to present our research in the clearest possible manner, and your insights are invaluable in this pursuit.\n\n**Bayesian Neural Networks (BNNs) for Model Quantization**: You're right in pointing out the uniqueness of utilizing BNNs for model quantization. We feel it is essential to highlight this, and we're pleased that you recognized its novelty.\n\n**Explanation of Table 1**: We agree that a more explicit explanation of Table 1 is needed and have amended this. in-ch refers to whether the approach uses a different codebook for each channel. So if we had a layer with weights with the dimensions: (in-ch, out-ch, filter-width, filter-height) then for each in-ch there would be a different set of unique weights used in the quantisation procedure. This both increases the complexity of inference (usually in the form of different scale/shift constants) and increases the memory requirements.\n\n   **Consistency in Terminology**: We had used the British spellings throughout the paper, but will swap this over to the American spelling to be consistent with the literature. \n\n**Training Costs of BNNs**: Given that multiple reviewers have questions around the training costs associated with Bayesian Neural Networks, we've incorporated a response addressing this aspect in the general rebuttal comments.\n\nOnce again, thank you for your review, we hope our revisions adequately address your concerns."}, {"review_id": "pDsWxvCzNt", "reviewer": "Reviewer_NcTB", "summary": "This paper presents a novel quantization scheme based on iterative training and clustering of the weights into a very limited and finite set of choices. The paper follows a Bayesian approach to assign weights to clusters. Experiments are demonstrated on ImageNet for ResNet and Transformer architectures. ", "questions": "See weaknesses reported above.", "limitations": "Limitation in the context of hardware has been stated by the authors. However, it would be nice to understand additional limitations w.r.t. training time compared to other quantization methods, and other features that play role in the method.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper addresses an important topic of improving quantization. The paper is very clearly written, easy to understand and the idea is simple and elegant. Experiments seem sufficient to demonstrate the power of the method.\n", "weaknesses": " I do not see any immediate weakness of the paper. One thing which could make the paper stronger is to explore the application on more complicated downstream tasks of for example DeiT or DeiT-small. This is just to understand the scalability of the method. Further, authors could also provide details on the added training costs to also challenges faced in the training involved for the approach. Additionally, detailed ablation is missing on how sensitive the results are w.r.t. the chosen prior, number of training+clustering combination steps, etc.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and valuable insights on our paper.\n\n**Training Costs**: We appreciate the concern raised about the training costs. Given its recurrent mention in the reviews, we've expanded on this aspect, providing a comparative analysis between PWFN and other prevalent quantization approaches as general comment. We hope this will shed light on the efficiency of our method in terms of training overheads.\n\n**Number of Training+Clustering Steps**: Your observation on the potential optimization of the number of training+clustering steps is apt. While we initially set the number of weight-fixing steps based on the success of WFN, we recognize the scope for fine-tuning. Post your review, we experimented by reducing the steps (down to 7) for CIFAR-10, which yielded encouraging results. However, it remains to be seen if this can be generalized to larger datasets. We acknowledge this as a potential area for future investigation.\n\n**Expansion to Other Domains**: We concur with your suggestion on diversifying the application domains. Our focus on image classification was largely inspired by preceding works, but we did venture into ViT-type architectures. Notably, with Bayesian Networks, PWFN stands out as the first approach to train image-net scale models. As for quantization, PWFN is one of only a few methods to successfully quantize DeiT-family models.\n\n**Future Research on BNN Quantization**: Your idea of expanding the scope of BNN Quantization to a broader range of tasks is indeed promising. We hope our work paves the way for subsequent research endeavors in this direction.\n\nThank you again for your constructive feedback."}, {"review_id": "vwnBvlkL9C", "reviewer": "Reviewer_tpmz", "summary": "This paper discusses the Weight-sharing quantization technique to reduce DRAM read costs. To address this issue, an iterative weight fixing scheme was employed, which involved alternating between network training and weight clustering. This paper proposes a BNN-based framework that utilizes a new initialization setting and regularization term. With this method, millions of weights can be represented by hundreds of unique values.", "questions": "(49) There is no explicit clarification for the abbreviation of PWFN.\nThere is no explanation for Figure 2 and 3 in the manuscript.\n(220) what is the meaning of j index?", "limitations": "The paper does not address the limitation of the method and any the negative societal impact.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "Probabilistic approach for weight-sharing quantization is novel and interesting.\nThe method improves the accuracy while minimizing the entropy and reducing the number of unique values used to represent the weights.\nThe prior initialization method using relative distance from powers-of-two consistently improves the accuracy, especially on resnet architectures.", "weaknesses": "There is no explanation in the manuscript about how the weight position described in the abstract was handled.\nLack of experimental support for regularization scheme (x)\nLack of clarification – No explanation about Figure 2 and 3 in the manuscript", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate the feedback and constructive criticisms provided by the reviewer. Your insights are invaluable to improving our manuscript. Here, we address the concerns raised:\n\n**Concern 1: Weight Position Handling**\n\nThe weight position is indeed utilised in our methodology when determining the cluster assignments. As you rightly pointed out, each weight is associated with an individual sigma value that aids in determining its cluster assignment. This unique approach ensures that weights, even with the same value but in different positions, can be assigned to distinct clusters. We realized the need for more clarity in our manuscript regarding this aspect and have amended the methodology section to elucidate this further.\n\n**Concern 2: Lack of Experimental Support and Clarification**\n\nThank you for pointing out the ambiguities associated with Figures 2 and 3. We've now expanded the captions for both figures to provide clarity. Specifically, Figure 2 showcases how increased alpha value (the scaling of the regularization term) plays a pivotal role in preventing sigma values from collapsing to zero. Avoiding this collapse is important, both in terms of using the network as a Bayesian Neural Network (if sigma is zero, then we have a point-estimate) and in our downstream use of using sigma to determine the quantization mapping. We hope this clarifies that experiments have been conducted and presented to explore the value of the regularization term in the methodology. Figure 3 has also been revised to enhance readability and comprehension.\n\n**Concern 3: Clarification on Abbreviations**\n\nOur sincere apologies for the oversight. We have now clearly defined the abbreviation as \"Probabilistic Weight Fixing Networks (PWFN)\" in the relevant sections of the manuscript.\n\n**Conclusion**\n\nWe believe our work on the bayesian weight-sharing quantization technique is a valuable contribution to the field, and we're confident that the changes we've made, following your feedback, will further clarify our methodology and its significance. Once again, thank you for your constructive feedback, and we hope our revisions address your concerns adequately."}, {"review_id": "D1hpAlDl27", "reviewer": "Reviewer_1QKt", "summary": "PWFN (Probabilistic Weight-Shared Fusion Network) is a technique that combines weight-sharing quantization with Bayesian neural networks to achieve highly compressed and quantized neural networks. It models each weight as a draw from a distribution, allowing for the quantification of uncertainty in weight values. PWFN consists of iterative stages of training and clustering to optimize the weight distributions and determine optimal cluster configurations.\n\nPWFN offers a probabilistic approach to weight-sharing quantization, providing enhanced compressibility, flexibility, and noise resilience in neural networks. It allows for uncertainty quantification and demonstrates improved performance compared to existing methods. ", "questions": "See Weaknesses", "limitations": "N/A", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "(1)  By representing weights as probability distributions, PWFN provides flexibility in weight values and enhances noise resilience. The distributions capture uncertainty in weight values, enabling robust performance even in the presence of noise.\n\n(2) PWFN incorporates a regularization term and iterative clustering procedure, leading to improved performance compared to other methods\n\n(3) By sampling weights from their distributions and observing changes in predictions, the model's confidence and uncertainty can be assessed.", "weaknesses": "(1) PWFN needs additional complexity in the training process including regularization terms and clustering iterations.\n\n(2) PWFN may require modifications to existing neural network frameworks to incorporate the probabilistic weight-sharing framework. It may also require additional expertise in Bayesian neural networks and variational inference techniques. For the reproducibility, the code needs to be open.\n\n(3) Modeling weights as distributions and performing sampling during training can introduce additional computational overhead compared to traditional point estimate-based methods. This may result in increased training time and inference latency.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful comments and constructive feedback on our submission. We appreciate the time and effort you've taken to review our paper.\n\n**Complexity in Training Process**: We understand your concern about the additional complexity introduced by PWFN during the training process and have added a general comment with regards to this which we hope alleviates your concerns. \n\n**Framework Modifications and Expertise**: We're pleased to share that we'll be releasing the code, and the process of mapping a traditionally trained network to a BNN will be automated - a user need only pass to the code a pre-trained model. This will significantly reduce the barriers to adoption, allowing those unfamiliar with Bayesian neural networks or variational inference techniques to still benefit from PWFN.\n\n**Computational Overhead**: You rightly pointed out the potential computational overhead when modeling weights as distributions. However, it's crucial to note that this overhead is mainly seen during training and isn't any larger than existing re-training quantization approaches. \nFor standard inference tasks, our method uses quantized point estimates for weights, ensuring no added latency. \nWhen leveraging the BNNs for uncertainty estimation there are additional costs, but these are comparable to running inference multiple times on the same network (aside from the random sampling step), and will likely have much less overhead when compared to ensembling separately trained models.   \n\nOnce again, thank you for your review, we hope our revisions address your concerns adequately."}, {"review_id": "sfG2n4MGLf", "reviewer": "Reviewer_jtUH", "summary": "This paper discusses a method of quantizing weights based on Bayesian neural networks (BNN) by clustering them as closely as possible to a set of powers-of-two values. Unlike the weight fix network (WFN) that quantizes by clustering the fixed weights of the network around the nearest centroid, the proposed Probabilistic weight fix network (PWFN) makes use of a probabilistic BNN model. Weights are assumed to follow a gaussian distribution. The $\\mu$'s and $\\sigma$'s of the weights are learned using Bayes-by-backprop (BBP). Quantization is performed iteratively by clustering around a specific centroid using the learned $\\mu$ and  $\\sigma$ from the model’s codebook.", "questions": "- Q1) How was 0.05 determined in equation 4? Is it a value that can be generally applied to all networks, or is it a hyperparameter that should be applied differently for each model?\n\n- Q2) Can you provide the standard deviation of accuracy in ensemble models? And can you provide the maximum and minimum accuracy of training?\n\n- Q3) How is the Top-1 accuracy in Table 1 determined, since the results may vary for each training?\n\n- Q4) It seems like using BNN would greatly increase training time; how much? Can you tell us how much the training time increases for ResNet-18 and DeiT compared to WFN?\n\n- Q5) How much time/MACs are required for training and inference?", "limitations": "Limitations are not discussed.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "By expressing the existing network as a weight-shared BNN, several advantages can be obtained:\n- It allows for uncertainty quantification.\n- Using the Mahalanobis distance to determine the centroid seems more reasonable than using the Euclidean distance or Relative movement distance.\n\nIn many existing papers, codebooks were maintained at the layer or channel level, requiring significant storage capacity. However, this paper achieves higher quantization efficiency by utilizing a codebook at the model level.", "weaknesses": "- The amount of computation is significantly increased compared to the conventional WFN. It is also harder to train than a regular network, which also increases the time for hyperparameter exploration.\n\n- Because of its complexity (BNN), it is unclear whether it can be applied to large models.\n\n- The choice of prior and the learning of mu and sigma can greatly influence the results. This can lead to inconsistent performance with each training and cases where performance drops significantly. Using ensembles seems like it would require too much inference time.\n\n- Sigma is set as the distance from the actual parameter to the power of two, but this could be different from the actual sensitivity of the parameter and could be a wrong choice of prior.\n\n- Reference format should be easily distingusiable from that of equation numbering. Since both uses parenthesis (), it looks quite odd at first sight. \n\n- According to the NeurIPS 2023 guide, all table numbers and titles should always appear before the table. https://media.neurips.cc/Conferences/NeurIPS2023/Styles/neurips_2023.pdf\n\n- Some symbols in the equations are insufficiently explained regarding their dimension, type, or meaning. Line 192 Θ, Line 193 S, Line 201 δ. Also, δ is duplicated with the threshold δ on Line 237.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the comprehensive feedback provided on our paper. Herein, we address your concerns systematically:\n\n**Q1) The Determination of 0.05 in Equation 4**\n\nThe value 0.05 was empirically determined through experiments on Cifar-10 using ResNet18. While this parameter demonstrated consistent performance across our tested networks, further exploration may optimize it for different architectures.\n\n**Q2)** Can you provide the standard deviation of accuracy in ensemble models? And can you provide the maximum and minimum accuracy of training?\n\nThank you for highlighting the importance of this metric. Here are the detailed results:\n\n| Model       | Max   | Min   | Std  |\n|-------------|-------|-------|------|\n| ResNet18    | 70.23 | 69.99 | 0.08 |\n| ResNet34    | 74.99 | 74.34 | 0.16 |\n| ResNet50    | 79.08 | 77.33 | 0.42 |\n| DeiT-Small  | 78.86 | 78.00 | 0.23 |\n| DeiT-Tiny   | 71.79 | 71.38 | 0.17 |\n\n\n**Q3)** How is the Top-1 accuracy in Table 1 determined, since the results may vary for each training?  \n\nFor the Top-1 accuracy, we fixed the quantized weights to their means (mu values). PWFN allows for both a compressed network for efficient inference and, optionally, an ensemble for uncertainty estimation.\n\n\n**Q4)** It seems like using a BNN would greatly increase training time; how much? Can you tell us how much the training time increases for ResNet-18 and DeiT compared to WFN?  \n\nIndeed, BNN adds an overhead, but it's manageable with current GPU resources. PWFN regularization also lessens the sigma-collapse challenge, making it feasible even for larger models. We've added a general rebuttal comment to further elaborate.\n\n\n**Q5)** How much time/MACs are required for training and inference?  \n\nFor inference in the quantized network, the time matches other quantized networks of similar compression rates. Uncertainty estimations introduce additional costs but only if deemed beneficial by the practitioner. These costs are comparable to running inference multiple times on the same network (aside from the random sampling step), and will likely have much less overhead when compared to ensembling separately trained models. \n\n**On the Complexity of BNNs**\n\nAs we show in the general comment, the overhead introduced training a BNN is manageable on consumer GPUs when using the PWFN regularisation term to avoid sigma-collapse. This means that PWFN can indeed be applied to large models. As far as we are aware, we are the first to successfully train a BNN on imagenet-scale problems, demonstrating strong performance on large and small models from the ResNet family as well as vision transformer models, and we see no reason why this can not be scaled further, we're excited to see how far training BNNs in this way can be pushed forward in future works.  \n\n**On Prior Choice and Consistency of Performance**\n\nWhile different priors might influence results, our choice is backed by previous works and targets optimal compression for hardware accelerators. The notion of inconsistency has been addressed in our response to Q2.\n\n**Clarifications on Symbols**\n\n    • Θ – this is the Heaviside step function. defined on Line 192.\n\n    • S – a constant as defined in Lines 194-196 and equation 3.\n\n    • δ – We've revised the text to clearly differentiate the context of its multiple usages.\n\n**Choice of Prior**\n\nWhile we acknowledge that other priors might optimize performance further, our choice is informed by previous works and the hardware efficiencies of mapping weights to powers-of-two.\n\n**Table Formatting**\n\nWe thank you for pointing out the format discrepancy. It's now corrected as per the NeurIPS 2023 guide.\n\n**Conclusion**\n\nWe value the feedback and believe the clarifications offered reinforce the strength and novelty of our methodology. We remain hopeful that our approach's advantages, as detailed in the paper and this rebuttal, and the utility of the work to the wider community can be seen.\n"}]

Author rebuttal:
### General Comments

Thank you to all the reviewers for their effort in reviewing our paper, kind comments, and suggestions for improvement. We have taken the comments and suggestions into consideration.

As a general confirmation, it is true that our work is the first to train Bayesian Neural Networks (BNNs) for the purpose of Quantization. To the best of our knowledge, we can make the even stronger claim that we are the first to train a BNN on ImageNet and achieve results that match (or exceed) standard training for the ResNet family of models, and exceed the SOTA accuracies for quantized transformer-based vision models.

Before addressing individual comments, we wanted to address a shared concern around the training costs incurred with training a BNN using the PWFN method. 

### Training Costs

It is indeed true that PWFN incurs additional training compared to the original baseline, but when we compare it to other post-training quantization methods:

| Method | Num of additional epochs |
|--------|--------------------------|
| ApoT   | 120                      |
| PWFN   | 27                       |
| WFN    | 27                       |
| LSQ    | 90                       |
| QviT   | 300                      |


We can see that PWFN requires substantially fewer additional training epochs than all methods but WFN.

WFN uses a regularizer that calculates the relative distance between all free weights and the existing cluster centers, and then penalizes weights depending on the distance to their closest center (in a soft way). This incurs computational costs in the backpropagation calculation for every iteration. In PWFN, we have a much simpler regularization term that penalizes sigma to increase - with costs that match that of l0 regularization. We do have memory overhead in terms of the number of parameters at training (sigma and mu) and the random number generation to sample, but this is only at training time. The simplicity of the regularization term also means that we experience a speed-up over the previous BNN training procedure outlined in the original bayes-by-backprop paper, making much more complex model-dataset pairings tractable.

This is not to say there are no costs; we find that a single training epoch with ResNet-18 on the ImageNet dataset takes ~1 hour 30mins on 4 consumer GPUs (GTX1080's) for PWFN, compared with 40 minutes for standard training and ~1 hour 20mins for WFN (but for WFN this increases as the number of clusters increases through training).

### Inference

At inference, the weights are fixed to mu values and treated in the usual way; there are no additional computation costs, and quantization speed-up on supported accelerators can be achieved. The core idea of PWFN is to reduce the number of unique weights and weight-space entropy so that the resultant networks can be used on accelerators (such as EIE) which use Huffman coding as a compression scheme. 

### Key Points 

Re-highlighting the core contributions of this work:

* This is the first work to train BNNs on imagenet-scale problems
* PWFN reduces the param count and entropies of transformer-based architectures with SOTA quantization performance 
* We achieve this with 27 epochs of additional training 
* Mapping to a BNN allows for both highly compressed networks and uncertainty estimation probing if required through sampling
* The point-estimates can be used in accelerator designs in conjunction with Huffman encoding to reduce computationally expensive DRAM reads
* We use a single codebook for the entire network whereas most previous works needed multiple codebooks for each layer/in-channel/attention head to maintain performance

Thank you again for taking the time to review our work and we hope to have covered all your points in the individual responses. 


Meta-review: Author rebuttal addressed most of the reviewer concerns and all reviewers are leaning accept after final discussion.
Decision: Accept (poster)