Efficient and Scalable Bayesian Neural Nets with Rank-1 Factors

By Michael W. Dusenberry, Ghassen Jerfel, Yeming Wen, Yi-An Ma, Jasper Snoek, Katherine A. Heller, Balaji Lakshminarayanan, Dustin Tran. 2020

240 citations recorded by the source corpus.

Bayesian neural networks (BNNs) demonstrate promising success in improving the robustness and uncertainty quantification of modern deep learning. However, they generally struggle with underfitting at scale and parameter efficiency. On the other hand, deep ensembles have emerged as alternatives for uncertainty quantification that, while outperforming BNNs on certain problems, also suffer from efficiency issues. It remains unclear how to combine the strengths of these two approaches and remediate their common issues. To tackle this challenge, we propose a rank-1 parameterization of BNNs, where each weight matrix involves only a distribution on a rank-1 subspace. We also revisit the use of mixture approximate posteriors to capture multiple modes, where unlike typical mixtures, this approach admits a significantly smaller memory increase (e.g., only a 0.4% increase for a ResNet-50 mixture of size 10). We perform a systematic empirical study on the choices of prior, variational posterior, and methods to improve training. For ResNet-50 on ImageNet, Wide ResNet 28-10 on CIFAR-10/100, and an RNN on MIMIC-III, rank-1 BNNs achieve state-of-the-art performance across log-likelihood, accuracy, and calibration on the test sets and out-of-distribution variants.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Efficiency & Compression, Probabilistic & Bayesian Methods

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Efficiency & Compression

In Probabilistic & Bayesian Methods

Citations in the corpus

Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness (2020, methodology): “Therefore we expect it to provide complementary benefits to approaches such as (efficient) ensembles and Bayesian neural networks [21, 42, 79] which marginalize over the representation parameters as well.”
Deterministic Neural Networks with Appropriate Inductive Biases Capture Epistemic and Aleatoric Uncertainty (2021, methodology): “Most well-known methods of uncertainty quantiﬁcation in deep learning [Blundell et al., 2015, Gal and Ghahramani, 2016, Lakshminarayanan et al., 2017, Wen et al., 2019, Dusenberry et al., 2020] require multiple forward passes at test time which can represent a signiﬁcant computational overhead.”
Hyperparameter Ensembles for Robustness and Uncertainty Quantification (2020, background): “Many mechanisms have been proposed to further foster diversity in ensembles of neural networks, e.g., based on cyclical learning rates [29] or Bayesian analysis [15].”
What Are Bayesian Neural Network Posteriors Really Like? (2021, methodology): “These methods include stochastic variational inference (Hoffman et al., 2013; Graves, 2011; Blundell et al., 2015; Kingma et al., 2015; Molchanov et al., 2017; Louizos & Welling, 2017; Khan et al., 2018; Zhang et al., 2018; Wu et al., 2018; Osawa et al., 2019; Dusenberry et al., 2020), dropout (Sriv”
A Review of Uncertainty Quantification in Deep Learning: Techniques, Applications and Challenges (2020, methodology): “[273] devised a rank-1 parameterization of BNNs and also utilized mixture approximate posteriors to capture multiple modes.”
Action and Perception as Divergence Minimization (2020, background): “, 2019), low rank posteriors (Izmailov et al., 2018; Dusenberry et al., 2020), and improved inference algorithms (Wen et al.”