Singular Bayesian Neural Networks

TL;DR

This work introduces singular Bayesian neural networks by factorizing weight matrices as $W=AB^{\top}$, which concentrates the posterior on the rank-$r$ manifold and induces structured correlations among weights. The authors provide theoretical guarantees via PAC-Bayes and Gaussian-complexity analyses and demonstrate practical benefits: competitive predictive performance with up to $15\times$ fewer parameters and improved out-of-distribution detection and calibration across MLPs, LSTMs, and Transformers. Empirically, low-rank Bayesian layers achieve strong uncertainty quantification while greatly reducing memory footprints, with efficiency gains that scale at larger model sizes. The results highlight a calibration–OOD tradeoff and show the potential of end-to-end low-rank variational inference as a principled path to scalable, uncertainty-aware deep learning.

Abstract

Bayesian neural networks promise calibrated uncertainty but require $O(mn)$ parameters for standard mean-field Gaussian posteriors. We argue this cost is often unnecessary, particularly when weight matrices exhibit fast singular value decay. By parameterizing weights as $W = AB^{\top}$ with $A \in \mathbb{R}^{m \times r}$, $B \in \mathbb{R}^{n \times r}$, we induce a posterior that is singular with respect to the Lebesgue measure, concentrating on the rank-$r$ manifold. This singularity captures structured weight correlations through shared latent factors, geometrically distinct from mean-field's independence assumption. We derive PAC-Bayes generalization bounds whose complexity term scales as $\sqrt{r(m+n)}$ instead of $\sqrt{m n}$, and prove loss bounds that decompose the error into optimization and rank-induced bias using the Eckart-Young-Mirsky theorem. We further adapt recent Gaussian complexity bounds for low-rank deterministic networks to Bayesian predictive means. Empirically, across MLPs, LSTMs, and Transformers on standard benchmarks, our method achieves predictive performance competitive with 5-member Deep Ensembles while using up to $15\times$ fewer parameters. Furthermore, it substantially improves OOD detection and often improves calibration relative to mean-field and perturbation baselines.

Figures (20)

• Figure 1: Weight correlation structure. Comparison of Full-Rank BBB (left) and Low-Rank Gaussian r=15 (right). Full-Rank exhibits diagonal correlations; Low-Rank captures block structure.
• Figure 2: Geometric distinction between mean-field and low-rank posteriors. 3D projection of weight space $\mathbb{R}^{m \times n}$. Left: Rank-$r$ manifold $\mathcal{M}_r$ (blue surface, dimension $r(m+n-r)$). Middle: Mean-field posterior $q_{\text{MF}}(W)$ has full-dimensional support (volume). Right: Low-rank posterior $q_{\text{LR}}(A,B)$ concentrates on the manifold surface.
• Figure 3: Empirical generalization bounds for low-rank Bayesian neural networks. PAC-Bayes (left) and Gaussian complexity (right) bounds use empirical values from trained LSTM model and training data. PAC-Bayes bound exhibits critical rank $r^* \approx 11$ transitioning from non-vacuous to vacuous ($>1$, dashed red line). Gaussian complexity decreases from 45.56 (full-rank) to 18.97 with rank reduction. Full-rank Bayesian models yield vacuous PAC-Bayes bounds regardless of configuration.
• Figure 4: Model comparison on MIMIC-III (averaged across 5 seeds). Low-Rank Gaussian r=15 (orange) achieves superior OOD detection. Deep Ensemble maintains better calibration and in-domain discrimination. Rank-1 multiplicative (green) achieves better calibration but weaker OOD detection. Full-Rank BBB (blue) shows balanced but moderate performance across metrics.
• Figure 5: Selective prediction on Beijing Air Quality LSTM. While Deep Ensemble (red) achieves best point predictions at 100% retention, Bayesian methods outperform when discarding the most uncertain samples. Low-Rank (green) achieves largest improvement (17.4% MAE reduction at 80% retention), demonstrating superior uncertainty quality for selective prediction.

Theorems & Definitions (85)

• Definition 3.1: Induced Posterior
• Lemma 3.2: Constrained Support
• Lemma 3.3: Measure Zero of $\mathcal{R}_r$
• Theorem 3.4: Singularity of the Induced Posterior
• proof
• Lemma 3.5: Covariance of Weight Entries
• Theorem 3.6: Bounded Loss Approximation
• Theorem 3.7: Decomposition of Approximation Error
• Theorem 3.8: Tighter Bounds for Low-Rank Posteriors
• Theorem 3.9: Gaussian-complexity bound for low-rank BNN predictors