Training Bayesian Neural Networks with Sparse Subspace Variational Inference
Authors: Junbo Li
Zichen Miao
Qiang Qiu
Ruqi Zhang
Venue: ICLR 2024 poster (2024)
Keywords: Bayesian neural networks
sparse Bayesian learning
variational inference
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We propose the first fully sparse Bayesian training framework that achieves state-of-the-art performance in the realm of sparse Bayesian neural networks.

Abstract: Bayesian neural networks (BNNs) offer uncertainty quantification but come with the downside of substantially increased training and inference costs. Sparse BNNs have been investigated for efficient inference, typically by either slowly introducing sparsity throughout the training or by post-training compression of dense BNNs. The dilemma of how to cut down massive training costs remains, particularly given the requirement to learn about the uncertainty. To solve this challenge, we introduce Sparse Subspace Variational Inference (SSVI), the first fully sparse BNN framework that maintains a consistently sparse Bayesian model throughout the training and inference phases. Starting from a randomly initialized low-dimensional sparse subspace, our approach alternately optimizes the sparse subspace basis selection and its associated parameters. While basis selection is characterized as a non-differentiable problem, we approximate the optimal solution with a removal-and-addition strategy, guided by novel criteria based on weight distribution statistics. Our extensive experiments show that SSVI sets new benchmarks in crafting sparse BNNs, achieving, for instance, a 10-20× compression in model size with under 3\% performance drop, and up to 20× FLOPs reduction during training. Remarkably, SSVI also demonstrates enhanced robustness to hyperparameters, reducing the need for intricate tuning in VI and occasionally even surpassing VI-trained dense BNNs.

Published as a conference paper at ICLR 2024
TRAINING BAYESIAN NEURAL NETWORKS WITH SPARSE SUBSPACE VARIATIONAL INFERENCE
Junbo Li', Zichen Miao?, Qiang Qiu", Ruqi Zhang'
Department of Computer Science 2. School of Electrical and Computer Engincering Purdue University, West Lafayette, IN 47907, USA (1 jujbbo, mi aoz, qqiu, ruqizi z @purdue.
ABSTRACT
Bayesian neural networks (BNNs) offer uncertainty quantification but come with the downside of substantially increased training and inference costs. Sparse BNNs have been investigated for efficient inference, typically by either slowly intro- ducing sparsity throughout the training or by post-training compression of dense BNNs. The dilemma of how to cut down massive training costs remains, par- ticularly given the requirement to learn about the uncertainty. To solve this challenge, we introduce Sparse Subspace Variational Inference (SSVI) the first fully sparse BNN framework that maintains consistently highly sparse Bayesian model throughout the training and inference phases. Starting from a randomly initialized low-dimensional sparse subspace, our approach alternately optimizes the sparse subspace basis selection and its associated parameters. While basis selection is characterized as a con-differentiabee problem, we approximate the op- timal solution with removal-and-adee andiiiiiooon strategy, guided by novel criteria based on weight distribution statistics. Our extensive experiments show that SSVI sets new benchmarks in crafting sparse BNNs, achieving, for instance, 10-20x com- pression in model size with under 3% performance drop, and up to 20x FLOPs reduction during training compared with dense VI training. Remarkably, SSVI also demonstrates enhanced robustness to hyperparameters, reducing the need for intricate tuning in VI and occasionally even surpassing VI-trained dense BNNs both accuracy and uncertainty metrics.
INTRODUCTION
Bayesian neural networks (BNNs) MacKay 1992)
infuse deep models with probabilistic methods
through stochastic variational inference (Hoffman et al. 2013 Kingma al., 2015 Molchanov
et al. 2017
Dusenberry et al. 2020p or Monte-Carlo-based methods (Welling Teh 2011 Chen
al. 2014
Zhang et al. 2020e: Cobb & Jalaian 2021 Wenzel et al. 2020 Zhao et al. 2019).
facilitating a systematic approach to quantify uncertainties. Such a capacity to estimate uncertainty
becomes important in applications where reliable decision-making is crucial (Filos al.) 2019
Feng et al. 2018 Abdullah et al. 2022). Unlike traditional neural networks, which offer point estimates, BNNs place distributions on weights and use full posterior distribution to capture model uncertainty and diverse data representation. This has been shown to greatly improve generalization
accuracy and uncertainty estimation across wide variety of deep learning tasks (Zhang et 2020e
Blundell et al. 2015 Wilson & Izmailov 2020 Vadera et al. 2022a)
However, shifting from deterministic models to BNNs brings its own set of complexities. The com- putational demands and memory requisites of BNNs, mainly stemming from sampling and integra- tion over these weight distributions, are notably increased especially for large-scale neural networks Zhang et al., 2020c b a). In standard deterministic neural networks, it has been observed that neu- ral architectures can often be significantly condensed without greatly compromising performance
[ranklle & Carbin 2018 Evei et al. 2020). Drawing inspiration from this, some BNN approaches
have leveraged sparsity- promoting priors, such as log-unin -og-uriform (Kingma et al. 2015 Molchanov et al., 2017) or spike-and-slaa (Deng et al. 2019 [Wang et al.) 2021), to gradually sparsify dur- ing training, or implement pruning after obtaining standard dense BNNs based on signal-to-noie noiise
ratios (Graavs 2011 Blundell et al. 2015 Ritter al. 2021)

---

