# Feature Preserving Shrinkage on Bayesian Neural Networks via the R2D2 Prior

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 41 of its 69 references that are in the map.
- First posted on arXiv 23 May 2025.

## Builds on (most cited first, 30 of 41)

- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding](https://mlsanity.com/paper/1510.00149/) (2015)
- [Evidential Deep Learning to Quantify Classification Uncertainty](https://mlsanity.com/paper/1806.01768/) (2018)
- [FastGCN: Fast Learning with Graph Convolutional Networks via Importance Sampling](https://mlsanity.com/paper/1801.10247/) (2018)
- [Predictive Uncertainty Estimation via Prior Networks](https://mlsanity.com/paper/1802.10501/) (2018)
- [Stochastic Gradient Hamiltonian Monte Carlo](https://mlsanity.com/paper/1402.4102/) (2014)
- [Hands-on Bayesian Neural Networks -- a Tutorial for Deep Learning Users](https://mlsanity.com/paper/2007.06823/) (2020)
- [Variational Dropout Sparsifies Deep Neural Networks](https://mlsanity.com/paper/1701.05369/) (2017)
- [Graph Neural Networks with convolutional ARMA filters](https://mlsanity.com/paper/1901.01343/) (2019)
- [Understanding Measures of Uncertainty for Adversarial Example Detection](https://mlsanity.com/paper/1803.08533/) (2018)
- [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)
- [Optimal Approximation with Sparsely Connected Deep Neural Networks](https://mlsanity.com/paper/1705.01714/) (2017)
- [A Systematic Comparison of Bayesian Deep Learning Robustness in Diabetic Retinopathy Tasks](https://mlsanity.com/paper/1912.10481/) (2019)
- [All You Need is a Good Functional Prior for Bayesian Deep Learning](https://mlsanity.com/paper/2011.12829/) (2020)
- [Stochastic gradient Markov chain Monte Carlo](https://mlsanity.com/paper/1907.06986/) (2019)
- [Model Selection in Bayesian Neural Networks via Horseshoe Priors](https://mlsanity.com/paper/1705.10388/) (2017)
- [A Comprehensive guide to Bayesian Convolutional Neural Network with Variational Inference](https://mlsanity.com/paper/1901.02731/) (2019)
- [Posterior Concentration for Sparse Deep Learning](https://mlsanity.com/paper/1803.09138/) (2018)
- [Topology-Aware Graph Pooling Networks](https://mlsanity.com/paper/2010.09834/) (2020)
- [Sparsity information and regularization in the horseshoe and other shrinkage priors](https://mlsanity.com/paper/1707.01694/) (2017)
- [Tractable Function-Space Variational Inference in Bayesian Neural Networks](https://mlsanity.com/paper/2312.17199/) (2023)
- [Radial Bayesian Neural Networks: Beyond Discrete Support In Large-Scale Bayesian Deep Learning](https://mlsanity.com/paper/1907.00865/) (2019)
- [Consistent Sparse Deep Learning: Theory and Computation](https://mlsanity.com/paper/2102.13229/) (2021)

## Abstract (arXiv)

Bayesian neural networks (BNNs) treat neural network weights as random variables, which aim to provide posterior uncertainty estimates and avoid overfitting by performing inference on the posterior weights. However, the selection of appropriate prior distributions remains a challenging task, and BNNs may suffer from catastrophic inflated variance or poor predictive performance when poor choices are made for the priors. Existing BNN designs apply different priors to weights, while the behaviours of these priors make it difficult to sufficiently shrink noisy signals or they are prone to overshrinking important signals in the weights. To alleviate this problem, we propose a novel R2D2-Net, which imposes the R^2-induced Dirichlet Decomposition (R2D2) prior to the BNN weights. The R2D2-Net can effectively shrink irrelevant coefficients towards zero, while preventing key features from over-shrinkage. To approximate the posterior distribution of weights more accurately, we further propose a variational Gibbs inference algorithm that combines the Gibbs updating procedure and gradient-based optimization. This strategy enhances stability and consistency in estimation when the variational objective involving the shrinkage parameters is non-convex. We also analyze the evidence lower bound (ELBO) and the posterior concentration rates from a theoretical perspective. Experiments on both natural and medical image classification and uncertainty estimation tasks demonstrate satisfactory performance of our method.