Training Bayesian Neural Networks with Sparse Subspace Variational Inference

Performance improvement:

As highlighted in your review, sparsity is key to training larger models within memory constraints. This leads to the question: Why not simply opt for a small dense model instead? The answer lies in the distinct advantages of large sparse models. Also shown by [5], “Sparse NNs keep the same model architectures with those dense large models, but only activate a small fraction of weights. Since state-of-the-art sparse NNs can match the utility performance of their dense counterparts (through either pruning or DST/PaI), sparsity can potentially be a free lunch to make training/inference cheaper while performing the same well.” This is also illustrated in our research, particularly evident in Figure 2 where we see a largely stable performance curve against sparsity level and hyperparameter selection. In contrast, a small dense model with a similar parameter count typically exhibits a noticeable drop in performance when compared to its larger, dense equivalent, as demonstrated in references [2,3,4]. Consequently, when considering models with comparable parameters, the superior performance of large sparse models makes them a preferable choice over small dense models in training scenarios.

A more contemporary and well-known instance highlighting the effectiveness of large sparse models is the adoption of Mixture-of-Expert (MoE) techniques. These techniques employ large sparse models not just during training but also throughout the inference process. They have been recognized as optimal solutions at the industrial level, as evidenced in high-profile projects such as Switch transformers [6] and Glam [7]. This trend underscores the practicality and efficiency of large sparse models in advanced applications, demonstrating their utility beyond the theoretical advantages of reduced FLOPs.

Abstract presentation

We would like to clarify that our paper is not only about pruning. Our paper is about how to find an optimal sparse subspace to do variational inference. To solve this problem, we optimize the sparse subspace and the corresponding variational model together. In this context, a removal(pruning)-based algorithm emerges as a natural solution within an alternative optimization schedule. We appreciate your input and will refine the abstract to provide a clearer description in the revised draft.

Initialization of weight mean

For initializing the mean values of weights added to the active set, we use a zero initialization.

About Figure 2 right

Your interpretation is essentially accurate. The phenomenon you described underscores a key advantage of our SSVI over conventional VI in Bayesian Neural Networks (BNNs). In Figure 2, we empirically verified this interpretation, as you correctly noted. It is important to recognize that training BNNs is more sensitive to initialization than training conventional neural networks. This makes our comparison between SSVI and VI particularly revealing. Consequently, the findings presented in Figure 2 right are neither trivial nor obvious and should be regarded as substantial.

Results of more sparsity levels for baseline algorithms

For our SSVI, we present outcomes at various sparsity levels in Figure 2, as you noticed. However, this does not extend to the baseline algorithms. Actually, all the baseline algorithms utilize implicit regularization to steer the model towards a level of sparsity that is not predefined, leading to a process that is frequently unstable and heavily dependent on precise hyperparameter tuning to achieve a satisfactory level of sparsity. Therefore, we first train the baselines following the settings in their original papers, then adopt the obtained sparsity level in our SSVI training to ensure a fair comparison.

We want to highlight that this is one of the advantages of our SSVI —- the sparsity can be flexibly assigned before training. We demonstrate the efficacy of SSVI across a range of sparsity levels and also explore its performance with various optimization schedules within the same sparsity level, as shown in Figures 2 and 3.

Influence of pruning criteria

Indeed, the choice of pruning criteria is crucial. It is important to note that global metrics like Spearman correlations may not be the most suitable metric for assessing the differences between various scores in this context. Since we utilize the score solely for selecting the top-K weights without incorporating it into the training process, and considering that K is significantly smaller than the total number of weights, a high Spearman score may not accurately reflect the divergence in the chosen weights.

Therefore, we have opted for the ratio of intersection over union (IoU). Each score selects the top-K weights for a fixed number K, and we calculate the intersection over union for all pairs among our five different scores.

In the following, we present the Intersection over Union (IoU) matrices for the five pruning criteria at the initial, middle, and final stages of training, respectively. Additionally, we include in the supplementary material the average of all 10 IoU values for different pairs across the training steps. The findings reveal that while the criteria show high similarity (reflected in high IoU scores) in the initial phase, they diverge as training progresses, potentially due to more precise uncertainty information and consequently more appropriate criteria. It is also notable that the traditional criteria, specifically $|\mu|$ and $\mathrm{SNR}{q_{\phi}}(\theta)$, exhibit a persistent similarity during the entire training period. On the other hand, the three novel criteria we introduced: $\mathbb{E}{q_\phi}|\theta|$, $\mathrm{SNR}{q_{\phi}}(|\theta|)$, and $\mathbb{E}{q_\phi}e^{\lambda|\theta|}$, show a marked deviation from these conventional metrics, emphasizing the innovative aspect of our methodology.