Published as a conference paper at ICLR 2024
As result, while they do mitigate the computational burden during inference, the substantial costs associated with training remain largely unaffected. Additionally, selecting optimal substructures (adera al. 2022bb may be challenging using the sho- pruning methods. In order to scale up BNNs to modern large-scale neural networks, a crucial question remains to be addressed:
How can we craft a fully sparse framework for both training and inference in BNNs?
In our pursuit of a solution, we depart from the conventional approach of designing complex sparse- promoting priors, instead, we advocate for a strategy that embeds sparsity into the posterior from the very beginning of training.
CIFAR-10
CIFAR-100
To this end, introduce Sparse Subspace Variational Inference (SSVI) framework, where we confine the parameters of the ap- proximate posterior to adaptive subspace, performing VI solely within this domain, which significantly reduces training overheads. A defining feature of SSVI is its learning strategy, where both the parameters of the posterior and the subspace are co-learned, em- ploying an alternating optimization technique. When updating the subspace, our method utilizes novel criteria that are constructed us- ing the weight distribution statistics, Empirical results show that SSVI establishes superior standards in efficiency, accuracy, and un- certainty measures compared to existing methods, as shown in Fig- ure Moreover, SSVI displays robustness against hyperparame- ter variations and consistently showcases enhanced stability during training, sometimes even outperforming VI-trained dense BNNs. We summarize our contributions as the following:
94.0
SSVI
SSVI
93.5
MeNN
75.0
BNN
93.0
72.5
92.5
67.5
91.5
91.0
90.5
62.5
58P
50P
00.0
30
10'
10'
10"
Traning CPs (x101))
Twiiiing F1OPs (x10%)
Figure 1: Flops v.s. Accuracy. Our SSVI achieves the best test accuracy with the lowest training FLOPs.
We propose SSVI, the first fully sparse BNN paradigm, reduce both training and testing costs while achieving comparable and even better performance compared dense BNNs.
We introduce novel criteria for evaluating the importance of weights in BNNs, leveraging this information to simultaneously optimize the sparse subspace and the posterior.
We set new benchmarks in performance metrics of BNNs, including efficiency, accuracy, and uncertainty estimation, achieving a 10-20x compression in model size with under 3% performance drop, and up 20x FLOPs reduction during training. Through exhaustive ablation studies, we further refine and highlight the optimal design choices of our SSVI framework. We release the code at https://ithub.com////222222225
1.1 RELATED WORK
Sparse BNNs Since the introduction of the Bayesian learning framework, extensive research has focused on deriving efficient posterior inference methods [Tipping 2001] Graves] 2011: Zhang et al.
2020c Svensson et al. 2016). The challenge amplifies when transitioning the domain of Bayesian neural networks (BNNs) MacKay 1992). To reduce computational costs, many studies have em- priors obtain sparse BNNs. Built upon variational inference, Molchanov inference.MMllchaovv et al. (2017) integrates variational dropout (Kingma et al.] 2013) with log-uniform prior, while Neklyudov et al. (201)) expands to node-sparsity. [ouizos et al. (017)) Ghosh et al. (2018) and Bai et al. (2020) further explore the use of half-Cauchy, Horseshoe and Gaussian- Dirac spike-and- slab priors respectively. Based on Markov chain Monte Carlo (MCMC), Deng et al. (201)); Wang et al. (202I) explore the use of Laplace-Gaustae ods perform incremental pruning during training and rely heavily on intricately designed priors. Another line of work focuses pruning fully trained dense BNNs (Graves) 2011 Blundell et al. 2013). thus the training costs still remain. The methodology most akin to ours is the recent work on masked BNN (mBNN) (Kong al. 2023), which uses similar birth-and-death mechanism for weight selection. However, mBNN requires a complex hierarchical sparse prior and uses MCMC to do posterior inference.
ployed pparse-promooiee
spike-and-slab prior. However, all these meth-
Fully Sparse DNNs pruning strategies due to the BNN training challenges, is common in standard neural networks to engage in fully sparse training, essentially initiating with sparse model. For instance, Rigl. (Evei
While current sparse BNNs mainly rely on post-training during training
2

---

Published as a conference paper at ICLR 2024
al. 2020p starts with sparse network, maintaining the desired sparsity throughout its training process by adjusting weights via dropping and growing. The criteria to determine such adjustments can derive from various factors, such as weight magnitude (Frankle & Carbin 2018 Evcı et al,
2020). gradient measures (Ström 2015. Evei et al. 2020p, and Hessian-based criteria (Hassibi &
Stork 1992 Singh & Alistarh 2020).
PRELIMINARY
Bayesian neural networks and variational inference We use the classical Bayesian neural net- work (BNN) formulation. Given a dataset X' = (X, Y), where X, Y. are the independent and dependent variables respectively, we want return the probability p(x[{)), where = (x, y) is new data point. Let 0 E IRd be the random variables of a Bayesian neural network, we have
N
p(x[2)) = p(x (((x(xx0))((xx)dd) =
N E p(x|XX g(o)), 8 ~ p(0)2)
where p() the posterior distribution. In this paper, we use variational inference (VT) (ordan et al. 1999) approximate the true posterior. Specifically, we find a distribution (s((), where are the distribution parameters, that minimizes (L){s(0)])p(0)) (O))))) VI typically adopts mean-field assumptions for both approximate posterior and prior p. which implies that (s(() = [ 96 (0;) and p(0) I[ p(O.). Thus, is equivalent to minimizing:
(L((s(0)|lp(0) - Egg log p(x (L((o(0.)lp(0.. Ey log p(
(I)
The expectation term is intractable and is usually approximated using Monte Carlo (MC) samples. In the backward process, the reparameterization trick (Kingma & Welling 2013) is applied to enable gradient backpropagation for Φ. We these standard assumptions and techniques our method.
SPARSE SUBSPACE VARIATIONAL INFERENCE
In this section, present framework, sparse subspace variational inference (SSVI). First, introduce the formulation and introduce an overview of the optimization in Section 3.1 Then we present the details of the alternative optimization strategy in Section 3.3.1 and 3.3.2
3.1 SSVI FORMULATION
In conventional VI settings, the parameter vector ó generally has a dimensionality that is at least double that of 6, accounting for both the mean and the variance. Consequently, there strong desire to seek sparse solutions to reduce the complexities arising from the augmented parameter space Tipping 200Tp. This holds particular significance when dealing with deep neural networks, where the model complexity is more pronounced (Molchanov et al. 2017 Neklyudov et al. 2017).
To address this, we introduce SSVI, wherein we confine the support qo to a subspace of
pre-assigned dimension s within Rdd, implying (0) = for 0 not belonging to S. To achieve sparse model, we consider We formulate SSVI as the following optimization problem
be subspace expanded directly from a subset of the axes of Rd.
min $.1
KL (00(0)(()) Egg log ((x(0)
I= I (n.....,,..) n,). (0,) 5(0,).
In this formulation, S is extended by subset of the axes of IRdd through the basis vectors {er, 1 Here, (0,)) refers to the marginal density function of 0, according to 9s (f), and 5(.) stands for the Dirac delta function. This means for { ras). the approximate posterior distribution of 0; shrinks to zero. For simplicity, we consider a Gaussian distribution for ds where Φ E dxdx
with o, (M, of) representing the mean and variance for < < Other mean-field variational distributions also be combined with our method. This suggests that, for i (nI.... .ns), di
3

---

Published as a conference paper at ICLR 2024
(0,0), 0), resulting in a sparse solution. Equivalently, write 2 the following optimization problem. Here @ means element-wise multiplication. The constrant conditions mean that when % is O, μ4 and o; also O.
min de bc (0
% K (go Eg log p(210)
(3)
s.t.
1 s, P= o² o² 0ny
The formulation 3 has three key advantages compared to previous sparse BNN methods (Kingma et al. 2013| Molchanov al. 2017 Neklyudov et al. 2017 Kong et al. 2023). which we discuss
in detail separately below.
Reduce training costs with constant sparsity rate Previous methods in both VI and MCMC- based sparse BNNs typically start with a dense model and traverse the full parameter space to eventually obtain sparse solutions on specific criteria. contrast, our SSVI approach optimizes parameters directly within a sparse subspace. This ensures a constant sparsity level of s/d throughout training, thereby considerably diminishing both computational and memory costs during training.
process finalized through post-training truncations based
Guarantee to achieve target sparsity levels Existing approaches in both VI and MCMC-based sparse BNNs depend on unpredictable sparsity regularization from the prior and post-hoe truncation based on threshold values. This leads to indeterminable sparsity levels, often requiring extensive tuning of the prior to achieve desired levels of sparsity. In contrast, our method provides a straight- forward means of controlling sparsity through the hyperparameter s. By constraining the search space to only subspaces with a dimension of 8, any solution pair (ø, 7) of guarantees the desired level of sparsity.
Allow the of Gaussian distribution for both prior and variational posterior Our formula- tion does not require an intricated design for either the prior or the posterior. Instead, we can simply adopt commonly used variational distributions, such Gaussian distributions, for both the prior and posterior distributions. Furthermore, this simplifies the computation by enabling the direct calcu- lation of the KL divergence term without necessitating complicated and expensive approximations, which are often required when using sparse-promoliie priors.
3.1.1 ALTERNATIVE OPTIMIZATION PROCEDURE
The optimization problem in 3 includes both á and 7. Denote 1 {i 7: = 1}. At the beginning. we initialize y by randomly selecting a subset 7 = {mI,.... n.. from {1.....d d) where indices Y are set 1. We also initialize 6 1 following techniques commonly used in VI-based BNN training. Given the discrete nature of 7, directly applying gradient descent for both variables is infeasible. To address this challenge, we propose to alternatively update and y at each iteration. The optimization for ø can be done via standard VI-based BNN training, leveraging gradient infor- mation. However, updating y represents a more sophisticated task. We first present our algorithm in Algorithm I and we will delve into the specifics of the design of the updating processes in the forthcoming sections.
3.2 UPDATE THE VARIATIONAL PARAMETER
Given a fixed the update for oi is simple and can be achieved using optimization methods like SGD. We obtain g++ from (0t,7) through M gradient descent steps, and each step is similar to the standard VI update.
During the forward process, we estimate Ex,, log ((2) using Monte Carlo samples from (0). For improved accuracy, distinct 0 samples are drawn for each input in batch. However, naive methods, which require a forward pass per sample, are inefficient. We thus employ the local repa- rameterization trick (LRT) (Kingma et al.] 2015 Molchanov al. 2017) to improve efficiency, achieving equivalent outputs in just two forward processes. Specifically, y, for a linear operation with inputs R and Gaussian matrix A E IRxx where B is batch size, p, q are the input and

---

Published as a conference paper at ICLR 2024
Algorithm Sparse Subspace Variational Inference (SSVI)
Require: BNN 0 IRd with prior p(0), variational distribution (o()), target sparsity s/d, re- placement rate {r,), inner update steps M, total steps T.
1: Random initialize (Ø", 77)) from the feasible set and set y =
2: for = ., do
3: # Update Φ
4: 5t.0 Initialize(s "',""'''"" ("
according to Section 3.2
5:
for m 0, O,.... M 1 do
6: Obtain @f using the gradient of end for
7:
8:
o++1 6r.
9:
# Update y.
11: Treppe Removal(y" (n,, @t+1 rt) according to Section 3.3.1
2f
Addition (eriiion(/hoommoo
Tremove: oft1, rr)) according to Section 3.3.2
12: end for
output dimensions, j and a define A's mean and variance, the output y is computed by
y F V(a (x
where e is sampled from standard Gaussian noise matching y's dimension. This method mirrors sampling distinct weights for each input z,,6 via just two forwards, rather than times. While the forward process matches the naive approach, LRT's backward process for differs and influences the SSVI design. Specifically, the gradient for with respect to the received loss becomes:
01
(
(x
do Oy (V
(5)
Please refer Appendix A for detailed discussion.
Establishing initial values for øz+1 updates We find that 5 exhibits linear dependency on σ. In the context of SSVI, the sparsity ratio of remains constant throughout the training process. (referring to the indices freshly updated to be non-zero), the value of of is zero. This implies that the straightforward application of gradient descent is infeasible the gradient persistently equals zero for of. Therefore, is crucial to design a strategy to set the initial value for of. different from the conventional sparse training in standard DNNs which does not require such initial values.
Consequently, for all /. 2 and satisfying 75 1-
We consider two potential strategies. The initial approach sets of to a very small number close to zero, which helps stabilize training while avoiding the zero gradient issue. Alternatively, of can be set to the average value of the non-zero σk. derived from the same module that is, the identical convolutional kernel or the same fully connected layer as of. Favorably, the latter strategy does not demand additional hyperparameters and has demonstrated superior stability in experimental settings, thereby being our method's s default choice.
3.3 UPDATE THE SPARSE SUBSPACE 7
The set of non-zero elements outlines the architecture of the subnetwork. Since y takes discrete values, it cannot be updated using regular gradient-basse based methods. Based on the existing (ø--1 7²), we propose to obtain (+1 through process selective removal and addition of indices in 7 This adjustment process takes inspiration from traditional sparse training techniques seen in regular, Bayesian neural network training (Frankle Carbin 2018) Evei et al. 20200. Notably, we introduce several novel techniques that especially compatible with BNN training
3.3. SUBSPACE REMOVAL
To derive g we begin by identifying and eliminating less significant indices in 7¹ based on di+1 (1t+1, ((+1)) E dR2x2 In standard DNN sparse training, common approach is to prune

---

Published as a conference paper at ICLR 2024
weights with small magnitudes, under the assumption that small weights tend to have negligible impact on the output. Building upon this insight, we derive several novel criteria for removing indices in the subspace. We omit the step t in the following.
|el| straightforward adaptation for BNNs is remove indices using the absolute values of the mean, i.e., [p], to set "/ with the smallest |py| to be zero. However, this strategy overlooks the critical information about uncertainty encoded o².
Criteria 1:
Criteria 2: $ (0) = |u|/o Another commonly used metric to decide the importance of weights in BNN is the Signil-to-Noise Ratio (SNR) (Kingma et al., 2015 Neklyudov et al. 2017), Le.11u//a. This metric considers both the magnitude and variance of the weights. Though previous works have utilized SNR for post-training pruning (Kingma al. 2015) Neklyudov et al. 2017
Louizos et al. 2017)
to the best of our knowledge, its application during training is novel.
Criteria 3: Egg, |0) We introduce a new criterion that considers both the mean and the uncertainty. Criteria |u|| essentially equates to gg However, in BNN, where weights are randomly sampled before each forward pass with given imputs, is more appropriate to use Eg, The reasoning behind this is that weights with smaller absolute values, on average, have lesser influence on the outputs. Notably, this criterion can be explicitly calculated when qs follows a Gaussian distribution. Specifically, for i-th dimension of weight O, with mean μt and variance o7 we have
Egg [0:1 = pi
2Φ
Os -1]
-1)
1
+
20, exp
1 20
where ((x)
1
exp
y²
dy
2x
V 2TT
2
(6)
The full derivation can be found in Appendix This criterion differs significantly from Criteria by expressing a theoretically-backe represented by |a), with the uncertainty dimensions from Criteria 2, denoted by lal/o. We highlight that the role of p/ar is deliberately bounded in this formula, either by the cumulative distribution function (-) or by the inverse of the exponential function 1/ cxp(()). Given that the uncertainty measure a can be challenging to estimate accurately during initial training phases, this imposed boundary on p/aa prevents from excessively skewing weight prioritization, thus helping to ad- dress potential issues of model instability. Besides the constrained behavior of p/a, the criteria prominently lean towards 14 because of its linear relationship. However, a concerning aspect high o; typically indicates weight's
new rule that combines the magnitude insights from Criteria 1,
the near-linear dependency o, This is problematic insignificance. To counteract this, we propose subsequent criteria that introduce regularization.
Criteria 4: SNR (10)) Although Egg [8] already contains uncertainty information, we can fur- ther consider the SNR of [0] as opposed to in Criteria 2. For 0;, we have
Varg. go (10,1) Eg 10:12 (Eg. a. 0, 0,1)² 0f (E 10.1) o² + 2 (Ego 10:1)
Hence, the SNR is
H4g
2Φ
1)
+
20 2
exp
a 20
SNR ((0a))
(7)
0²? + /4
2Φ Mi
20 22x
cxp
2o
Given that the denominator of progresses linearly with respect to a, it counteracts the unfavorable characteristic seen in 5 Additionally, y retains its prominent role as the denominator behaves sub- linearly with μ. whereas the numerator exhibits a linear relationship.
Criteria 5 & 6: Egg e\ and SNRG0 (E-101) When examining weight importance in the net- work, one could go beyond the straightforward metrics like the expectation and SNR of ( and delve into more general functions f. These functions should be monotonic within the range [0, +oo) to ensure that 9 with greater absolute values always has high importance. As practical example, we adopt the function f((Ø) exp(2)0)), where À is positive hyperparameter. This choice results

---

Published as a conference paper at ICLR 2024
Table Comparisons with dense BNNs. Our method achieves 10-20x savings in both training costs and model complexity via keeping high-sparsity subspaces throughout the training. Here we show the results with two different sparsity levels as inputs of SSVI.
Dataset
Model
Accuracy
BCE
Sparsity
Training FLOPs
Dense
0%
Ix(1.1 10")
CIFAR-10
w/ SSVI
94.58
0.005
90%
0.10x
w/
93.69
Donse BNN
78.20
0.0016
0%
1x(1 10"T)
CIFAR-100
W/ SSVI
75.81
0.0017
90%
0.10x
SSVI
73.83
95%
0,05x
in an explicit calculation of both the expectation and SNR of f((0)) under a Gaussian distribution
without necessitating further approximations. For À > O, the derived outcomes are outlined below:
Egg 101
1 ( Ao
2²2² tAp $ (- " a
\o
2²/² AAp
SNR,g (e~181)=
(8)
x20²
x2,2 -Ap
1 ( + Aor)
+hx + 2 (- +x0)
(5 + 22o) 62A268 + 2 (- + 2Ar) c2AR02 -22
@ ( + Ar)) 22 +Ap (- +
23,²
Criteria Selection Criteria SNR((0]) has a compelling theoretical explanation, and our empiri- cal results in Section support its superior and stable performance during training compared to other criteria. Consequently, we choose as the default for SSVI.
3.3.2 SUBSPACE ADDITION
After pruning spaces using one of the proposed criteria, it is essential to reintroduce some of the removed spaces back to y. Drawing inspiration from traditional sparse training in DNNs, we select important weights by their gradient's absolute magnitude. Specifically, for BNN training with
stochastic batch, denoting the loss function in Bb (z,) where 0 is sample for stochastic batch r with batch size B, the selection criterion is given by:
B
1
Eng
Ex
(af),
B
==1
To compute the two expectations, we propose three criteria, using Monte Carlo approximation:
1. Apply a one-step Monte-Carlo for (6, z): B E.
2. Apply one-step Monte-Carlo for with 0 set to μi
B Dl Volur(a.))
3. Apply multi-step Monte-Carlo for (0, K Σ
Volan
Ablation studies Section g indicate that Criteria 1, while being the most straightforwardd delivers performance on par with the other two. Thus, we default to Criteria 1 for SSVI.
4 EXPERIMENTS
4.1
MAIN RESULTS
We conduct experiments using CIFAR-10 10 and CIFAR-100 datasets (Krizhessyy
2009) and ResNet- KL warm-up (Molchanov et al. 2017; Sønderby et al. 2016). a technique widely applied in VI-based BNNs. Specifically, auxiliary coefficient introduced for the KL. term in Eq. I gradually increasing
18 (He et al. 2016) as the backbone networks. To facilitate stable training,
7

---

Published as a conference paper at ICLR 2024
Table Comparison with previous sparse BNNs on CIFAR-IO and CIFAR-100. SSVI achieves the
best accuracy and uncertainty quantification with the lowest computation and memory costs.
Dataset
Algurithm
Accuracy
Loss
ECE
Sparsity
Complexity
Training FL.Pss (x:10
VD
Louizos et al.
2017
90.6
0.298
0.009
76.8%
2.60M
67.1
SBP
Ivokivady et
2017
90.1
0.301
0.009
81.1%
2.12M
67.1
CIFAR-IO
AET
91.84
0.308
0.023
90.0%
1.2M
383.4
mBNN
Kong ct al 0055
93.14
0.237
0.008
96.4%
0.41M
22.8
SSV1
93.74
0,218
0
95.0%
0.56M
5.7
VD> Lauizes
2017
60.0
0.0060
58,
4.6M
SBP
Nokivadoy & el 201
62.3
1,834
0.0053
61 21.2%
4.34M
67.1
CIFAR-100
ALB Deng et alg 2019
70.26
1,103
0.0110
90%
1.2M
383.4
mBNN poing SSVI
73.42
1,078
0.0024
87.8%
1.37M
25.9
75.81
0.0017
1.12M
from 9max during training. More training details are shown in Appendix C For inference, we
consider accuracy, loss, and expected calibration error (ECE) (Guo et al. 2017) as the metrics, which are obtained by five samples from the posterior distribution following prior work KKng et al. (2023) We also do evaluations on distribution shift benchmarks in Appendix D.I. For all the experiments, we train with three different seeds and average the results.
We first demonstrate the results of the reduction in training costs and model complexity compared with dense BNNs in Table T SSVI achieves comparable test accuracy and ECE with drastically reduced model size and training costs. For example, SSVI achieves 20x model size compression and FLOPs reduction with under 2% accuracy drop and 0.002 ECE increase on CIFAR- 100 and 10x model size compression and FLOPs reduction with under 3% accuracy drop and 0.0001 ECE in- crease on CIFAR- 100. We also compare our method with previous sparse BNN methods in Table The results show that SSVI significantly outperforms all previous works across all metries on both datasets, achieving the best accuracy and uncertainty quantification with the lowest computation and memory costs. It strongly indicates that our method, optimizing both the subspace and variational parameters, can lead to better sparse solutions than using complicated sparse ppasse-poomoting priors.
4.2 ANALYSIS AND ABLATION STUDIES
In this section, we perform a thorough analysis of SSVI's properties and present ablation studies focusing on its core designs. Here we train ResNet-18 on the CIFAR- 100 dataset, which is more challenging than CIFAR- 10, and can help distinguish different algorithm designs.
Flexible target sparsity Different from previous works (Molchanov et al., 2017; Kong et al.) 2023), SSVI can flexibly assign the target sparsity before training. In Figure 2 we plot the results on broad range of sparsity from 50% to 1%. The plots show that our SSVI is robust to different sparsity levels with minimal accuracy drop. For a clearer perspective, we contrast our results with RigL. (Evci al. 2020), method closely aligned with our approach in the domain of sparse training for traditional DNNs. While both methods employ similar subspace training techniques, our criteria tailored for BNNs. The results indicate that while the performance trajectory of our BNNs in relation to sparsity closely mirrors that of RigL, our approach is more robust to varying sparsity levels, as evidenced by lesser decline in accuracy.
SSVI's edge over VI While SSVI inherently optimizes parameters within a subspace, making it more efficient than standard VI, it even occasionally outperforms VI. This superior performance is largely due to its enhanced resilience to hyperparameter variations. To illustrate, consider VI with Gaussian distributions where the parameters are represented as o = (M, o). The optimization process requires an initialization 6d = ((4",")) where o is typically sampled from a Gaussian with minimal mean and variance. Ensuring small σ important for preventing training failures in the initial phase. Our results further verify the instability of BNN training with VI to initial values, highlighting the need for precise hyperparameter tuning. In contrast, BNNs trained with SSVI demonstrate increased tolerance to deviations in the initial configurations, thus simplifying the intricacies associated with hyperparameter optimization.
Our findings presented in Figure 2 We use SSVI with sparsity 0.1 with different removal crite- ria. By adjusting the initialization mean of o from 0.001 to 0.01 while keeping all other settings unchanged for both VI and SSVI, we clearly mance deteriorates immediately. In contrast, SSVI maintains robustness across all different dropping criteria and even outperforms VI significantly reduced training expense.
that the mean marginally increases, VI's perfor-
8

---

Published as a conference paper at ICLR 2024
SSVI RigL with varying densities
SSVI with different dropping criteria
Vi-10
SSV1-0.1 jil
55y1-0.1 pile
Accuracy
74
Accuracy
5891 0.1 810)
76
SSVi-0.1 zul
66
SSV
1.
0.008 Mean of initialization
Density
Figure 2: Left: the performance of SSVI for BNNs compared to Rigl. (Evei et al. 2020 for standard NNs across various sparsity levels. SSVI demonstrates superior robustness to different sparsity lev- els. Right: the results for VI at full density and SSVI at 0.1 density using different removal criteria. We see that VI is heavily dependent warmup, whereas SSVI offers more consistent training across all drop criteria without additional tricks, and outperforms VI at initialization value 0.01. Notably, among all criteria, SNR((6)) stands out, leveraging uncertainty information in a theoretically backed manner.
precise hyperparameter tuning, especially concerning KL
SSVI with different MC stops
SSVI with d eifferent initializotion
70
24
72
1-step
50
mean
5-step
10
68
20 stop MC
10-5
0.5
1.5
Sparsity
Sparsity
Figure 3: Left: ablation results concerning the number of MC steps, indicating minimal differences between them. Right: ablations related to initialization, highlighting the advantage of initializing with the mean.
Ablation study on removal and addition criteria In Section 3.3.1 we introduce several criteria tailored for BNNs using weight distribution statistics. Figure compares accuracy for these criteria. We see that SNR (10) achieves the best accuracy on all sparsity levels, aligning with our analysis in Section B.3.I We also calculate the Intersection over Union (IoU) among the weights removed based on each criteria pair. See Appendix D.2 for details. Figure 3 shows the results using different MC steps, highlighting that multi-step MC is not better than the 1-step MC. Thus, advocate for the one-step MC due to its superiority in both efficacy and performance.
Ablation study on updating a In Section 3.2 we highlighted the necessity for a specialized & before its update using gradient descent, as
design in initializing the variance parameters opposed to starting with zero, as driven by the gradient outlined in ( Two strategies were proposed: one starts with small value close zero, and the other uses the average of the remaining non-zero values. The comparison results of these strategies, while holding other factors the same, are showcased in Figure g the optimal hyperparameter setting for the first method appears to closely mirror the results when ø is initialized using the mean of non-zero values. Therefore, we adopt the mean- mean-based initialization approach, which also eliminates the necessity for extra hyperparameters.
CONCLUSION
In this paper, we propose Sparse Subspace Variational Inference (SSVI), the first fully sparse frame- work for both training and inference in Bayesian neural networks (BNNs). Specifically, SSVI op- timizes the sparse subspace and variational parameters alternately, employing novel criteria based on weight distribution statistics for the removal and addition within the subspace. We conduct com- prehensive experiments to demonstrate that SSVI establishes new benchmarks in sparse BNNs, by significantly outperforming previous methods in terms of both performance and efficiency. We also perform in-depth analysis to explore SSVI's properties and identify optimal design choices for cri- teria. As look ahead, our future endeavors will focus refining the optimization process. This could involve designing new weighting functions for the SNR removal criteria and exploring new addition criteria that utilize higher-order information of the weight distribution.

---

Published as a conference paper at ICLR 2024
REFERENCES
Abdullah A Abdullah, Masoud M Hassan, and Yaseen T Mustafa. A review on bayesian deep learning in healthcare: Applications and challenges. IEEE Access, 10:36538-36528 2022.
Jincheng Bai, Qifan Song, and Guang Cheng. Efficient variational inference for sparse deep learning with theoretical guarantee. Advances in Neural Information Processing Systems, 33:466-476, 2020.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International conference on machine learning, pp. 1613-1622. PMLR, 2015.
Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In International conference on machine learning, pp. 1683-1691. PMLR, 2014.
Adam D Cobb and Brian Jalaian. Scaling hamiltonian monte carlo inference for bayesian neural networks with symmetric splitting. In Uncertainty in Artificial Intelligence, pp. 675-685, PMLR, 2021.
Wei Deng, Xiao Zhang, Faming Liang, and Guang Lin. An adaptive empirical bayesian method for sparse deep learning, Advances in neural information processing systems, 32, 2019.
Michael Dusenberry, Ghassen Jerfel, Yeming Wen, Yian Ma, Jasper Snoek, Katherine Heller, Balaji Lakshminarayanan, and Dustin Tran. Efficient and scalable bayesian neural nets with rank-1 factors. International conference on machine learning, pp. 2782-2792. PMLR, 2020.
Utku Evci, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen. Rigging the lottery: Making all tickets winners. In International Conference on Machine Learning, pp. 2943-2952. PMLR, 2020.
Di Feng, Lars Rosenbaum, and Klaus Dietmayer. Towards safe autonomous driving: Capture un- certainty in the deep neural network for lidar 3d vehicle detection. In 2018 21st international conference on intelligent transportation systems (ITSC), pp. 3266-322 3273. IEEE, 2018
Angelos Filos, Sebastian Farquhar, Aidan Gomez, Tim GJ Rudner, Zachary Kenton, Lewis Smith, Milad Alizadeh, Arnoud De Kroon, and Yarin Gal. systematic comparison of bayesian deep learning robustness in diabetic retinopathy tasks. arXiv preprint 11X1v:1912.10481, 2019.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint 1803:0880003000 2018.
Soumya Ghosh, Jiayu Yao, and Finale Doshi-Velez. Structured variational learning of bayesian neural networks with horseshoe priors. In International Conference on Machine Learning, pp- PMLR, 2018
Alex Graves. Practical variational inference for neural networks. Advances neural information processing systems, 24, 1.11.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In International conference on machine learning, pp. 1321-1330. PMLR, 2017.
Babak Hassibi and David Stork. Second order derivatives for network pruning: Optimal brain surgeon. Advances in neural information processing systems, 5, 1992.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common cor- ruptions and perturbations. arXiv preprint arXiv: 190091223.,,, 2019.
Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic variational infer- ence. Journal of Machine Learning Research, 2013.
10

---

Published as a conference paper at ICLR 2024
Michael Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning, 37:183-2233 233,, 1999.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes.
arXiv preprint
arXiv: 14,611111100111 2013.
Durk P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameteri- zation trick. Advances in neural information processing systems, 28, 2015.
Insung Kong, Dongyoon Yang, Jongjin Lee, Ilsang Ohn, Gyuseung Back, and Yongdai Kim. Masked bayesian neural networks: Theoretical guarantee and its posterior inference. arXiv preprint arXiv:2305. 14765, 2023.
A Krizhevsky. Learning multiple layers of features from tiny images. Master's thesis, University of Tront, 2009.
Christos Louizos, Karen Ullrich, and Max Welling, Bayesian compression for deep learning. Ad- vances neural information processing systems, 30, 2017.
David JC MacKay, A practical bayesian framework for backpropagation networks. Neural compu- tation, 4(3):448-477 1992.
Dmitry Molchanov, Arsenii Ashukha, and Dmitry Vetrov. Variational dropout sparsifies deep neural networks. In International Conference on Machine Learning, pp. 2498 2507. PMLR, 2017.
Kirill Neklyudov, Dmitry Molchanov, Arsenii Ashukha, and Dmitry P Vetrov. Structured bayesian pruning via log-normal multiplicative noise. Advances in Neural Information Processing Systems, 30, 2017.
Hippolyt Ritter, Martin Kukla, Cheng Zhang, and Yingzhen Li. Sparse uncertainty representation in deep learning with inducing weights. Advances in Neural Information Processing Systems, 34: 6515-6528, 2021.
Sidak Pal Singh and Dan Alistarh. Woodfisher: Efficient second-order approximation for neural network compression. Advances in Neural Information Processing Systems, 13:18098-18109, 2020.
Casper Kaae Sønderby, Tapani Raiko, Lars Maaløe, Søren Kaae Sønderby, and Ole Winther. How train deep variational autoencoders and probabilistic ladder networks. arXiv preprini arXiv:1602.02282.. 3(2), 2016.
Nikko Ström. Scalable distributed dnn training using commodity gpu cloud computing. 2015.
Andreas Svensson, Arno Solin, Simo Särkkä, and Thomas Schön. Computationally efficient bayesian learning of gaussian process state space models. In Artificial Intelligence and Statis- fics, pp. 213-221, PMLR, 2016.
Michael E Tipping. Sparse bayesian learning and the relevance vector machine. Journal of machine learning research, ((lmm):211-24:: 2001.
Meet Vadera, Jinyang Li, Adam Cobb, Brian Jalaian, Tarek Abdelzaher, and Benjamin Marlin. Ursabench: system for comprehensive benchmarking of bayesian deep neural network models and inference methods. Proceedings of Machine Learning and Systems, 4:217-237, 2022a.
Meet Vadera, Adam D Cobb, Brian Jalaian, and Benjamin M Marlin. Impact of parame- ter sparsity on stochastic gradient meme methods for bayesian deep learning. arXiv preprint 2.07022202.03770, 2022b.
Yating Wang, Wei Deng, and Guang Lin. Bayesian sparse learning with preconditioned stochastic gradient meme and its applications. Journal of Computational Physics, 432:1 10134, 2021.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-I1), pp. 681-688, 2011.
11

---

Published as a conference paper at ICLR 2024
Florian Wenzel, Kevin Roth, Bastiaan S Veeling, Jakub Swiętkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How good the bayes posterior in deep neural networks really? arXiv preprint htth:2002.022003. 2020.
Andrew G Wilson and Pavel Izmailov. Bayesian deep learning and probabilistic perspective of generalization. Advances in neural information processing systems, 33:4697-44708, 2020.
Jianyi Zhang, Ruiyi Zhang, Lawrence Carin, and Changyou Chen. Stochastic particle-opimization sampling and the non-asymptotic convergence theory. In International Conference on Artificial Intelligence and Statistics, pp. 1877-1887. PMLR, 2020a.
Jianyi Zhang, Yang Zhao, and Changyou Chen.
Variance reduction in stochastic particle- optimization sampling. In International Conference on Machine Learning, pp. 11307-11316. PMLR, 2020b.
Ruqi Zhang, Feder Cooper, and Christopher De Sa. Amagold: Amortized metropolis adjustment for efficient stochastic gradient meme. In International Conference on Artificial Intelligence and Statistics, pp. 2142-2131 PMLR, 2020c.
Ruqi Zhang, A Feder Cooper, and Christopher De Sa. Asymptotically optimal exact minibatch metropolis-hastinge Advances in Neural Information Processing Systems, 33:19500-19510, 2020d.
Ruqi Zhang, Chunyuan Li, Jianyi Zhang, Changyou Chen, and Andrew Gordon Wilson. Cycli- cal stochastic gradient meme for bayesian deep learning. International Conference on Learning Representations, 2020e.
Yang Zhao, Jianyi Zhang, and Changyou Chen. Self-adversarially learned bayesian sampling. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 5893-5900, 2019.
12


Reviews:
[{"review_id": "aeUDN8F2Cx", "reviewer": "Reviewer_DFxu", "summary": "In this paper, the authors present an approach to train sparse-BNNs using sparse subspace variational inference. The authors show that their approach gives computation gains during training as well as during test-time inference.", "questions": "- Is eq 2 being really used somewhere directly? \n\n- What's the model architecture used in experiments? \n\n- Apart from implementation, can the authors highlight the diff between their work and that of Vadera et al.? \n\n- How would you extend this to other approximate Bayesian inference techniques, apart from mean-field VI? \n\n- It'll be useful if the authors can demonstrate the quality of uncertainty metrics coming out of the BNNs in their approach on downstream tasks. See [3] for example.\n\nReferences\n\n[3] Vadera, M., Li, J., Cobb, A., Jalaian, B., Abdelzaher, T., & Marlin, B. (2022). URSABench: A system for comprehensive benchmarking of Bayesian deep neural network models and inference methods. Proceedings of Machine Learning and Systems, 4, 217-237.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper is well-motivated. Using BNNs in practice naively requires lot of compute which can hamper its adoption \n\n- The experimental results provided by the paper are impressive. Having training time benefits is extremely useful. \n\n- I also appreciate the fact that the authors have provided the implementation for their work.", "weaknesses": "- I think related relevant work is missing [1, 2]. Especially the work by Vadera et al.[1], looks at sparse SGHMC for BNNs, which could be easily extended to VI and when doing so would look similar to what the authors have proposed. The authors should highlight the difference between their work and existing work. \n\n- The empirical results seem insufficient. The authors haven't mentioned the model architecture used in their experiments. Also, to really emphasize the usefulness of the proposed approach, it would be important to have empirical results on an expanded set of datasets + models + tasks (including UQ tasks such as OOD detection, misclassification detection, etc.)\n\n- I think the presentation can be improved. The authors introduce complex notations, but after reading through section 3 twice, I believe that it can be greatly simplified. For e.g., where exactly is eq 2 reconciled in the algorithm? Based on my reading, the paper would have been fine to totally exclude eq 2.  \n\n\nReferences\n\n[1] Vadera, M. P., Cobb, A. D., Jalaian, B., & Marlin, B. M. (2022). Impact of Parameter Sparsity on Stochastic Gradient MCMC Methods for Bayesian Deep Learning. arXiv preprint arXiv:2202.03770.\n\n[2] Ghosh, S., Yao, J., & Doshi-Velez, F. (2018, July). Structured variational learning of Bayesian neural networks with horseshoe priors. In International Conference on Machine Learning (pp. 1744-1753). PMLR.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ruh0bwLm67", "reviewer": "Reviewer_KZuX", "summary": "The paper proposes a novel inference method for sparse Bayesian neural networks based on variational inference, namely Sparse Subspace Variational Inference (SSVI). The idea, that is different to the previous sparse BNNs, is to explicitly keep the number of non-zero weights in the neural network fixed while updating which weights are included or excluded as non-significant. This is in contrast to the existing literature, where sparsity is induced by special priors and, firstly, still starts with a dense model that is sparsified during training, which still make inference during training computationally expensive. Secondly, sparsity-inducing priors do not allow to set a specific level of sparsity, which may require a user to play with hyperparameters of the prior to get the desired level of sparsity. The proposed method has a level of sparsity as a hyperparameter and it is fixed from the start of the training, saving computational efforts from the start. Experiments in CIFAR-10/100 show the superiority of the proposed method in comparison to other sparse NNs in accuracy, loss, ECE, and training FLOPs.", "questions": "I like the paper, my minor weakness findings are listed above, but there is nothing the authors should or could address during the rebuttal.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "* Interesting new idea for sparse BNNs, that:\n* ... is shown to be working empirically\n* ... has a solid theoretical model that is intuitive\n* ... brings significant benefits in terms of computational cost\n\nOriginality: The idea presented in the paper appears to be novel. \n\nQuality: The proposed method is well defined and all the inference steps seem to be sound. The experiments are well designed and executed. Ablation study is presented to different aspects of the method. \n\nClarity: The paper is very well written and easy to follow. \n\nSignificance: I believe the paper is of extreme interest for the Bayesian deep learning community as it addresses the problem of scaling inference of BNNs that is a known issue of the concept-appealing approach.", "weaknesses": "All of the below is not major weakness points.\n\nOriginality: A bit of context in terms of BNNs is missing. The review of the existing methods solely (though understandably) focusses on sparse models, leaving behind other approaches of making efficient Bayesian inference for NNs. E.g., Cobb, A.D. and Jalaian, B., 2021. Scaling Hamiltonian Monte Carlo inference for Bayesian neural networks with symmetric splitting. In Uncertainty in Artificial Intelligence (pp. 675-685). PMLR.\n\nQuality: The empirical evaluation of the model though done rather thoroughly but only on 2 medium-size/complexity datasets from the same domain on 1 architecture. It would be interesting to see more experiments. \n\nClarity: Though the model is mostly well-written, there are some moments that are missing. See details below. \n\nSpecific comments:\n1.\tAbstract. 20x compression, 20x FLOPs reduction – in comparison to what?\n2.\tAbstract. “surpassing VI-trained dense BNNs” – in terms of what?\n3.\tFigure 1. Not defined names of baselines.\n4.\tEq. (3) it is better to include explicitly what is the operator in the last 2 equations\n5.\tAround eq. (4). What are p, B, q, W?\n6.\tEq. (5). What is l?\n7.\tTable 1. What is the difference between row 2 and 3 and then 5 and 6?\n8.\tIt is unclear about eq. (11) and (12) (from the appendix), which one is used at the end?\n\nMinor:\n1.\tAblation study on updating \\phi. ‘… in Figure 3. the optimal …’ -> The", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TcQUTNLYl0", "reviewer": "Reviewer_SCNX", "summary": "The paper proposes a procedure for training sparse Bayesian neural networks with a factorized variational distribution. It does so by alternating between optimizing the parameters of a subset of the weight distributions, with the remaining ones zeroed out, and updating the subset of non-zero weights by adding to and removing from that subset based on a range of criteria. The experiments show improved accuracy and calibration on CIFAR10/100 over methods from the literature at higher sparsity.\n\nOverall, I would say that the paper introduces some reasonable but not ground-breaking technical material, however it is let down by a severely limited evaluation, omitting most standard uncertainty estimation benchmarks. With some further issues around overclaiming novelty, I would argue for rejecting the paper in its current form.\n\nPOST REBUTTAL UPDATE:\nThank you for the extensive updates. These largely address my core concerns (limited evaluation, lack of demonstration of difference between pruning scores), so I will raise my score. I would suggest at least referencing Appendix D from the main text and if possible moving some figures to the experiment section.\n\nI did not see an update to the paper reflecting W4. I completely agree that the method has limitations, my point is only that it is not a post-hoc pruning approach and the paper shouldn't misrepresent this.", "questions": "- how are the mean values initialized for weights that are added to the active set? I assume 0? Do you reuse the old value if a weight had been temporarily removed?\n- I'm not sure I follow the discussion of Figure 2 right. For the initialization of the variance of a mean-field Gaussian, identical considerations apply as for the initialization of weights in deterministic networks, so loosely speaking we want to scale the sum of variance of the means and the initial value of the variance parameter inversely with the number of units to avoid the activations diverging with increasing depth. So to me it seems natural, that dense VI would reach this critical threshold before a pruned variance, as the latter is simply removing terms from a positive sum and thus decreasing variance. Am I missing something/misinterpreting the figure?\n- Can you demonstrate any practical benefits from sparsity during training/the FLOP reduction translating to real time speedups?\n\nFor suggestions see weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- the paper is overall clear in what it does and why. It shouldn't be too difficult to implement the method based on the description.\n- selection criteria for weight pruning may be of independent interest.\n- there is a decent range of ablation studies on various aspects of the proposed method.\n- performance on the datasets and metrics under consideration is better than that of the baselines.", "weaknesses": "1. the experiments are quite bare-bones for a BNN paper, there is no evaluation of predictive uncertainty besides calibration -- we don't need a Bayesian approach do well on this metric. I would either suggest adding e.g. a temperature scaling baseline applied to a sparse deterministic net or (preferably) the usual out-of-distribution and distribution shift benchmarks.\n2. primarily testing at a single sparsity level as in Table 2 also seems a bit limited to me. In my view, there are broadly two possible goals when using sparsity: opitimizing sparsity at a given performance level, e.g. close to optimal, or optimizing performance at a given sparsity level. I would have liked to see more figures in the style of Figure 2 left and Figure 3 to cover both of these settings also for the baselines.\n3. I would have liked to see a bit more in-depth investigation of the pruning criteria, e.g. a plot of Spearman correlations between the preferred score and the others throughout training or a correlation matrix at various stages (say beginning, halfway through and end of training). I must say that I am not overly convinced that they matter too much, the variation of accuracy in Fig 2 seems to be only about 0.5% (although see questions). So I think it might be worth saving the page discussing the criteria in fewer of more thorough experiments.\n4. the paper makes some rather inaccurate claims vs the existing literature. In particular, it is not the first paper introducing a \"fully sparse BNN framework that maintains a consistently sparse Bayesian model through- out the training and inference\", this statement also applies to the (Ritter et al., 2021) paper, which is incorrectly cited as a post-hoc pruning paper (the paper does use post-hoc pruning as an optional step to further increase sparsity, but the core low-rank parameterization is maintained throughout training). This doesn't affect the contribution of course, but prior work needs to be contextualized correctly.\n5. I don't really see the need to make such claims in the first place, it is not obvious that sparsity in training is desirable. Of course it may be the case that a larger network that would not fit into memory without sparsity performs better, but then this needs to be demonstrated (or like-wise any hypothetical training speed increases resulting from a reduced number of FLOPs - in the age of parallelized computation, that is a mostly meaningless metric if it cannot be shown that a practical implementation can lead to actual cost savings).\n6. the abstract is simultaneously wordy and vague. I did not know what the paper was doing specifically after reading it, even though it really isn't hard to describe the method in 1 or 2 sentences. I would say that the low-rank/basis terminology led me in the wrong direction of thinking and a pruning-based description would have been clearer, but this may of course differ for readers with a different background.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)