Efficient and Scalable Bayesian Neural Nets with Rank-1 Factors

Abstract

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.

Published in the International Conference on Machine Learning (ICML) 2020. Code available at https://github.com/google/edward2

References in corpus (15)

• Adam: A Method for Stochastic Optimization
• What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?
• Wide Residual Networks
• Weight Uncertainty in Neural Networks
• Variational Dropout and the Local Reparameterization Trick
• Deep Ensembles: A Loss Landscape Perspective
• Bayesian Compression for Deep Learning
• Variational Dropout Sparsifies Deep Neural Networks
• Multiplicative Normalizing Flows for Variational Bayesian Neural Networks
• A Simple Baseline for Bayesian Uncertainty in Deep Learning
• BatchEnsemble: An Alternative Approach to Efficient Ensemble and Lifelong Learning
• Structured and Efficient Variational Deep Learning with Matrix Gaussian Posteriors
• Bayesian Layers: A Module for Neural Network Uncertainty
• Model Selection in Bayesian Neural Networks via Horseshoe Priors
• Bayesian Inference for Large Scale Image Classification

Cited by in corpus (18)

• A Review of Uncertainty Quantification in Deep Learning: Techniques, Applications and Challenges
• DVERGE: Diversifying Vulnerabilities for Enhanced Robust Generation of Ensembles
• Hyperparameter Ensembles for Robustness and Uncertainty Quantification
• Bayesian Deep Ensembles via the Neural Tangent Kernel
• Bayesian Neural Network Priors Revisited
• Action and Perception as Divergence Minimization
• BNNpriors: A library for Bayesian neural network inference with different prior distributions
• Repulsive Deep Ensembles are Bayesian
• Global inducing point variational posteriors for Bayesian neural networks and deep Gaussian processes
• The Dual Information Bottleneck
• Unpacking Information Bottlenecks: Unifying Information-Theoretic Objectives in Deep Learning
• Amortized Conditional Normalized Maximum Likelihood: Reliable Out of Distribution Uncertainty Estimation
• On Stein Variational Neural Network Ensembles
• Informative Bayesian Neural Network Priors for Weak Signals
• Training on Test Data with Bayesian Adaptation for Covariate Shift
• Quantum Bayesian Neural Networks
• URSABench: Comprehensive Benchmarking of Approximate Bayesian Inference Methods for Deep Neural Networks
• Improving Uncertainty Calibration via Prior Augmented Data