—-----------------IoU matrix at the beginning of training—-----------------

We greatly appreciate the extensive comments. We have carefully addressed each of the weaknesses and questions raised in your feedback. In light of these comprehensive revisions, we kindly hope that you will consider raising your score.

Experiments about OOD detection and distribution shift benchmark

We extended our experiments to two distribution shift benchmarks including both Out-of-Distribution (OOD) detection tasks and robustness benchmarks. The results demonstrate that our SSVI consistently improves over the strongest baseline, mBNN, significantly across all sparsity levels, metrics, and tasks. The results are presented both in the subsequent sections and in the updated Appendix D.

For OOD detection tasks, we computed the area under the receiver operating characteristic (AUROC) and the area under the precision-recall curve (AUPR) using our models, which were trained using SSVI on CIFAR-10/100 and tested on CIFAR-100/10 and SVHN, following the methodology of previous studies [1]. In this context, we employed the mBNN algorithm, the strongest previous work, as our baseline, ensuring equal sparsity for a fair and meaningful comparison.

We proceed to assess our model, trained using SSVI on CIFAR-10/100, on the respective corrupted versions, CIFAR-10/100-C. Both datasets encompass 15 distinct corruptions, each classified into 5 corruption levels. The outcomes are presented in the following table, illustrating curves for the average accuracy, loss, and Expected Calibration Error (ECE) across all 15 different corruptions for each level.

Calibration metric

While we understand your perspective, we respectfully disagree with the notion that calibration is not valuable. Calibration is widely recognized as a crucial metric in Bayesian deep learning literature, particularly as an indicator of a model's uncertainty capabilities. This perspective is not only supported in the paper you referred to [1] but also in our most closely related baseline [8]. Additionally, calibration is considered a significant metric in various studies investigating the posterior of Bayesian models, as seen in references [9, 10].

References

[1] 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.

[2] Li, Zhuohan, et al. "Train big, then compress: Rethinking model size for efficient training and inference of transformers." International Conference on machine learning. PMLR, 2020.

[3] ​​Evci, Utku, et al. "Rigging the lottery: Making all tickets winners." International Conference on Machine Learning. PMLR, 2020.

[4] Liu, Shiwei, et al. "The unreasonable effectiveness of random pruning: Return of the most naive baseline for sparse training." arXiv preprint arXiv:2202.02643 (2022).

[5] Liu, Shiwei, and Zhangyang Wang. "Ten lessons we have learned in the new" sparseland": A short handbook for sparse neural network researchers." arXiv preprint arXiv:2302.02596 (2023).

[6] Fedus, William, Barret Zoph, and Noam Shazeer. "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity." The Journal of Machine Learning Research 23.1 (2022): 5232-5270.

[7] Du, Nan, et al. "Glam: Efficient scaling of language models with mixture-of-experts." International Conference on Machine Learning. PMLR, 2022.

[8] Kong, Insung, et al. "Masked Bayesian Neural Networks: Theoretical Guarantee and its Posterior Inference." arXiv preprint arXiv:2305.14765 (2023).

[9] Wenzel, Florian, et al. "How good is the bayes posterior in deep neural networks really?." ICML 2020.

[10] Izmailov, Pavel, et al. "What are Bayesian neural network posteriors really like?." International conference on machine learning. PMLR, 2021.

[11] Kurtz, Mark, et al. "Inducing and exploiting activation sparsity for fast inference on deep neural networks." International Conference on Machine Learning. PMLR, 2020.

[12] Kurtic, Eldar, et al. "The optimal bert surgeon: Scalable and accurate second-order pruning for large language models." arXiv preprint arXiv:2203.07259 (2022).

[13] Gale, Trevor, et al. "Sparse gpu kernels for deep learning." SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020.

[14] 2020 Nvidia a100 tensor core gpu architecture. https://www.nvidia.com/content/dam/en- zz/Solutions/Data-Center/nvidia-ampere-architecture-whitepaper.pdf .

—----------------IoU matrix at the middle of training—----------------

—-----------------IoU matrix at the end of training—-----------------

Comparison with Ritter et al. 2021

