Training Bayesian Neural Networks with Sparse Subspace Variational Inference

TL;DR

This work tackles the substantial computational burden of Bayesian neural networks by introducing Sparse Subspace Variational Inference (SSVI), a fully sparse framework that maintains a high-sparsity Bayesian model during training and inference. By constraining the posterior to a subspace of dimension $s$ within $\,\mathbb{R}^d$ and alternately optimizing the subspace and variational parameters, SSVI achieves dramatic reductions in training FLOPs and model size while preserving or improving accuracy and uncertainty metrics. The authors develop novel removal/addition criteria based on weight distribution statistics, with a rigorously designed optimization procedure and the Local Reparameterization Trick to boost efficiency. Empirical results on CIFAR-10/100 with ResNet-18 show 10–20× sparsity-induced reductions and up to 20× training FLOPs reductions, outperforming previous sparse BNN approaches and sometimes even surpassing VI-trained dense BNNs in both accuracy and uncertainty calibration. Overall, SSVI enables scalable, uncertainty-aware Bayesian learning for large networks by coupling fully sparse training with robust subspace-based variational inference.

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

Figures (5)

• Figure 1: Flops v.s. Accuracy. Our SSVI achieves the best test accuracy with the lowest training FLOPs.
• Figure 2: Left: the performance of SSVI for BNNs compared to RigL evci2020rigging for standard NNs across various sparsity levels. SSVI demonstrates superior robustness to different sparsity levels. 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 on precise hyperparameter tuning, especially concerning KL 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, $\mathrm{SNR}(|\theta|)$ stands out, leveraging uncertainty information in a theoretically backed manner.
• 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.
• Figure 4: Performance analysis on CIFAR-10-C and CIFAR-100-C. SSVI shows consistently better results across all metrics under corruption.
• Figure 5: Average IoU of all 10 pairs between 5 different criteria.