# Training Bayesian Neural Networks with Sparse Subspace Variational Inference

arXiv: https://arxiv.org/abs/2402.11025 · HTML: https://mlsanity.com/paper/2402.11025/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 46 references that are in the map.
- First posted on arXiv 16 Feb 2024.

## Builds on (most cited first, 29 of 29)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://mlsanity.com/paper/1803.03635/) (2018)
- [Rigging the Lottery: Making All Tickets Winners](https://mlsanity.com/paper/1911.11134/) (2019)
- [Variational Dropout and the Local Reparameterization Trick](https://mlsanity.com/paper/1506.02557/) (2015)
- [Bayesian Deep Learning and a Probabilistic Perspective of Generalization](https://mlsanity.com/paper/2002.08791/) (2020)
- [Stochastic Gradient Hamiltonian Monte Carlo](https://mlsanity.com/paper/1402.4102/) (2014)
- [Variational Dropout Sparsifies Deep Neural Networks](https://mlsanity.com/paper/1701.05369/) (2017)
- [Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning](https://mlsanity.com/paper/1902.03932/) (2019)
- [Efficient and Scalable Bayesian Neural Nets with Rank-1 Factors](https://mlsanity.com/paper/2005.07186/) (2020)
- [Bayesian Compression for Deep Learning](https://mlsanity.com/paper/1705.08665/) (2017)
- [Towards Safe Autonomous Driving: Capture Uncertainty in the Deep Neural Network For Lidar 3D Vehicle Detection](https://mlsanity.com/paper/1804.05132/) (2018)
- [A Systematic Comparison of Bayesian Deep Learning Robustness in Diabetic Retinopathy Tasks](https://mlsanity.com/paper/1912.10481/) (2019)
- [Scaling Hamiltonian Monte Carlo Inference for Bayesian Neural Networks with Symmetric Splitting](https://mlsanity.com/paper/2010.06772/) (2020)
- [Structured Bayesian Pruning via Log-Normal Multiplicative Noise](https://mlsanity.com/paper/1705.07283/) (2017)
- [Efficient Variational Inference for Sparse Deep Learning with Theoretical Guarantee](https://mlsanity.com/paper/2011.07439/) (2020)
- [Structured Variational Learning of Bayesian Neural Networks with Horseshoe Priors](https://mlsanity.com/paper/1806.05975/) (2018)
- [Stochastic Particle-Optimization Sampling and the Non-Asymptotic Convergence Theory](https://mlsanity.com/paper/1809.01293/) (2018)
- [An Adaptive Empirical Bayesian Method for Sparse Deep Learning](https://mlsanity.com/paper/1910.10791/) (2019)
- [Computationally Efficient Bayesian Learning of Gaussian Process State Space Models](https://mlsanity.com/paper/1506.02267/) (2015)
- [Sparse Uncertainty Representation in Deep Learning with Inducing Weights](https://mlsanity.com/paper/2105.14594/) (2021)
- [Asymptotically Optimal Exact Minibatch Metropolis-Hastings](https://mlsanity.com/paper/2006.11677/) (2020)
- [AMAGOLD: Amortized Metropolis Adjustment for Efficient Stochastic Gradient MCMC](https://mlsanity.com/paper/2003.00193/) (2020)
- [Variance Reduction in Stochastic Particle-Optimization Sampling](https://mlsanity.com/paper/1811.08052/) (2018)
- [Masked Bayesian Neural Networks : Theoretical Guarantee and its Posterior Inference](https://mlsanity.com/paper/2305.14765/) (2023)
- [Self-Adversarially Learned Bayesian Sampling](https://mlsanity.com/paper/1811.08929/) (2018)
- [Bayesian Sparse learning with preconditioned stochastic gradient MCMC and its applications](https://mlsanity.com/paper/2006.16376/) (2020)
- [Impact of Parameter Sparsity on Stochastic Gradient MCMC Methods for Bayesian Deep Learning](https://mlsanity.com/paper/2202.03770/) (2022)

## Built on by (most cited first, 7 of 7)

- [Kubrick: Multimodal Agent Collaborations for Synthetic Video Generation](https://mlsanity.com/paper/2408.10453/) (2024)
- [BMRS: Bayesian Model Reduction for Structured Pruning](https://mlsanity.com/paper/2406.01345/) (2024)
- [A Bayesian Interpretation of Adaptive Low-Rank Adaptation](https://mlsanity.com/paper/2409.10673/) (2024)
- [Coeff-Tuning: A Graph Filter Subspace View for Tuning Attention-Based Large Models](https://mlsanity.com/paper/2503.18337/) (2025)
- [Stochastic Weight Sharing for Bayesian Neural Networks](https://mlsanity.com/paper/2505.17856/) (2025)
- [Bayesian Lottery Ticket Hypothesis](https://mlsanity.com/paper/2602.18825/) (2026)
- [Compressed Active Subspaces for Scalable Bayesian Inference](https://mlsanity.com/paper/2609.19539/) (2026)

## Abstract (arXiv)

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 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 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-20x compression 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 on both accuracy and uncertainty metrics.