Although (Ritter et al., 2021) achieve parameter reduction by storing the inducing weights $U$ that have much lower dimensions, they still require sampling full-dimension original weights $W$ in each forward pass, as stated in Section 3.2 of their paper. As a result, their computation costs in both training and inference remain similar or even larger than the dense BNNs. In contrast, our method directly enforces sparsity in $W$, which has a substantial reduction of computation costs theoretically. Besides, the sampling rule in (Ritter et al., 2021) requires a Gaussian parameterization, whereas our method is more flexible and can generalize beyond Gaussian parametrization. Yet, we still thank you for the comment and will change the description of (Ritter et al., 2021) in the revision.

Why do we need sparsity in training

We note that the concern raised is not unique to our work on sparse Bayesian Neural Networks (BNNs) but extends to the broader field of sparse neural networks. The advantages of having sparsity in training have been addressed in previous discussions, such as [5]. Here we formulate the answer about why we need sparsity in training from two directions: speedup and performance.

Real-time speedup:

First, we want to clarify that sparsity in training can lead to actual, real-time performance improvements on modern hardware, beyond just theoretical advantages suggested by FLOPs reduction. As shown by [5], “Another notable success was recently demonstrated by DeepSparse [11, 12], which successfully deploys large-scale BERT-level sparse models on modern Intel CPUs, obtaining 10× model size compression with < 1% accuracy drop, 10× CPU-inference speedup with < 2% drop, and 29× CPU-inference speedup with < 7.5% drop.” Additionally, there are considerable efforts, such as those by NVIDIA, to develop advanced GPU kernels like Sputnik [13] and the NVIDIA Ampere Architecture [14], which “built the momentum to better support finer-grained sparsity” and thereby increase real-time speedup. We will explore how to achieve real wall-time speedup with the specific hardware in our future work.

Thank you once again for your comprehensive review of our paper. We have addressed each of your comments in the five parts above and are confident that our responses have resolved all your concerns. As the close of the discussion period nears, we welcome any further questions you may have or would appreciate if you consider increasing the score for our submission. We are grateful for the time and effort you have invested in this review process!

Thank you for your supportive and thoughtful comments! We are grateful for your acknowledgment of the significance, novelty, and potential impact of this work. We answer your questions below.

More previous works in BNNs

In the introduction section of our revised paper, we have incorporated the work you referenced, along with other pertinent studies. Specifically, we have expanded our discussion to include additional works on stochastic variational inference [2,3,4] and Monte Carlo methods [5,6,7], beyond those already cited in the current version.

More experiments

We extended our experiments to two distribution shift benchmarks including both Out-of-Distribution (OOD) detection tasks and robustness benchmarks. The results demonstrate that our SSVI consistently improves over the strongest baseline, mBNN, significantly across all sparsity levels, metrics, and tasks. The results are presented both in the subsequent sections and in the updated Appendix D.

For OOD detection tasks, we computed the area under the receiver operating characteristic (AUROC) and the area under the precision-recall curve (AUPR) using our models, which were trained using SSVI on CIFAR-10/100 and tested on CIFAR-100/10 and SVHN, following the methodology of previous studies [1]. In this context, we employed the mBNN algorithm, the strongest previous work, as our baseline, ensuring equal sparsity for a fair and meaningful comparison.

We proceed to assess our model, trained using SSVI on CIFAR-10/100, on the respective corrupted versions, CIFAR-10/100-C. Both datasets encompass 15 distinct corruptions, each classified into 5 corruption levels. The outcomes are presented in the following table, illustrating curves for the average accuracy, loss, and Expected Calibration Error (ECE) across all 15 different corruptions for each level.

We agree that applying SSVI to more ML tasks, beyond standard models and applications, is an interesting question and we are happy to explore it further in our future work.

Clarity

Thank you for your thorough comments! We have made the suggested improvements to the paper. Please review the updated version. Here is a summary of responses to your questions:

1. The "20x compression, 20x FLOPs reduction" is benchmarked against the densely trained VI.
2. Our approach excels over VI-trained dense BNNs in terms of both accuracy and uncertainty metrics, such as the Expected Calibration Error (ECE).
3. The experimental section contains references and descriptions of the baselines.
4. The operator is element-wise matrix multiplication.
5. $B$ is batch size, $p,q$ are the input and output dimensions, $W$ is the weight matrix.
6. $l$ is the loss function received from the backward process.
7. Row 2 and 3 present the outcomes of our SSVI under two different sparsity levels, which are adjustable inputs in our SSVI.
8. We use equation (12) in our SSVI.

References

[1] 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.

[2] Hoffman, Matthew D., et al. "Stochastic variational inference." Journal of Machine Learning Research (2013).

[3] Dusenberry, Michael, et al. "Efficient and scalable bayesian neural nets with rank-1 factors." International conference on machine learning. PMLR, 2020.

