SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We next introduce how a cluster center  $c_{k}$  is defined and how the mapping  $\mu_{i}\mapsto c_{k}\in \pmb{c}$  is performed. Let  $R = \{-\frac{1}{2^b},\dots , - \frac{1}{2^{j + 1}}, - \frac{1}{2^j},0,\frac{1}{2^j},\frac{1}{2^{j + 1}},\dots \frac{1}{2^b}\}$  be the set of all powers-of-two up to a precision  $b$ . For a weight to be a desired additive power of two, a sum over at most  $\omega$  elements of  $R$  is defined to be a cluster center of order  $\omega$ . Formally, for  $\mathcal{P}(R)$  the power set of  $R$ ,
$$
\boldsymbol {c} ^ {\omega} = \left\{\sum_ {i \in r} i \mid r \in \mathcal {P} (R) \wedge | r | \leq \omega \right\}. \tag {5}
$$
PwFN begins with order  $\omega = 1$ , the powers-of-two up to precision  $b$  as the proposal cluster set  $c^{\omega}$ . Next, for each weight  $w_{i} = (\mu_{i},\sigma_{i})$  in the network, the value of  $\sigma_{i}$  is used to determine how far away they are from each of the cluster centers using:
$$
D _ {\text {p r o b}} \left(w _ {i}, c _ {j}\right) = \frac {\left| \mu_ {i} - c _ {j} \right|}{\sigma_ {i}}. \tag {6}
$$
Interpret this Mahalanobis distance as: "how many sigmas (standard deviations) away is cluster  $c_{j} \in \pmb{c}^{\omega}$  from weight  $w_{i}$ ". At iteration stage  $t$ , for each free weight we define  $c_{*}^{\omega}(i) = \min_{c \in c^{\omega}} D_{\mathrm{prob}}(w_{i}, c)$  as the cluster center that is the fewest sigmas away from  $w_{i} \in W_{\mathrm{free}}^{t}$ . We denote by  $n_{k}^{\omega}$  the number of weights with the smallest  $D_{\mathrm{prob}}$  to cluster  $c_{k}^{\omega}$ , i.e.,  $n_{k}^{\omega} = \sum_{i} \mathbb{I}[c_{k}^{\omega} = c_{*}^{\omega}(i)]$ . We then take the index  $k^{*}$  of the cluster with the most number of weights nearest to a cluster:  $k^{*} = \operatorname{argmax}_{k} n_{k}^{\omega}$ . Thus,  $c_{k^{*}}^{\omega} \in c^{\omega} = (c_{1}^{\omega}, \ldots, c_{k}^{\omega})$  is the cluster with the most number of weights nearest to it.
We then order the weights in  $W_{\mathrm{free}}^{t}$  by their distance to  $c_{k^*}^{\omega}$ . In detail, for  $W_{\mathrm{free}}^{t} = [w_{1}, \ldots, w_{i}, \ldots, w_{n}]$ , we reorder the weights by permuting the indices  $w_{i}' = w_{\pi(i)}$  where  $\pi: [1, \ldots, n] \to [1, \ldots, n]$  is a permutation,  $i \mapsto \pi(i)$ . The ordered list  $[w_{1}', \ldots, w_{n}']$  satisfies
$$
D _ {\text {p r o b}} \left(w _ {i} ^ {\prime}, c _ {k ^ {*}} ^ {\omega}\right) \leq D _ {\text {p r o b}} \left(w _ {i + 1} ^ {\prime}, c _ {k ^ {*}} ^ {\omega}\right) \tag {7}
$$
Next, we need to determine how many of these weights we should assign to cluster  $c_{k^*}^\omega$ . To do so, we define a threshold  $\delta$  and we take the first  $\ell(\delta)$  weights from  $[w_1', \dots, w_n']$  such that:
$$
\frac {1}{\ell (\delta)} \sum_ {i = 1} ^ {\ell (\delta)} D _ {\operatorname {p r o b}} \left(w _ {i} ^ {\prime}, c _ {k ^ {*}} ^ {\omega}\right) \leq \delta . \tag {8}
$$
As long as this is possible with  $\ell(\delta) > 0$ , we have identified both a cluster  $c_{k^*}^\omega$  and set of weights  $[w_1', \ldots, w_{\ell(\delta)}'$  which can be moved from  $W_{\mathrm{free}}^t$  to  $W_{\mathrm{fixed}}^{t+1}$ . We map the weights in  $[w_1', \ldots, w_{\ell(\delta)}'] = [(\mu_1', \sigma_1'), \ldots, (\mu_{\ell(\delta)}', \sigma_{\ell(\delta)}')]$  to a single weight  $w_{k^*} = (\mu_{k^*}, \sigma_{k^*})$  corresponding to cluster  $c_{k^*}^\omega: \mu_{k^*} = c_{k^*}^\omega$  and  $\sigma_{k^*} = \operatorname{std}([\mu_1', \ldots, \mu_{\ell(\delta)}'])$  where  $\operatorname{std}$  computes the standard deviation of its argument. This process is then repeated, finding the next most popular cluster until  $Np_t$  weights are assigned a cluster. If  $\ell(\delta) = 0$ , before enough weights are assigned in iteration  $t$ , then we have not been able to find any cluster centers  $c_j \in c^\omega$  which are close enough to any weight, i.e.,  $D_{\mathrm{prob}}(w_i, c_j) > \delta$  for all weights  $w_i \in W_{\mathrm{free}}^t$  and  $c_j = c_{k^*}$ . In this case, we set  $\omega \gets \omega + 1$  and  $\delta \gets 2\delta$  giving us a higher-order additive powers-of-two sets and less restrictive  $\delta$  value threshold. Since  $|c^{\omega+1}| > |c^\omega|$ , this increase in  $\omega$  makes more cluster centers available during the next clustering attempt.
Putting it All Together. Putting the training and clustering stages together, we have a process for training a neural network whose weights are from a 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  $T$  iterations, with the fraction  $p_t$  of weights increasing with each  $t$ $p_{t+1} > p_t$  until all of the weights are moved from  $W_{\text{free}}$  to  $W_{\text{fixed}}$  at iteration  $T$  where  $p_T = 1$ . We give the full algorithm in the Appendix.
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 a separate codebook for each layer, in-channel and attention head respectively.
Model	Method	Separate Codebook	Top-1 (Ensemble)	Entropy	Params
Layer	In-ch	Attn	Top-1
ResNet-18	Baseline	-	-	-	68.9	23.3	10756029
LSQ	✓	X	-	68.2	-	-
APoT	✓	X	-	69.9	5.7	1430
WFN	X	X	-	69.7	3.0	164
PWFN (no prior)	X	X	-	69.3 (69.6)	1.7	143
PWFN	X	X	-	70.0 (70.1)	2.5	155
ResNet-34	Baseline	-	-	-	73.3	24.1	19014310
LSQ	✓	X	-	71.9	-	-
APoT	✓	X	-	73.4	6.8	16748
WFN	X	X	-	73.0	3.8	233
PWFN (no prior)	X	X	-	73.5 (74.4)	1.2	147
PWFN	X	X	-	74.3 (74.6)	1.8	154
ResNet-50	Baseline	-	-	-	76.1	24.2	19915744
LSQ	✓	X	-	75.8	-	-
WFN	X	X	-	76.0	4.1	261
PWFN (no prior)	X	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+	✓	✓	X	77.8	-	-
Q-ViT	✓	✓	✓	78.1	11.3	3066917
Q-ViT (w/o FL-Bias)	✓	✓	✓	78.1	10.4	257149
PWFN (no prior)	X	X	X	78.0 (78.3)	2.7	352
PWFN	X	X	X	78.1 (78.5)	2.7	356
DeiT-Tiny	Baseline	-	-	-	72.9	15.5	5481081
LSQ+	✓	✓	X	68.1	-	-
Q-ViT	✓	✓	✓	69.6	11.5	1117630
Q-ViT (w/o FL-Bias)	✓	✓	✓	69.6	10.5	128793
PWFN (no prior)	X	X	X	71.4 (71.6)	2.8	300
PWFN	X	X	X	71.2 (71.5)	2.8	296
DenseNet161	Baseline	-	-	-	77.8	17.1	26423159
PWFN	X	X	X	77.6 (78.0)	1.1	125
We conduct our experimentation on the ImageNet dataset with a wide range of models: ResNets (18,34,50) [He et al., 2016], DenseNet-161 [Huang et al., 2017] and the challenging DeiT (small and tiny) [Touvron et al., 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-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  $\delta = 1$ ,  $\alpha = 2^{-11}$  which we found using grid-search on the Cifar-10 dataset and works surprisingly well in all settings. For all our experiments we train using 4x RTX8000 GPUs and a batch-size of 128. For the ensemble results, we sample 20 times different weights using the learned weights' distributions and report the mean accuracy.
We compare PwFN against a range of quantization approaches where the model weights have been made available so that we can 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	27
LSQ	90
QviT	300
For the ResNet family, we compare against the current state-of-the-art APoT [Li et al., 2019b]  ${}^{1}$  and WFN [Subia-Waud and Dasmahapatra, 2022]  ${}^{2}$  . For the transformer models, there has only been one work released, Q-Vit [Li et al., 2022]  ${}^{3}$  , 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 a straightforward regularization term that encourages an increase in  $\sigma$ , and its computational cost is comparable to that of 11 regularization. While our approach does lead to greater memory demands due to the additional  $\sigma$  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 1 we can see the set of results. PwFN demonstrates superior entropy, unique parameter count and top-1 accuracy across the board. In addition to 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 to 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 a 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 a 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 et 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 a network's 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. This observation suggests that having a 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.
<https://github.com/yhhhli/APoT_Quantization
$^{2}$ https://github.com/subiawaud/Weight_Fix_Networks
<sup>3</sup>https://github.com/YanjingLi0202/Q-ViT
Figure 5: The maximum (top left) and mean (bottom left) relative distance a weight moves to a 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.
Percentage Clustered
This work proposed PwFN, a training and clustering methodology that can both scale BNNs to complex model-dataset combinations and then use the weight uncertainties to inform quantization decision-making. The result is 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 a sampled ensemble giving noticeable accuracy gains. An exciting direction for future work is to explore how the uncertainty estimations and the out-of-distribution performance of neural networks could be enhanced using PwFN to train Bayesian Neural Networks.
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]
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-29378, 2021.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International conference on machine learning, pages 1613-1622. PMLR, 2015.
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 arXiv:1702.03044, 2017.
Shuchang Zhou, Yuxin Wu, Zekun Ni, Xinyu Zhou, He Wen, and Yuheng Zou. Dorefa-net: Training low bitwidth convolutional neural networks with low bitwidth gradients. arXiv preprint arXiv:1606.06160, 2016.
Bayes-by-backprop and PwFN The Bayesian posterior neural network distribution  $P(\boldsymbol{w}|\boldsymbol{D})$  is approximated by a distribution  $Q(\boldsymbol{w}|\theta)$  whose parameters  $\theta$  are trained using back-propagation, Bayes-by-backprop (BPP). The approximation is achieved by minimizing the Kullback-Leibler (KL) divergence  $D_{KL}[Q\|P]$  between  $P$  and  $Q$  to find the optimal parameters  $\theta^{*}$ . These parameters  $\theta^{*}$  instantiate the means  $\mu_{i}$  and variances  $\sigma_{i}^{2}$  of the PwFN.
$$
\theta^ {*} = \arg \min  _ {\theta} K L [ Q (\boldsymbol {w} | \theta) \| P (\boldsymbol {w} | D) ], \text {w h e r e}
$$
$$
K L [ Q (\boldsymbol {w} | \theta) \| P (\boldsymbol {w} | \boldsymbol {D}) ] = \mathbb {E} _ {Q (\boldsymbol {w} | \theta)} \log \left(\frac {Q (\boldsymbol {w} | \theta)}{P (\boldsymbol {w} | \boldsymbol {D})}\right) = \mathbb {E} _ {Q (\boldsymbol {w} | \theta)} \log \left(\frac {Q (\boldsymbol {w} | \theta) P (\boldsymbol {D})}{P (\boldsymbol {D} | \boldsymbol {w}) P (\boldsymbol {w})}\right). \tag {9}
$$
The  $P(D)$  term does not contribute to the optimization and is dropped, leaving
$$
\begin{array}{l} \theta^ {*} = \arg \min  _ {\theta} \mathbb {E} _ {Q (\boldsymbol {w} | \theta)} [ \log Q (\boldsymbol {w} | \theta) - \log P (\boldsymbol {D} | \boldsymbol {w}) - \log P (\boldsymbol {w}) ], \\ \approx \arg \min  _ {\theta} \sum_ {m} \underbrace {\log Q \left(\boldsymbol {w} ^ {m} \mid \theta\right) - \log P \left(\boldsymbol {w} ^ {m}\right)} _ {\text {p r i o r d e p e n d e n t}} - \underbrace {\log P (\boldsymbol {D} \mid \boldsymbol {w} ^ {m})} _ {\text {d a t a d e p e n d e n t}}, \tag {10} \\ \end{array}
$$
where the expectation value is approximated by samples  $\pmb{w}^m\sim Q(\pmb {w}|\theta)$  drawn from  $Q(\pmb {w}|\theta)$  each of which instantiates a neural network.
Gradient descent over each  $\pmb{w}^m$  that instantiates a neural network is made possible by the reparametrization trick [Kingma and Welling, 2013, Blundell et al., 2015]. The idea is to regard each sample  $\pmb{w}^m = \pmb{\mu} + \epsilon_m\pmb{\sigma}$  where  $\epsilon_{m}\sim p(\epsilon)$  is a random draw from some distribution that we take to be an isotropic Gaussian:  $p(\epsilon) = \mathcal{N}(0,I)$  with  $I$  the  $N$ -dimensional identity matrix for the  $N$  weights of network  $W$ . These weights  $\pmb{w}^m$  are used in the forward pass through the network while parameters  $\pmb{\mu}$  and  $\pmb{\sigma}$  are trainable. Then, for any function  $f(\pmb{w})$  we have  $\mathbb{E}_{Q(\pmb{w}|\theta)}[f(\pmb{w})] = \mathbb{E}_{p(\epsilon)}[f(\pmb{w})]$ , so that
$$
\frac {\partial}{\partial \theta} \mathbb {E} _ {Q (\boldsymbol {w} | \theta)} [ f (\boldsymbol {w}, \theta) ] = \frac {\partial}{\partial \theta} \mathbb {E} _ {p (\epsilon)} [ f (\boldsymbol {w}, \theta) ] = \mathbb {E} _ {p (\epsilon)} [ \frac {\partial f (\boldsymbol {w} , \theta)}{\partial \boldsymbol {w}} \frac {\partial \boldsymbol {w}}{\partial \theta} + \frac {\partial f (\boldsymbol {w} , \theta)}{\partial \theta} ].
$$
The terms are all calculable, allowing us to draw from a distribution for each weight  $\boldsymbol{w}^m$  and backpropagate to the underlying distribution parameters  $\theta \coloneqq (\pmb{\mu},\pmb{\sigma})$ . For  $\boldsymbol{w}^m = \boldsymbol{\mu} + \epsilon_m\boldsymbol{\sigma}$ , the derivatives are  $\frac{\partial\boldsymbol{w}^m}{\partial\boldsymbol{\mu}} = I$ , and  $\frac{\partial\boldsymbol{w}^m}{\partial\boldsymbol{\sigma}} = \epsilon_mI$ , making the respective gradients
$$
\nabla_ {\boldsymbol {\mu}} = \frac {\partial f (\boldsymbol {w} , \boldsymbol {\mu} , \boldsymbol {\sigma})}{\partial \boldsymbol {w}} + \frac {\partial f (\boldsymbol {w} , \boldsymbol {\mu} , \boldsymbol {\sigma})}{\partial \boldsymbol {\mu}} \text {a n d} \nabla_ {\boldsymbol {\sigma}} = \frac {\partial f (\boldsymbol {w} , \boldsymbol {\mu} , \boldsymbol {\sigma})}{\partial \boldsymbol {w}} \epsilon_ {m} + \frac {\partial f (\boldsymbol {w} , \boldsymbol {\mu} , \boldsymbol {\sigma})}{\partial \boldsymbol {\sigma}}. \tag {11}
$$
PWFN Clustering Algorithm
Prior Initialization
In addition to the prior initialization described in main paper, we added an reweighting determined by the size of the  $\sigma$  values in the network. Using the definition of  $v = D_{\mathrm{rel}}(\mu_i,\mathrm{pow}2_u(\mu_i))$  we re-weight by the third quartile  $\tilde{v}_{0.75}$  and re-write the initialization as:
$$
f \left(\mu_ {i}\right) = 0. 0 0 2 5 \times \frac {D _ {\text {r e l}} \left(\mu_ {i} , \operatorname {p o w} 2 _ {u} \left(\mu_ {i}\right)\right) \times D _ {\text {r e l}} \left(\mu_ {i} , \operatorname {p o w} 2 _ {d} \left(\mu_ {i}\right)\right)}{\tilde {v} _ {0 . 7 5}}, \tag {12}
$$
while  $|W_{\mathrm{fixed}}^{t + 1}|\leq Np_t$  do  
fixed  $\mathbf{\Pi}_{\mathrm{new}}\gets []$   
while fixed new is empty do  
Increase the order  $\omega \leftarrow \omega +1$   
Increase  $\delta \delta \leftarrow 2\delta$ $c^{\omega}\gets \{\sum_{i\in r}i|r\in \mathcal{P}(R)\wedge |r|\leq \omega \}$   
for each  $i = 1\dots ,|W_{\mathrm{free}}^{t + 1}|$ $c_{*}^{\omega}(i)\gets \min_{c\in C^{\omega}}D_{\mathrm{prob}}(w_i,c)$   
for each cluster center  $c_k^\omega \in C^\omega$ $n_k^\omega \gets \sum_i\mathbb{I}[c_k^\omega = c_*^\omega (i)]$ $k^{*}\gets \arg \max_{k}n_{k}^{\omega}$   
Sort:  $[w_1',\ldots ,w_N']\gets [w_1,\ldots ,w_N],w_i' = w_{\pi (i)},\pi$  permutation  
where  $D_{\mathrm{prob}}(w_i',c_{k*}^{\omega}) < D_{\mathrm{prob}}(w_{i + 1}'',c_{k*}^{\omega})$ $i\gets 1$  mean  $\leftarrow D_{\mathrm{prob}}(w_1',c_{k*}^{\omega})$   
while mean  $\leq \delta$  do  
 $\begin{array}{rl} & {\mathrm{fixed}_{\mathrm{new}}\gets w_i'}\\ & {\mathrm{mean}\leftarrow \frac{i}{i + 1}\times \mathrm{mean} + \frac{1}{i + 1}\times D_{\mathrm{prob}}(w_{i + 1}'',c_{k*}^{\omega})}\\ & {i\leftarrow i + 1} \end{array}$   
end  
 $\delta_{t = t}\gets \beta \times \delta_{t = t}$   
end  
Assign all the weights means  $\mu_i\in \mathrm{fixed}_{\mathrm{new}}$  to cluster center  $c_{*}^{\omega}(i)$  and set each of the  $\sigma_{i}\in \mathrm{fixed}_{\mathrm{new}}$  to be the variance of the weight means in fixed new. Finally, move them from  $W_{\mathrm{free}}^{t + 1}$  to  $W_{\mathrm{fixed}}^{t + 1}$
Algorithm 1: Clustering  $Np_{t}$  weights at the  $t^{th}$  iteration in PwFN.
and clamp the values to be within the range  $[2^{-30}, 0.05]$  giving us our initial variance values.
$$
\sigma_ {i} = \max  (0. 1, \min  (f (\mu_ {i}), 2 ^ {- 3 0}). \tag {13}
$$
Figure 6: Here we compare the  $\mu$  vs  $\sigma$  values for all weights in a given layer at initialization (left) and after PwFN convergence and clustering (right).