Stochastic Weight Sharing for Bayesian Neural Networks

arXiv:2505.17856v1 [cs.LG] 23 May 2025 Stochastic Weight Sharing for Bayesian Neural Networks Moule Lin Shuhao Guan Weipeng Jing Goetz Botterweck Andrea Patane Lero, Trinity College Dublin University College Dublin Northeast Forestry University Lero, Trinity College Dublin Lero, Trinity College Dublin Abstract While offering a principled framework for uncertainty quantification in deep learning, the employment of Bayesian Neural Networks (BNNs) is still constrained by their increased computational requirements and the convergence difficulties when training very deep, state-of-the-art architectures. In this work, we reinterpret weightsharing quantization techniques from a stochastic perspective in the context of training and inference with Bayesian Neural Networks (BNNs). Specifically, we leverage 2D-adaptive Gaussian distributions, Wasserstein distance estimations, and alpha-blending to encode the stochastic behaviour of a BNN in a lower-dimensional, soft Gaussian representation. Through extensive empirical investigation, we demonstrate that our approach significantly reduces the computational overhead inherent in Bayesian learning by several orders of magnitude, enabling the efficient Bayesian training of large-scale models, such as ResNet-101 and Vision Transformer (VIT). On various computer vision benchmarks—including CIFAR-10, CIFAR-100, and ImageNet1k—our approach compresses model parameters by approximately 50×and reduces model size by 75%, while achieving accuracy and uncertainty estimations comparable to the state-of-the-art. 1 INTRODUCTION Bayesian Neural Networks (BNNs) promise to combine the representational capacity of deep learning with principled uncertainty estimations enabled by means of Bayesian learning theory (Hinton and Neal,1995). Arguably, this combination makes them particularly appealing for safetyProceedings of the 28th International Conference on Artificial Intelligence and Statistics (AISTATS) 2025, Mai Khao, Thailand. PMLR: Volume 258. Copyright 2025 by the author(s). critical machine learning applications where the quantification of uncertainty is of paramount importance (Forsberg et al.,2020). Indeed, they have been widely employed in scenarios like e-Health (Marcos et al.,2010), robust control (Wicker et al.,2024), autonomous driving (Michelmore et al.,2020), human-in-the-loop applications (Treiss et al.,2021), automated diagnosis (Billah and Javed,2022) and many others (Lampinen and Vehtari,2001;Bharadiya, 2023;Vehtari and Lampinen,1999). Unfortunately, though, the principled treatment of uncertainty comes at the price of an increased pressure on computational resources, including the model size (×2 in the common case of mean-field Variational Inference (Blundell et al.,2015)), and the inference time, increased by an order of magnitude as multiple forward passes are needed (Hinton and Neal,1995). Therefore, despite their potential, the use of BNNs in edge-AI and resource-constrained applications is still very limited (Bonnet et al.,2023). While recent works have investigated the development of techniques to tackle the aforementioned challenges, these are generally limited to the application of methods originally developed for deterministic neural networks (NNs) (Ferianc et al.,2021;Park et al.,2021;Chien and Chang,2023), or the usage of the Bayesian paradigm at training time for model-order reduction purposes but without the uncertainty estimation at inference time (Van Baalen et al.,2020;Guo, 2018;Perrin et al.,2024;Subia-Waud and Dasmahapatra, 2024). In this work, we present a quantisation technique specifically tailored to capture the stochastic behaviour of BNNs. More specifically, we design a stochastic weight-sharing quantisation method, called 2DGBNN, based on dynamically adaptive mini-batch 2D Gaussian Mixture Models and predicated on optimising weight distributions through metrics derived from Wasserstein distances (Chizat et al., 2020;De Palma et al.,2021), network gradients, and intraclass variance. Our technique works by reinterpreting standard weight-sharing (Subia-Waud and Dasmahapatra, 2024) from a 2D perspective, accounting for both the mean and variance of BNN’s parameters in the case of meanfield Variational Inference (VI), and by giving it a stochastic semantic. At each training step, the current sum-total of Stochastic Weight Sharing for Bayesian Neural Networks network parameters is clustered using standard parameterfree techniques, and a mini-batch approach is used for the estimation of Gaussian distributions on a parameter space accounting for (possibly) millions of parameters. Representatives for each cluster are then selected, and alphablending techniques are used for sampling parameter realisations during forward passes through the network architecture. Thanks to its simplicity, the method can be seamlessly integrated into commonly used Bayesian approximation techniques based on Variational Inference (Blundell et al.,2015;Gal and Ghahramani,2016;Minka,2001; Welling and Teh,2011). We perform an extensive empirical investigation on the effectiveness of our method in training large-scale models and in reducing the computational footprint of BNNs. We utilize four widely used image classification datasets (i.e., MNIST (LeCun et al.,1998), CIFAR10 (Krizhevsky,2009), CIFAR-100 (Krizhevsky,2009) and ImageNet1K (Deng et al.,2009)) and test the model results in four widely employed neural network architectures, including ResNet-18, ResNet-50, ResNet-101 (He et al., 2016) and Vision Transformer (ViT) (Dosovitskiy et al., 2021). Through our approach, we can reduce the number of trainable parameters in the network by up to 50×, while obtaining accuracy and uncertainty metrics comparable to the state-of-the-art.1 This paper makes the following main contributions: •We introduce a stochastic weight-sharing technique specifically-tailored to BNNs and that employs Wasserstein distance, gradients, and within-class variance in order to improve model efficiency. •We empirically demonstrate that our stochastic weight-sharing method compares favourably against quantisation methods employed for BNNs in terms of further reducing the computational requirements and better preserving accuracy and uncertainty metrics. •In a variety of architectures, including ResNet-18, ResNet-50, ResNet-10 and ViT, we show how our technique can reduce model size by up to just a quarter of the original size, while working on par with stateof-the-art techniques for large-scale approximate BNN training. 2 RELATED WORK Computational efficiency is one of the long-standing issues concerning the application of BNNs to edge-AI and embedded systems (Bonnet et al.,2023). Indeed, there is a vast section of the literature that aims to tackle the issue 1To support reproducibility, our code is available at https: //github.com/moulelin/2DGBNN from a variety of different angles. One such closely related area is that of quantisation of the BNN’s parameters (Guo, 2018), where the precision of the latter is reduced to minimise their computational footprint. Works have adapted techniques initially developed for deterministic NNs (Ferianc et al.,2021;Subedar et al.,2021;Lin et al.,2023;Dong et al.,2022;Ullrich et al.,2017) or developed new techniques that take into account the distributional behaviour of BNNs’ parameters (Chien and Chang,2023;Park et al., 2021;Yang et al.,2020a). While these techniques increase the computational efficiency of a given BNN architecture, by developing a weight-sharing technique (and therefore not only quantising but also in effect reducing the number of BNN parameters) the method we introduce is able to match their behaviour in standard BNN benchmarks, while at the same time, it allows for training of large scale models (Hernández-Lobato and Adams,2015). Several works have looked at reducing the number of parameters in BNNs, either by applying pruning techniques (Sharma and Jennings,2021;Beckers et al.,2023;Roth and Pernkopf,2018) or using low-rank approximations (Doan et al.,2024;Dusenberry et al.,2020;Swiatkowski et al., 2020). The latter work by reparameterising BNN weights and biases using a lower rank representation, enabling them to scale BNN inference to large models such as ResNet-50 (Dusenberry et al.,2020) and ViT (Doan et al.,2024) architectures, and are as such closely related to our work in that a smaller common representation is found for BNN parameters. The number of final parameters is still though generally higher than their full deterministic counterpart, and therefore cannot be used for edge-AI applications. In Section 5, we will observe that, albeit at the price of a small reduction in accuracy, our method reduces the number of parameters of 3 orders of magnitude. Finally, a number of works have looked into applying Bayesian techniques for quantisation of deterministic NNs (Subia-Waud and Dasmahapatra,2024;Louizos et al., 2017;Van Baalen et al.,2020;Perrin et al.,2024;Achterhold et al.,2018;Soudry et al.,2014;Yang et al.,2020b), including the application of weight-sharing techniques (Roth and Pernkopf,2018;Subia-Waud and Dasmahapatra,2024; Nowlan and Hinton,2018). While these techniques provide encouraging results on the suitability of Bayesian theory for increasing the efficiency of deep learning, being specifically tailored to deterministic neural networks they cannot be applied to BNNs. 3 BAYESIAN NEURAL NETWORKS We consider a neural network architecture 𝑓𝐰∶ℝ𝑛→ℝ𝑚 parameterised by a vector of weights and biases 𝐰∈ℝ𝑛𝑤, which we refer to collectively as parameters of the neural network. Bayesian learning of neural networks begins by placing a prior distribution, 𝑝(𝐰), over the networks’ param- Moule Lin, Shuhao Guan, Weipeng Jing, Goetz Botterweck, Andrea Patane eter. This is often assumed to be encoded through a vector of independent Gaussian distributions, one for each weight andbiasinthe BNN (Blundell etal.,2015). Thisprior belief is then updated given a dataset’s evidence through the application of the Bayesian learning rule. Let = {(𝑥𝑖, 𝑦𝑖)}𝑛 𝑖=1 denote the full training dataset, 𝐗= (𝑥1,…, 𝑥𝑛)the combined vector of training inputs and 𝐲= (𝑦1,…, 𝑦𝑛)their corresponding outputs, then the posterior distribution on the weight is computed as: 𝑝(𝐰|𝐗,𝐲) = 𝑝(𝐲|𝐗,𝐰)𝑝(𝐰) 𝑝(𝐲|𝐗),(1) where 𝑝(𝐲|𝐗,𝐰)is the likelihood and 𝑝(𝐲|𝐗)is the model evidence. Finally, given a test point 𝑥∗, the BNN’s posterior predictive distribution on 𝑥∗is defined by: 𝑝(𝑦∗|𝑥∗,𝐗,𝐲) = ∫𝑝(𝑦∗|𝑥∗,𝐰)𝑝(𝐰|𝐗,𝐲)𝑑𝐰.(2) Unfortunately, neither Equation (1) nor Equation (2) can generally be computed exactly (Hinton and Neal,1995). Therefore a variety of approximate Bayesian inference techniques have been developed in the literature, with the two most prominent classes of approaches being based on Monte Carlo algorithms (Hinton and Neal,1995) or on Variational Inference methods (Blundell et al.,2015). While the former provides the gold standard in terms of approximation accuracy in small architectures, Variational Inference (VI) guarantees better scalability and is therefore the focus of this paper. Briefly, VI works by approximating the true posterior, 𝑝(𝐰|𝐗,𝐲), by optimising the KL divergence over a simpler, parameterised distribution, 𝑞(𝐰), often a multidimensional Gaussian distribution with diagonal covariance. The predictive distribution of Equation (2) is then approximated by sampling multiple times from 𝑞(𝐰), and averaging the results. Despite the approximation, however, VI BNNs still come with several limitations that impede their deployment in practice. First, even in the case of diagonal Gaussian distributions, the number of parameters in the BNN is doubled compared to their deterministic counterpart. Furthermore, approximating the predictive distribution requires multiple sampling procedures and multiple forward passes through the network so that their computational time is orders of magnitude higher than, again, their deterministic counterpart. Finally, despite its greater flexibility, standard Variational Inference struggles to learn very deep BNNs and it is generally limited to more traditional architecture and smallto-medium-size datasets. In the following, we develop a weight-sharing quantisation scheme targeted at BNNs to tackle these limitations. 4 2D GAUSSIAN BAYESIAN NEURAL NETWORK Consider the vector 𝐰of the BNN’s weights, where, at each step of the training process, each weight, 𝑤𝑖𝑖= 1,…, 𝑛𝐰, is distributed accordingly to a given Gaussian distribution (𝜇𝑤𝑖, 𝜎𝑤𝑖). We denote with 𝑓𝑢𝑙𝑙 the full set of weight distributions. Our stochastic weight sharing technique aims at finding a set of 2D Gaussian distributions (which we collectively denote as ws)(𝜇1,Σ1),…,(𝜇𝑘,Σ𝑘),2with 𝑘≪𝑛𝐰, and such that ws can be used to approximates the behaviour of 𝑓𝑢𝑙𝑙 in terms of resulting accuracy and uncertainty. Briefly, we do this by first modelling all the hyperparameters of the 𝑓𝑢𝑙𝑙 distributions through a Gaussian Mixture Model (GMM) (Section 4.2), and then applying alphablending to sample weights from the resulting realisations of the GMM (Section 4.3). Additionally, 2DGBNN implements several steps informed by best practice in quantisation for deterministic NNs for further reducing the shared number of weights, including outliers detection (Section 4.1), cluster dimensionality reduction and the merging of similar distributions (Section 4.2). Finally, we will present the overall algorithm for 2DGBNN in Section 4.4. 4.1 Outliers vs. Inliers Classification The key observation behind the weight-classification stage of our algorithm is that not all the weights of a neural network have an equal impact on the output. Taking inspiration from quantisation techniques for deterministic neural networks (Subedar et al.,2021), we, therefore, do not quantise extreme values in the BNN as those are, likely, particularly influential in the final result. Specifically, we partition the full weight vector winto two separate vectors, w𝑖𝑛 and w𝑜𝑢𝑡, and only apply weight-sharing to the former. We do this by using two different criteria. Mean Threshold: Weights associated with a mean with an absolute value greater than a threshold 𝜏(e.g., 𝜏= 0.2) are classified as outliers, i.e.: 𝑤𝑖∈𝐰is an outlier if |𝜇𝑖|> 𝜏, A discussion of how we chose parameters like 𝜏is provided in Appendix E. Gradient Threshold: Weights associated with gradient magnitudes exceeding the threshold that places them within the top 1% during backpropagation are also categorized as outliers, as their high gradient values likely signify their 2As standard we use 𝜎to denote the one-dimensional standard deviation in the case of 1d Gaussian, and Σto denote the multidimensional covariance in the case of multidimensional Gaussian. Stochastic Weight Sharing for Bayesian Neural Networks substantial impact on model performance, i.e.: 𝑤𝑖∈𝐰is an outlier if |∇𝑤𝑖|is top 1% of {|∇𝑤𝑗|}𝑛𝐰 𝑗=1. Algorithm 1 2DGBNN Input: NN architecture 𝑓𝐰, training data = {(𝐗,𝐲)} – 𝜏𝑤, 𝜏𝑑, 𝜏𝑔, 𝜏𝑣algorithm thresholds – BNN prior 𝑝(𝐰). Output: Stochastic weight-sharing trained BNN. Stage 1: Initialise GMM 1: Initialise 𝜇,𝜎according to 𝑝(𝐰) 2: for each epoch do ⊳Pre-training 3: Sample weights: 𝐰=𝜇+𝜎 ⊙ 𝜖,𝜖∼(0,𝐈) 4: Update 𝜇,𝜎by training on  5: end for 6: if |𝑤𝑖|> 𝜏𝑤or |∇𝑤𝑖|in top 1% then 𝑤𝑖is outlier 7: else 𝑤𝑖is inlier ⊳§4.1 8: end if 9: Learn GMM on inlier params Θ𝑖𝑛 (Equation (3)) Stage 2: Refine GMM 10: for each inlier weight 𝑤𝑖do 11: Perform §4.3 check on Mahalanobis distance 12: if Outside 95th percentile then 13: 𝑤𝑖is assigned to multiple clusters. 14: else 𝑤𝑖is assigned only to the closest Gaussian. 15: end if 16: end for 17: Apply alpha-blending for ellipse points (Eq. 7) 18: repeat 19: for each pair (1,2)in GMM do 20: if 𝑊(1,2)< 𝜏𝑑,Δ𝑔< 𝜏𝑔,Δ𝑣< 𝜏𝑣then 21: Merge 1,2using Eqs. (5), (6) 22: end if 23: end for 24: until No more Gaussians can be merged Stage 3: Final BNN Training 25: for each epoch do 26: for each weight 𝑤𝑖do 27: if 𝑤𝑖is inlier then 28: Sample 𝑤𝑖∼∑𝐾 𝑘=1 𝜋𝑘(𝜇𝑘,Σ𝑘) 29: else 30: Use 𝑤𝑖∼(𝜇𝑤𝑖, 𝜎2 𝑤𝑖) 31: end if 32: end for 33: Minimising step for  (, 𝑞)of Eq. (10) 34: end for 4.2 2DGBNN training Given the set of distributions of the inlier weights, which we denote as 𝑖𝑛, we proceed by clustering their means and variances on the 2-dimensional 𝜇-𝜎plane. We do this by learning a Gaussian Mixture Model (GMM) of the form: 𝑝((𝜇, 𝜎)) = ∑𝐾 𝑘=1 𝜋𝑘((𝜇, 𝜎)|𝜇𝑘,Σ𝑘)over the set of points Θ𝑖𝑛 = {(𝜇𝑤𝑖, 𝜎𝑤𝑖)}𝑛𝐰𝑖𝑛 𝑖=1 . Due to the large volume of points involved in the learning of the GMM (typically, millions of weights), we rely on mini-batch learning for GMMs (Li et al.,2014). This is achieved by sampling random mini-batches ⊂Θ𝑖𝑛, and iteratively minimising the log-likelihood over the mini-batch: min ∑ (𝜇𝑖,𝜎𝑖)∈log (𝐾 ∑ 𝑘=1 𝜋𝑘((𝜇𝑖, 𝜎𝑖) ∣ 𝜇𝑘,Σ𝑘)).(3) After the Initial GMM learning, we perform two further reduction steps based on the number of points around each Gaussian, and on the distance between pairs of Gaussians. Cluster Size Reductions: During the initial Gaussian clustering stage, clusters associated with fewer than 30 weights are identified. Weights within these small clusters are treated as outliers due to their lack of representation within the broader weight distribution. Overall, we empirically find that approximately 1.8% of the total weights of a neural network are generally allocated as outliers. Merging Gaussians: We merge together Gaussian distributions that are very close to each other. We do this by relying on the distance between Gaussians and their gradients. Specifically, we compute the Wasserstein-2 distance between pairs of distributions as (Jacobs et al.,2023): 𝑊2((𝜇𝑖,Σ𝑖),(𝜇𝑗,Σ𝑗))2=||𝜇𝑖−𝜇𝑗||2 2 +Tr(Σ𝑖+ Σ𝑗− 2(Σ1∕2 𝑖Σ𝑗Σ1∕2 𝑖)1∕2)(4) If the distance between two Gaussians is less than a given threshold 𝛾, then we inspect the gradient of the network in the weight associated to the cluster centroid and its variance. If those are smaller than two given threshold 𝜎and 𝛼 then we proceed by merging the two Gaussians into one.3 The merger is executed using the following equations (Agueh and Carlier,2011;Takatsu,2011): 𝜇merged =𝜇1+𝜇2 2(5) Σmerged =Σ1+ Σ2 2+1 8(𝜇1−𝜇2)(𝜇1−𝜇2)𝑇 +1 2(Σ1∕2 1Σ2Σ1∕2 1)1∕2 (6) Thatensuresthat the newlyformed Gaussian componentaccurately reflects the collective distribution characteristics of the initial components while maintaining minimal internal variation. 3Details about the thresholds we use in our experiments can be found in the Appendix. Moule Lin, Shuhao Guan, Weipeng Jing, Goetz Botterweck, Andrea Patane Table 1: Comparison of 2DGBNN and competitive techniques superscripts indicate matchingarchitectures)onImageNet1k dataset. We also provide the number of outliers, ellipses, and Gaussians derived by our method. Architecture Method Accuracy ↑NLL ↓ECE ↓#Outliers #Ellipses #Gaussians #Parameters(M) ↓/ Compression Ratio(%) ↑ ResNet-18 Mutual BNN (Pham et al.,2024) 67.7 1.327 0.1300 - - - 23.4M / - 2DGBNN(ours) 68.1 1.253 0.019 23013 10885 2217 0.038M /99% ResNet-50 Deep Ensembles (Lakshminarayanan et al.,2017) 77.5 0.877 0.0305 - - - 146.7M / - Rank-1 BNN(Dusenberry et al.,2020) 77.3 0.886 0.0166 - - - 26.0M / - ATMC (30 samples) (Heek and Kalchbrenner,2019) 77.5 0.883 - - - - 768.0M / - MCMC (9 samples) BNN (Zhang et al.,2019) 77.1 0.888 - - - - 230.4M / - 2DGBNN(ours) 75.1 0.961 0.029 37172 56873 3250 0.101M/99% ResNet-101 2DGBNN(ours) 75.50 0.969 0.023 53641 4311 2464 0.063M / 99% VIT-B-16 76.01 0.901 0.064 9765 338329 5440 0.359M / 98% 4.3 𝛼-blending (Multi-Clusters) for Weights Before the final sampling step, we reassess inlier weights 𝐰𝑖𝑛 by computing their squared Mahalanobis distances to cluster means using, i.e., 𝐷2= (𝑤𝑖−𝜇𝑖)⊤Σ−1 𝑘(𝑤𝑖−𝜇𝑖). If a weight’s 𝐷2exceeds 5.991,4we reassess its cluster assignment: We do this by relying on 𝛼-blending (Mildenhall et al.,2021). Specifically, for each the 𝑤𝑖∈𝐰𝑖𝑛 we compute the subset of GMM’s component (𝜇𝑘,Σ𝑘), for 𝑘= 1,…, 𝑛𝑖such that the above condition on the 𝐷2is met. We then sample the final value of the weight by the resulting distributions: 𝑝(𝑤𝑖) = 𝑛𝑖 ∑ 𝑘=1 𝛼𝑘(𝜇𝑘,Σ𝑘)(7) where 𝛼𝑘is the mixing coefficient, computed as the pdf of (𝜇𝑖, 𝜎2 𝑖)according to (𝜇𝑘,Σ𝑘). 4.3.1 Combined Variational Formulation Finally, we observe that our stochastic weight-sharing technique can be seamlessly integrated within the ELBO formulation for VI training (Nowlan and Hinton,2018;Zhang et al.,2018). Formally, we assume that 𝐰𝑜𝑢𝑡 and 𝐰𝑖𝑛 are vectors of pairwise independent weights,5which allow us 4Corresponding to the 95th percentile of the 𝜒2 2distribution which models Mahlanobis distance of multidimensional Gaussians. 5This is true for the variational distribution but it is an approximation for the true posterior. to bound the variational objective as it follows: (, 𝑞) = 𝔼𝑞(𝐰)[log 𝑝(𝐲∣𝐗,𝐰)]−KL(𝑞(𝐰)‖𝑝(𝐰))=(8) 𝔼𝑞(𝐰)[log 𝑝(𝐲∣𝐗,𝐰)]−KL(𝑞(𝐰𝑖𝑛)‖𝑝(𝐰𝑖𝑛))− KL(𝑞(𝐰𝑜𝑢𝑡)‖𝑝(𝐰𝑜𝑢𝑡))≈𝔼𝑞(𝐰)[log 𝑝(𝐲∣𝐗,𝐰)](9) −KL(∑𝜋𝑘𝑘‖𝑝(𝐰𝑖𝑛))−KL (𝑞(𝐰𝑜𝑢𝑡)‖𝑝(𝐰𝑜𝑢𝑡)) ≥𝔼𝑞(𝐰)[log 𝑝(𝐲∣𝐗,𝐰)]−∑ 𝑘 𝜋𝑘KL(𝑘‖𝑝(𝐰𝑖𝑛)) ⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟ GMM KL divergence − KL(𝑞(𝐰𝑜𝑢𝑡)‖𝑝(𝐰𝑜𝑢𝑡)) ⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟ Outliers KL divergence ∶=  (, 𝑞),(10) where the equality in Equation (8) is due to the pairwise independence assumption between components of 𝐰𝑖𝑛 and 𝐰𝑜𝑢𝑡, the approximation of Equation (9) is due to the GMM approximation of the inlier weights, and the final inequality is due to the convexity of the KL divergence. Notice that the resulting value loss function,  is an upper bound on the original loss so that its minimisation by means of gradient descent guarantees the improvement of the latter. 4.4 Overall Methodology The overall methodology is presented in pseudocode form in Algorithm 1. 2DGBNN combines its component parts in three stages. In the first stage, the BNN is initialised with the given prior, a pre-training step is performed and the resulting BNN is used to initialise the GMM clustering. In stage 2, the initial GMM clustering obtained is refined by merging close Gaussians, and by performing 𝛼-blending. Finally the BNN is trained by optimising the variational objective on a combination of full weight distributions (for the outliers) and GMM-based weight-sharing (for the inliers). Stochastic Weight Sharing for Bayesian Neural Networks (a) Density and Frequency for ResNet50 in ImageNet and CIFAR-100 (b) Weights Scatter ImageNet(above) and CIFAR-100 (below) Figure 1: Weight distribution for the BNN prior to stochastic-sharing. Panel (a) shows the density plot for the second convolutional layer of ResNet-50 when trained on ImageNet (in blue) and CIFAR-100 (in red). Panel (b) shows the corresponding scatter plots, including lines for the 1% and 99%. 5 EXPERIMENTS To validate the effectiveness and scalability of 2DGBNN, we conduct comprehensive experiments using various NN architectures on benchmark image classification datasets. This section details the datasets, models, experimental setup, results, and analysis of our findings. Specifically, we evaluate our method on four common image classification benchmarks: MNIST, CIFAR-10, CIFAR-100, and ImageNet1k; as well as four widely used NN architectures: ResNet-18, ResNet-50 and ResNet-101 and a Vision Transformer (ViT). The hyperparameters used and the details of the training are given in Appendix. Throughout this section, we compare our technique against the results obtained by Deep Ensembles (Lakshminarayanan et al., 2017), Rank-1 BNN (Dusenberry et al.,2020), MCMC BNN (Zhang et al.,2019), ATMC (Heek and Kalchbrenner,2019), Mutual BNN (Pham et al.,2024), F-SGVB-LRT Nguyen et al. (2024), ABNN (Franchi et al.,2024), LPBNN (Franchi et al.,2023), IR (Kim et al.,2023), SSVI (Li et al.,2024) and mBCNN (Kong et al.,2023). We evaluate the resulting models in terms of Accuracy, Negative Log-Likelihood (NLL, which measures the model’s uncertainty in its predictions), and Expected Calibration Error (ECE, which measures the calibration of predicted probabilities (Guo et al.,2017)). Each experiment is conducted three times with different random seeds, and we report the average results. We conduct experiments with all the aforementioned comparison models, totaling 13 2DGBNN experiments, which include the base models of all comparison methods. Additionally, we perform two quantisation comparison experiments (Section 5.2) and an ablation study (Section 5.3). 5.1 Performance Evaluation The results obtained with 2DGBNN and those of the stateof-the-art are listed in Tables 1,2and 3for ImageNet1k, CIFAR-100 and CIFAR-10 respectively along with details on the computations performed by 2DGBNN. In the case of ImageNet1k (Table 1) the comparison is performed against Deep Ensembles (Lakshminarayanan et al., 2017), Rank-1 BNN (Dusenberry et al.,2020), MCMC BNN with 9 samples (Zhang et al.,2019), ATMC (Heek and Kalchbrenner,2019), and Mutual BNN (Pham et al., 2024). We observe that, in all cases, our method successfully reduces the number of parameters by 3 or 4 orders of magnitudes.While the accuracy is reduced by around 2% in the ResNet-50 case, we do obtain comparable uncertainty estimation as evaluated by NLL and ECE. Similar results we obtain on the CIFAR-100 dataset (Table 2), comparing against Deep Ensembles (Lakshminarayanan et al.,2017), Rank-1 BNN(Dusenberry et al.,2020), FSGVB-LRT (Nguyen et al.,2024) and ABNN (Franchi et al.,2024). Our method achieves a substantial reduction in model parameters while maintaining or improving performance compared to other methods at the price of approximately 2% when compared to Deep Ensembles and Rank-1 BNN. For instance, for ResNet-18 we use only 0.019M parameters, which is drastically lower than the 23.4M parame- Moule Lin, Shuhao Guan, Weipeng Jing, Goetz Botterweck, Andrea Patane Table 2: Comparison of 2DGBNN and competitive techniques (superscripts indicate matching architectures) on CIFAR100 dataset. We also provide the number of outliers, ellipses, and Gaussians derived by our method. Architecture Method Accuracy ↑NLL ↓ECE ↓#Outliers #Ellipses #Gaussians #Parameters(M) ↓/ Compression Ratio(%) ↑ ResNet-18 F-SGVB-LRT (Nguyen et al.,2024) 70.1 1.121 0.036 - - - 23.4M / - SSVI (Li et al.,2024) 75.8 - 0.001 - - - 2.32M / 90% mBCNN (Kong et al.,2023) 73.7 1.004 0.002 - - - 2.86M / 87.8% 2DGBNN(ours) 74.7 1.053 0.038 14624 260 2387 0.019M /99% WRN-28-10 Deep Ensembles (Lakshminarayanan et al.,2017) 82.7 0.666 0.021 - - - 146M / - Rank-1 BNN (Dusenberry et al.,2020) 82.4 0.689 0.012 - - - 36.6M / - LP-BNN (Franchi et al.,2023) 79.3 - 0.0702 - - - 26.8M / 63% 2DGBNN(ours) 80.5 0.798 0.0432 40354 341 2390 0.045M /99% ResNet-50 ABNN (Franchi et al.,2024) 74.20 0.828 4.5 - - - 54.2M / - 2DGBNN(ours) 78.1 0.986 0.107 247591 330 1980 0.251M /99% ResNet-101 2DGBNN(ours) 78.4 0.834 0.066 45240 348 3199 0.052M /99% Table 3: Comparison of 2DGBNN and competitive techniques (superscripts indicate matching architectures) on CIFAR-10 dataset. We also provide the number of outliers, ellipses, and Gaussians derived by our method. Architecture Method Accuracy ↑NLL ↓ECE ↓#Outliers #Ellipses #Gaussians #Parameters(M) ↓/ Compression Ratio(%) ↑ ResNet-18 F-SGVB-LRT (Nguyen et al.,2024) 90.31 0.262 0.014 - - - 23.4M / - SSVI (Li et al.,2024) 93.74 - 0.006 - - - 1.17M / 95% mBCNN (Kong et al.,2023) 93.20 0.220 0.008 - - - 0.93M / 96% 2DGBNN(ours) 91.72 0.305 0.019 123310 67 1569 0.018M /99% WRN-28-10 Deep Ensembles (Lakshminarayanan et al.,2017) 96.2 0.143 0.020 - - - 146M / - Rank-1 BNN (Dusenberry et al.,2020) 96.3 0.128 0.008 - - - 36.6M / 50.8% LP-BNN (Franchi et al.,2023) 95.0 - 0.009 - - - 26.8M / 63% 2DGBNN(ours) 95.2 0.142 0.012 39395 365977 3950 0.413M /99% ResNet-50 ABNN (Franchi et al.,2024) 95.01 0.160 1.0 - - - 54.2M / 25% 2DGBNN(ours) 93.84 0.223 0.012 129640 01628 0.132M /99% ResNet-101 2DGBNN(ours) 92.78 0.270 0.015 45240 348 3199 0.052M /99% ters used by the F-SGVB-LRT model, yet we achieve competitive accuracy. Finally, Table 3lists analogous results in the context of CIFAR-10, comparing against IR (Kim et al.,2023), F-SGVB-LRT (Nguyen et al.,2024), ABNN (Franchi et al.,2024) and LP-BNN (Franchi et al.,2023). Additionally to the architectures used for comparisons, the tables report results for ResNet-101 and ViT. In these architectures too, 2DGBNN is able to reduce the number of parameters while obtaining accuracy and uncertainty metrics on par with that of state-of-the-art techniques across the remaining architectures. Interestingly, observing the training results obtained we notice how the distribution of weights in models trained on smaller datasets (CIFAR-100 and CIFAR-10) tends to cluster near zero, as depicted in Figure 1(b) (down). Conversely, in ImageNet1k the weight distribution is broader, as can be seen from Figure 1(a)) comparing empirical distributions obtained on CIFAR-100 and ImageNet1k. Notice how this translates to, for example, the ResNet-50 model trained on ImageNet1k dataset to have a significantly greater number of Gaussian and Ellipse weights than when trained on the CIFAR-100 dataset. 5.2 Comparison against Quantisation We now compare 2DGBNN against quantisation techniques applied to BNNs (Subedar et al.,2021). For this purpose, we remove the pretraining stage of 2DGBNN so to mimic the “vanilla” BNN training employed by Subedar et al. (2021). We use a Gaussian prior with a mean of 0 and a standard deviation of 0.1. Notice that these experiments are limited to CIFAR-10 and MNIST as the vanilla training of BNNs used in Subedar et al. (2021) does not scale to the larger architectures and datasets analysed in the previous section. The comparative results are presented in Table 4. Accuracy values areverysimilaracrosstheboard, whileourtechnique obtains significantly better uncertainty metrics, except for NLL in the case CIFAR-10. Stochastic Weight Sharing for Bayesian Neural Networks Table 4: Comparison against the quantisation technique of Subedar et al. (2021) on CIFAR-10 and MNIST. Datasets Algorithm Quantisation technique Accuracy ↑NLL ↓ECE ↓#Parameters (MB: Megabyte) #Outliers #Ellipses #Gaussians #Parameters CIFAR-10 BNNs Quantization (Subedar et al.,2021) ResNet-20 (INT8 SIGMA4) 90.92 0.266 1.778 - - - 0.87 MB ResNet-20 (INT8 SIGMA2) 90.85 0.273 2.547 - - - 0.72 MB ResNet-20 (INT8 SIGMA1) 90.96 0.266 0.711 - - - 0.54 MB 2DGBNN ResNet-20(without pretrained) 90.91 0.303 0.040 74181 3634 142 0.644MB (1.62MB) ResNet-20(with pretrained) 91.04 0.303 0.037 14624 260 2387 0.020M /(0.71MB) MNIST BNNs Quantization (Subedar et al.,2021) ResNet-20 (INT8 SIGMA4) 99.36 0.020 0.215 - - - 0.10 MB ResNet-20 (INT8 SIGMA2) 99.32 0.024 0.277 - - - 0.08 MB ResNet-20 (INT8 SIGMA1) 99.34 0.027 0.351 - - - 0.06 MB 2DGBNN ResNet-20(without pretrained) 99.52 0.013 0.001 3206 581 237 0.092MB (0.403MB) Table 5: Ablation Study: Impact of Outliers and Ellipses on CIFAR-10 Using ResNet-20 Configuration #Outliers #Ellipses Accuracy (%) ↑NLL ↓ECE↓ 2DGBNN ✓ ✓ 90.92 0.265 0.040 Without Ellipse ✓– 90.43 0.304 0.043 Without Outliers – ✓90.18 0.308 0.026 Without Outliers and Ellipse – – 90.01 0.318 0.029 In terms of the model size (here compared in Megabytes), the two techniques compare similarly when it comes to the size of trainable parameters (corresponding to the value reported not in brackets for 2DGBNN), with quantisation having a slight edge when only 1 bit is used for encoding the standard deviation. Notice, however, that in small NNs (like the one here analysed) our techniques incur significant storage overhead in that we need to keep an index (encoded in uint8) that assigns each inlier weight to its cluster. When this value is added (size reported in brackets in the Table) quantisation has a significant advantage over our storage requirements. While techniques such as Huffman coding or multi-level index tables can potentially reduce the size of the index vector by several factors, we leave further investigations to future work, and here notice that despite maintaining full precision on the workings of the BNN, weightsharing quantisation can already obtain comparable results to int8 quantisation. We notice that the two methods are complimentary, and int8 quantisation can further reduce the storage requirements of the outlier weights and GMMs. 5.3 Ablation Study Table 5presents an ablation study on CIFAR-10 using ResNet-20 to explore how outliers and ellipses contribute to the performance of 2DGBNN. When both outliers and ellipses are included, the model achieves an accuracy of 90.92%, with the lowest NLL of 0.265 and ECE of 0.040. However, removing either component significantly impacts performance. Without ellipses, the accuracy drops by 0.49% to 90.43%, and excluding outliers reduces it slightly further by 0.25% to 90.18%. When both are removed, the accuracy drops by 0.87% to a low of 90.05%. 6 CONCLUSIONS We have presented a stochastic weight-sharing quantisation technique based on GMMs specifically tailored to BNNs. In an extensive empirical evaluation, we have seen how our technique can significantly reduce the effective number of parameters of a BNN while obtaining results on par with state-of-the-art in large datasets and architectures such as ImageNet1k and ViT. Future work will explore how to integrate our method into a fully Bayesian framework and the application of further quantisation for the outlier weights. We have presented a stochastic weight-sharing quantisation technique based on GMMs specifically tailored to BNNs. In an extensive empirical evaluation, we have seen how our technique can significantly reduce the effective number of parameters of a BNN while obtaining results on par with state-of-the-art in large datasets and architectures such as ImageNet1k and ViT. Future work will explore how to integrate our method into a fully Bayesian framework and the application of further quantisation for the outlier weights. Moule Lin, Shuhao Guan, Weipeng Jing, Goetz Botterweck, Andrea Patane 7 ACKNOWLEDGEMENTS This publication has emanated from research jointly funded by European Union’s Horizon Europe 2021–2027 framework programme, Marie Skłodowska-Curie Actions, Grant Agreement No. 101072456 and Taighde Éireann – Research Ireland under grant number 13/RC/2094_2. References Achterhold, J., Koehler, J. M., Schmeink, A., and Genewein, T. (2018). Variational network quantization. In International conference on learning representations. Agueh, M. and Carlier, G. (2011). Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis, 43(2):904–924. Beckers, J., Van Erp, B., Zhao, Z., Kondrashov, K., and De Vries, B. (2023). Principled pruning of bayesian neural networks through variational free energy minimization. IEEE Open Journal of Signal Processing. Bharadiya, J. P. (2023). A review of bayesian machine learning principles, methods, and applications. International Journal of Innovative Science and Research Technology, 8(5):2033–2038. Billah, M. E. and Javed, F. (2022). Bayesian convolutional neural network-based models for diagnosis of blood cancer. Applied Artificial Intelligence, 36(1):2011688. Blundell, C., Cornebise, J., Kavukcuoglu, K., and Wierstra, D. (2015). Weight uncertainty in neural networks. In Proceedings of the 32nd International Conference on Machine Learning, pages 1613–1622. Bonnet, D., Hirtzlin, T., Majumdar, A., Dalgaty, T., Esmanhotto, E., Meli, V., Castellani, N., Martin, S., Nodin, J.-F., Bourgeois, G., et al. (2023). Bringing uncertainty quantification to the extreme-edge with memristorbased bayesian neural networks. Nature Communications, 14(1):7530. Chien, J.-T. and Chang, S.-T. (2023). Bayesian asymmetric quantized neural networks. Pattern Recognition, 139:109463. Chizat, L., Roussillon, P., Léger, F., Vialard, F.-X., and Peyré, G. (2020). Faster wasserstein distance estimation with the sinkhorn divergence. Advances in Neural Information Processing Systems, 33:2257–2269. De Palma, G., Marvian, M., Trevisan, D., and Lloyd, S. (2021). The quantum wasserstein distance of order 1. IEEE Transactions on Information Theory, 67(10):6627– 6643. Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. (2009). Imagenet: A large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248–255. IEEE. Doan, B. G., Shamsi, A., Guo, X.-Y., Mohammadi, A., Alinejad-Rokny, H., Sejdinovic, D., Ranasinghe, D. C., and Abbasnejad, E. (2024). Bayesian low-rank learning (bella): A practical approach to bayesian neural networks. arXiv preprint arXiv:2407.20891. Dong, R., Tan, Z., Wu, M., Zhang, L., and Ma, K. (2022). Finding the task-optimal low-bit sub-distribution in deep neural networks. In International Conference on Machine Learning, pages 5343–5359. PMLR. Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. (2021). An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations. Dusenberry, M., Jerfel, G., Wen, Y., Ma, Y., Snoek, J., Heller, K., Lakshminarayanan, B., and Tran, D. (2020). Efficient and scalable bayesian neural nets with rank-1 factors. In International conference on machine learning, pages 2782–2792. PMLR. Ferianc, M., Maji, P., Mattina, M., and Rodrigues, M. (2021). On the effects of quantisation on model uncertainty in bayesian neural networks. In Uncertainty in Artificial Intelligence, pages 929–938. PMLR. Forsberg, H., Lindén, J., Hjorth, J., Månefjord, T., and Daneshtalab, M. (2020). Challenges in using neural networks in safety-critical applications. In 2020 AIAA/IEEE 39th Digital Avionics Systems Conference (DASC), pages 1–7. IEEE. Franchi, G., Bursuc, A., Aldea, E., Dubuisson, S., and Bloch, I. (2023). Encoding the latent posterior of bayesian neural networks for uncertainty quantification. IEEE Transactions on Pattern Analysis and Machine Intelligence. Franchi, G., Laurent, O., Leguéry, M., Bursuc, A., Pilzer, A., and Yao, A. (2024). Make me a bnn: A simple strategy for estimating bayesian uncertainty from pre-trained models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12194– 12204. Gal, Y. and Ghahramani, Z. (2016). Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pages 1050–1059. PMLR. Guo, C., Pleiss, G., Sun, Y., and Weinberger, K. Q. (2017). On calibration of modern neural networks. In International conference on machine learning, pages 1321– 1330. PMLR. Guo, Y. (2018). A survey on methods and theories of quantized neural networks. arXiv preprint arXiv:1808.04752. He, K., Zhang, X., Ren, S., and Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of Stochastic Weight Sharing for Bayesian Neural Networks Table 9: Deterministic neural networks used as prior for the Initialisation 2DGBNN in the experiments discussed in Section 5.1. Method Dataset Accuracy ↑NLL ↓ECE ↓#Parameters (M: million) ImageNet-1k ResNet-18 69.70 1.265 0.027 11.7M ResNet-50 76.10 0.989 0.035 25.6M ResNet-101 77.30 0.936 0.936 44.5M VIT-B-16 81.07 0.856 0.056 86.0M CIFAR-100 ResNet-18 77.10 1.038 0.114 11.7M ResNet-50 79.20 0.950 0.054 25.6M ResNet-101 80.02 0.849 0.095 44.5M WRN-28-10 81.41 0.766 0.045 53.6M CIFAR-10 ResNet-20 91.84 0.246 0.031 0.27M ResNet-18 93.21 0.201 0.022 11.7M ResNet-50 94.81 0.211 0.010 25.6M ResNet-101 94.70 0.849 0.095 44.5M WRN-28-10 95.92 0.131 0.010 53.6M The table is structured to highlight the performance across multiple architectures and datasets, facilitating a direct comparison. Higher accuracies and lower NLL and ECE values indicate better model performance. For instance, VIT-B-16 on ImageNet-1k achieves an accuracy of 81.07% with the lowest ECE of 0.056 among its dataset counterparts. Similarly, WRN-28-10 shows superior performance on CIFAR-10 with the highest accuracy of 95.92% and a remarkably low ECE of 0.010. The number of parameters, reported in millions, also provides insight into the model complexity, ranging from 0.27M for ResNet-20 on CIFAR-10 to 86.0M for VIT-B-16 on ImageNet-1k. The 2DGBNN models on ImageNet1k showed a decrease in accuracy, with our ResNet-50 configuration dropping from the benchmark high of 77.5% to 75.10%. This reduction was coupled with a substantial decrease in the number of parameters—from models requiring up to 25.6 million parameters to just 0.101 million. Despite these changes, the increases in NLL and ECE were minimal and within acceptable ranges, indicating that the models maintain satisfactory predictive performance and calibration despite the reduced complexity. Similar trends were observed in the CIFAR-100 dataset, where our WRN-28-10 model’s accuracy decreased from 81.41% to 80.5%, while significantly reducing the parameter count to only 0.045 million from 53.6 million. The NLL and ECE metrics, although slightly elevated, remained competitive, affirming the effective uncertainty estimation capabilities of the models despite their reduced complexity. G Additional Background KL Divergence between Gaussians The KL divergence between two Gaussian distributions 𝑞(𝑤) = (𝑤|𝜇𝑞, 𝜎2 𝑞)and 𝑝(𝑤) = (𝑤|𝜇𝑝, 𝜎2 𝑝)is given by: KL(𝑞(𝑤)‖𝑝(𝑤)) = 1 2(𝜎2 𝑞 𝜎2 𝑝 +(𝜇𝑝−𝜇𝑞)2 𝜎2 𝑝 − 1 + ln 𝜎2 𝑝 𝜎2 𝑞)(11) This formula can be used to compute the KL divergence terms in the ELBO expression for both inliers and outliers. Moule Lin, Shuhao Guan, Weipeng Jing, Goetz Botterweck, Andrea Patane Expected Log-Likelihood Computation The expected log-likelihood term involves an expectation over the variational posterior: 𝔼𝑞(𝐰)[log 𝑝(𝐲|𝐗,𝐰)]=∫𝑞(𝐰) log 𝑝(𝐲|𝐗,𝐰)𝑑𝐰(12) In practice, this integral is intractable and is approximated using Monte Carlo sampling. H Wasserstein-based 2D Gaussian Merging In the context of comparing two Gaussian distributions, the Wasserstein distance provides a meaningful way to measure the distance between probability distributions. Specifically, for two 2D Gaussian distributions (𝜇1,Σ1)and (𝜇2,Σ2), where 𝜇1, 𝜇2are the means and Σ1,Σ2are the covariance matrices, the Wasserstein-2 distance is given by the following formula: 𝑊2((𝜇1,Σ1),(𝜇2,Σ2))2=||𝜇1−𝜇2||2 2+Tr(Σ1+ Σ2− 2(Σ1∕2 1Σ2Σ1∕2 1)1∕2)(13) Where ||𝜇1−𝜇2||2is the Euclidean distance between the means of the two distributions, Tr is the trace operator, which sums the diagonal elements of a matrix, Σ1∕2 1refers to the matrix square root of the covariance matrix Σ1. The Wasserstein distance threshold is set at 1.5 × 10−7, with further discussion on the I.1. If the distance between Gaussian components falls below this, gradient information is further considered, ensuring a precise analysis of component similarity. This method ensures that the newly formed Gaussian component accurately reflects the collective distribution characteristics of the initial components while maintaining minimal internal variation. I Sub-module Discussion and Hyper-parameter Configuration I.1 Effectiveness of 2D Gaussian Merging Discussion In this section, we discuss the effectiveness of the merging of 2D Gaussian distributions in our method. Experimental results analysing its effect are listed in Table 10. For the ResNet20 model on the CIFAR-10 dataset, after merging the 2D Gaussian, the accuracy improved from 88.77% to 91.02%, the NLL decreased from 0.3792 to 0.3066, and the ECE reduced from 0.0500 to 0.0374. For the ResNet18 model on the CIFAR-100 dataset, although the improvement is smaller, the accuracy still increased from 74.50% to 74.63%, and the NLL decreased from 1.0555 to 1.0535. These results demonstrate the effectiveness of our 2D Gaussian merging method in enhancing model performance and calibration. Table 10: Performance of ResNet20 on CIFAR-10 and ResNet18 on CIFAR-100 before and after merging Gaussian distributions. Model Dataset Method Accuracy (%) NLL ECE ResNet20 CIFAR-10 Before Merging Gaussian 88.77 0.3792 0.0500 After Merging Gaussian 91.02 0.3066 0.0374 ResNet18 CIFAR-100 Before Merging Gaussian 74.50 1.0555 0.0391 After Merging Gaussian 74.63 1.0535 0.0400 I.2 Data Augmentation To improve the robustness and generalization of our models, we applied a series of data augmentation techniques during training on the ImageNet-1K, CIFAR-100, and CIFAR-10 datasets. Table 11 summarises the specific augmentation methods used for each dataset. For the ImageNet-1K dataset, we applied random resized cropping to obtain images of size 224× 224 pixels. This was followed by random horizontal flipping with a probability of 50% to augment the dataset with mirrored images. Color jittering was used to adjust the brightness, contrast, saturation, and hue of the images with factors of 0.4, 0.4, Stochastic Weight Sharing for Bayesian Neural Networks Table 11: Summary of data augmentation techniques applied to each dataset. Dataset Data Augmentation Techniques ImageNet-1K Random resized crop to 224 × 224 pixels, Random horizontal flip, Color jittering (brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1), Conversion to tensor, Normalization CIFAR-100 Conversion to tensor, Padding of 4 pixels (reflection mode), Random crop to 32 × 32 pixels, Random horizontal flip, Conversion to tensor, Normalization CIFAR-10 Random crop to 32 × 32 pixels, Random horizontal flip, Conversion to tensor, Normalization 0.4, and 0.1, respectively. The images were then converted to tensors and normalised using the standard mean and standard deviation values for ImageNet. In the case of the CIFAR-100 dataset, we started by converting the images to tensors. We then padded the images with 4 pixels on each side using reflection mode to preserve edge information. After padding, we performed a random crop to 32 × 32 pixels, followed by random horizontal flipping to introduce mirror variations. The images were converted back to tensors and normaliseded accordingly. For the CIFAR-10 dataset, the augmentation process involved a random crop to 32 × 32 pixels, which helps in teaching the model to be invariant to translations. We also applied random horizontal flipping to include mirrored versions of the images. Finally, the images were converted to tensors and normalised to standardise the input data. I.3 Hyperparameters for training the deterministic network In our experiments, similar to those conducted by previous researchers, we standardised the hyperparameters across all models in the initial stage. This approach was applied to various models including ResNet-18, ResNet-50, ResNet-101, and VIT. The hyperparameters used are as follows: Hyperparameter Variable Default Value Batch Size -b 256 Warm-up Phases -warm 2 Learning Rate -lr 0.1 Resume Training -resume False Total Epochs -EPOCH 250 Milestones -MILESTONES [30, 60, 90, 120, 150, 200] Weight Decay -MultiStepLR 5e-4 Training Mean –TRAIN_MEAN (0.5071, 0.4865, 0.4409) Training Std -TRAIN_STD (0.2673, 0.2564, 0.2761) Table 12: Summary of hyperparameters in the neural network training configuration The table summarizes the standardised hyperparameters used across all models in our neural network training configurations, mirroring settings from previous research. It outlines common parameters such as batch size, warm-up phases, Moule Lin, Shuhao Guan, Weipeng Jing, Goetz Botterweck, Andrea Patane learning rate, and total epochs, alongside specific settings like weight decay and learning rate milestones. I.4 Discussion the Initialisation of 𝜎 We experimented with several different initialisation methods for our models, including initialisation via a specific function as detailed by Lee et al., random generation, and Gaussian distribution, among others. According to our experimental results, we ultimately adopted the following initialisation methods for our neural network parameters. For the weight parameter weight_sigma, we used the Xavier uniform initialisation with a gain of 0.01, defined as: 𝐖∼(−𝑔 √𝑛in +𝑛out ,𝑔 √𝑛in +𝑛out )(14) where 𝑔= 0.01,𝑛in is the number of input units, 𝑛out is the number of output units, and (𝑎, 𝑏)denotes a uniform distribution between 𝑎and 𝑏. For the bias parameter bias_sigma, we initialised it using a normal distribution with a mean of 0.0 and a standard deviation of 0.001: 𝐛∼(0,0.0012)(15) In our neural network, we assign different learning rates to different parameters. Table 13 summarises these hyperparameters. Table 13: Summary of Hyperparameters Used in Our Experiments Hyperparameter Symbol Value Learning rate for weight and bias 𝜇 𝜂weight_mu 1 × 10−4 Learning rate for weight and bias 𝜎 𝜂weight_sigma 1 × 10−2 We set the learning rates for the parameters as follows: the learning rates for weight_mu and bias_mu are both 𝜂weight_mu = 1 × 10−4; the learning rates for weight_sigma and bias_sigma are both 𝜂weight_sigma = 1 × 10−3. I.5 Hyper-parameters for training Bayesian Neural Network In our experiments, we utilise several hyperparameters that are crucial for the performance and convergence of our Bayesian neural network (BNN) model. These hyperparameters are carefully selected based on empirical studies to balance computational efficiency and model accuracy. Table 14 summarises the hyperparameters used in our experiments. In the KMeans clustering algorithm, we initially use 𝐾= 2000 clusters. Outliers in the weight values are identified using a threshold 𝑇𝑤= ±0.2. Any weight value exceeding this threshold is considered an outlier. This threshold is chosen based on the empirical distribution of the weights after initial training. Similarly, outliers in the gradients are identified by selecting the top 𝑃𝑔= 1% of gradient magnitudes. A minimum cluster size of 𝑁min = 30 is enforced to ensure statistical significance in the clustering results. Clusters with fewer than 𝑁min samples are considered invalid and their associated weights are treated as outliers. This prevents the model from being influenced by clusters that may represent noise or insignificant patterns. During the BNN training, we use a learning rate of 𝜂BNN = 1 × 10−5, which is lower than the initial learning rate used in the preliminary training. The smaller learning rate is necessary to accommodate the Bayesian updates and to ensure that the posterior distributions over the weights converge properly. In the predictive function, we draw 𝑁𝑠= 30 samples from the posterior distribution to estimate the predictive mean and uncertainty. The Expected Calibration Error (ECE) is computed using 𝑁𝑏= 15 bins. The Mahalanobis distance threshold 𝑇𝑀= 5.991 corresponds to the chi-squared distribution value with 2 degrees of freedom at the 95% confidence level. The Mahalanobis distance for a data point 𝑥with respect to a Gaussian distribution with mean 𝜇and covariance Σis calculated as: 𝐷2 𝑀= (𝑥−𝜇)⊤Σ−1(𝑥−𝜇).(16) Stochastic Weight Sharing for Bayesian Neural Networks Table 14: Summary of Hyperparameters Used in Our Experiments Hyperparameter Symbol Value Initial learning rate for 𝜇 𝜂𝜇1 × 10−4 Number of epochs 𝐸200 Number of clusters in KMeans 𝐾6000 Outlier threshold (weight value) 𝑇𝑤±0.2 Outlier threshold (gradient percentile) 𝑃𝑔Top 1% Minimum samples per cluster 𝑁min 20 Learning rate in BNN training 𝜂BNN 1 × 10−5 Number of samples in predictive function 𝑁𝑠30 Number of bins in ECE computation 𝑁𝑏15 Wasserstein distance threshold 𝑇𝑊1 × 10−2 Mahalanobis distance threshold 𝑇𝑀5.991 Number of nearest Gaussians 𝑘5 Points with a Mahalanobis distance greater than 𝑇𝑀are considered outliers. In handling outliers, we consider the 𝑘= 5 nearest Gaussian components for each outlier point. This allows us to reassign outlier weights to the most probable Gaussian components based on their proximity in the parameter space.