[4] Welling, Max, and Yee W. Teh. "Bayesian learning via stochastic gradient Langevin dynamics." Proceedings of the 28th international conference on machine learning (ICML-11). 2011.

[5] Chen, Tianqi, Emily Fox, and Carlos Guestrin. "Stochastic gradient hamiltonian monte carlo." International conference on machine learning. PMLR, 2014.

[6] Cobb, Adam D., and Brian Jalaian. "Scaling Hamiltonian Monte Carlo inference for Bayesian neural networks with symmetric splitting." Uncertainty in Artificial Intelligence. PMLR, 2021.

[7] Wenzel, Florian, et al. "How good is the bayes posterior in deep neural networks really?." ICML 2020.

Thanks again for your careful review and acknowledgment of our response and detailed reply!

The characterization of Ghosh et al., 2019 [1]

We want to clarify that we actually both agree on the point that the sparsity is induced by priors for VI, as highlighted in both our first-round rebuttal: “[1] depends on complexly designed priors to induce sparsity in the model, ” and also in our related works section in the original paper “many studies have employed sparse-promoting priors to obtain sparse BNNs.” and “Louizos et al. (2017), Ghosh et al. (2019) and Bai et al. (2020) further explore the use of half-Cauchy, Horseshoe and Gaussian-Dirac spike-and slab priors respectively.”

We are not suggesting that dense training leads to sparsity. Our reference to "dense training" is to highlight our SSVI's superiority and distinction from previous works, including Ghosh et al., 2019 [1]. These studies induce sparsity through the use of specific priors. However, as noted in Section 4.3 of Ghosh et al., 2019 [1], “The Horseshoe and its regularized variant provide strong shrinkage towards zero for small $w_{kl}$. However, the shrunk weights, although tiny, are never actually zero. ”, this means their training is still dense so they still have a high training cost. In contrast, our SSVI approach achieves sparsity right from the start, leading to reductions in both training and inference costs.

About Vadera et al.

We have already added this paper Vadera et al. to our introduction section. We also agree that it’s an interesting direction to extend the method from MCMC to VI as outlined in Vadera et al., and evaluate the sparse structure of our SSVI in comparison to prior works using approaches from Vadera et al. However, this is beyond the scope of our current work. We will investigate this systematically in our future work.

Extention to other approximate Bayesian inference techniques

Bayesian neural networks based on other approximate Bayesian inference techniques, such as non-mean-field VI and MCMC, often involve distinct assumptions and approximations, such as those outlined in [4,5], which are beyond the scope of our current work as a conference submission. However, we appreciate the significance of this direction and plan to explore it in our future work.

Benefits of uncertainty metric in downstream tasks

We believe that the additional experiments on out-of-distribution (OOD) detection and distribution shift tasks, as outlined above and detailed in Appendix D.1, effectively address this question. The experiments we have are similar to those in [3]. Our SSVI surpasses all baseline models in terms of ECE for standard tasks, AUROC and AUPR for OOD tasks. This improvement in uncertainty metrics contributes to the enhanced accuracy on corrupted datasets, as evidenced by our results with CIFAR-10/100-C.

References

[1] 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.

[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.

[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.

[4] Nguyen, Son, et al. "Structured dropout variational inference for Bayesian neural networks." Advances in Neural Information Processing Systems 34 (2021): 15188-15202.

[5] Zhang et al. Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning. ICLR 2020.

It would be important to have empirical results on an expanded set of datasets + models + tasks

We extended our experiments to two distribution shift benchmarks including both Out-of-Distribution (OOD) detection tasks and robustness benchmarks. The results demonstrate that our SSVI consistently improves over the strongest baseline, mBNN, significantly across all sparsity levels, metrics, and tasks. The results are presented both in the subsequent sections and in the updated Appendix D.

For OOD detection tasks, we computed the area under the receiver operating characteristic (AUROC) and the area under the precision-recall curve (AUPR) using our models, which were trained using SSVI on CIFAR-10/100 and tested on CIFAR-100/10 and SVHN, following the methodology of previous studies [1]. In this context, we employed the mBNN algorithm, the strongest previous work, as our baseline, ensuring equal sparsity for a fair and meaningful comparison.

We proceed to assess our model, trained using SSVI on CIFAR-10/100, on the respective corrupted versions, CIFAR-10/100-C. Both datasets encompass 15 distinct corruptions, each classified into 5 corruption levels. The outcomes are presented in the following table, illustrating curves for the average accuracy, loss, and Expected Calibration Error (ECE) across all 15 different corruptions for each level.

Notations and Equation 2

Our intention behind introducing all notations is to provide a clearer illustration of both our intuition and algorithm. Regarding Equation 2, we want to emphasize that it serves as the objective function (loss function) for our Sparse Subspace Variational Inference (SSVI). This is similar to Equation 1, which is the objective function for standard Variational Inference (VI). Equation 2 is a slight modification of Equation 1, representing our goal to constrain the variational space to a sparse subspace. The proposed algorithm SSVI is then developed to solve the optimization problem stated in Equation 2.

Moreover, Equation 2 serves as the general form for sparse subspace VI. If we choose the common Gaussian as the variational distribution, it will result in a special case as shown in Equation 3. It is worth noting that other distribution families are possible but they may introduce unnecessary complexity and lead to alternative forms based on Equation 2.

In summary, Equation 2 is important to our framework, providing the foundation for the problem this paper tackles and guiding subsequent developments in our approach. We will revise this point in the final version to improve the clarity of the paper.

We are thankful for your valuable comments and inquiries which have contributed to a clearer understanding of key elements in our work. We believe the following explanations address all identified weaknesses and answer all the questions raised. If you find our response satisfactory, we kindly hope you will consider raising your score.

Comparison with previous works [1, 2]

We want to clarify that our work fundamentally differs from [1, 2]. Our approach involves developing a new framework for identifying the optimal sparse subspace, ensuring sparsity throughout the training process. However, for [1], we quote what they mentioned at the beginning of their Section 3.1:

“ We emphasize that the contribution of this work is not to propose new methods for deriving sparse sub-networks, but rather to evaluate the impact of sparse sub-network structures on MCMC-based inference methods”

This indicates that [1] focuses on applying MCMC to pre-existing sparse neural networks to assess their impact on MCMC inference, which significantly differs from our work. While it is possible to modify their approach for VI, such as implementing VI on these pre-existing networks, our work diverges entirely from this route in two points: Our work aims to develop a novel algorithm for choosing sub-networks instead of just using existing methods. In contrast to the two-step method of first developing a sparse network and then conducting MCMC in [1], our approach integrates the selection of sub-networks directly into the VI training process, creating a more efficient, single-phase procedure.

Similar to several studies we discussed in the related works section that differ from our SSVI, [2] employs a one-shot pruning method following dense training, leading to significantly increased training costs compared with our SSVI. This one-shot approach tends to result in sub-optimal sub-networks, unlike our method which iteratively removes and adds elements to pinpoint a better sub-network. Furthermore, [2] depends on complexly designed priors to induce sparsity in the model, which can result in unpredictable levels of final sparsity. In contrast, our approach allows for precise control over the final sparsity level right from the start of the training process. We have included [1,2] in the related works in the revision.

The authors haven't mentioned the model architecture used in their experiments.

We have mentioned the model architecture used in experiments, ResNet-18, in the first line of Section 4.1 on page 7. This model architecture is standard and has been widely used in the sparse Bayesian Neural Network (BNN) literature such as mBNN (Kong et al., 2023) and AEB (Deng et al., 2019). In fact, many papers on sparse BNNs (including Kong et al., 2023 and Deng et al., 2019) used ResNet as the only model architecture for large-scale Bayesian deep learning experiments. Therefore, we also chose ResNet to ensure a fair comparison.

I thank the authors for engaging in the rebuttal. I'll respond to their rebuttal part by part.

Furthermore, I believe the above responses are fair - however, I don't quite agree with the characterization of on work by Ghosh et al., 2019 [1]. The crux of the work is to use sparsity inducing priors for VI - the sparsity isn't induced by dense training. Could the authors elaborate more on this?

Finally, while I agree that the work of Vadera et al. doesn't really develop novel techniques for finding optimal sparse substructure, I still believe that it's an important baseline to compare against.

References

[1] 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.

Thank you once more for acknowledging our rebuttal and the additional experiments we conducted. In your latest response, you raised four points:

1. The characterization of Ghosh et al., 2019 [1] – we have clarified this and confirmed that our views are in sync with yours, indicating no divergence between us on this aspect.
2. Simplification of Equation 2 – we have already made these simplifications in our revised paper.
3. Comparison with Vadera et al. – we acknowledge this valuable direction and plan to include it in our future research.
4. Extension to other BNNs – this direction too will be addressed in our upcoming work.

We are confident that we have adequately addressed these comments, as well as those from your first-round review. If you find our responses to be satisfactory, we would appreciate it if you would consider raising your score for our submission. Thank you for your time and valuable input!