Bayesian Layers: A Module for Neural Network Uncertainty

Abstract

We describe Bayesian Layers, a module designed for fast experimentation with neural network uncertainty. It extends neural network libraries with drop-in replacements for common layers. This enables composition via a unified abstraction over deterministic and stochastic functions and allows for scalability via the underlying system. These layers capture uncertainty over weights (Bayesian neural nets), pre-activation units (dropout), activations ("stochastic output layers"), or the function itself (Gaussian processes). They can also be reversible to propagate uncertainty from input to output. We include code examples for common architectures such as Bayesian LSTMs, deep GPs, and flow-based models. As demonstration, we fit a 5-billion parameter "Bayesian Transformer" on 512 TPUv2 cores for uncertainty in machine translation and a Bayesian dynamics model for model-based planning. Finally, we show how Bayesian Layers can be used within the Edward2 probabilistic programming language for probabilistic programs with stochastic processes.

arXiv commentCode available at https://github.com/tensorflow/tensor2tensor

References in corpus (16)

• Auto-Encoding Variational Bayes
• Caffe: Convolutional Architecture for Fast Feature Embedding
• Neural Discrete Representation Learning
• MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems
• Density estimation using Real NVP
• Pyro: Deep Universal Probabilistic Programming
• GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration
• MADE: Masked Autoencoder for Distribution Estimation
• Image Transformer
• The Reversible Residual Network: Backpropagation Without Storing Activations
• Flipout: Efficient Pseudo-Independent Weight Perturbations on Mini-Batches
• Bayesian Recurrent Neural Networks
• Deep Gaussian Processes for Regression using Approximate Expectation Propagation
• Uncertainty Estimations by Softplus normalization in Bayesian Convolutional Neural Networks with Variational Inference
• Mesh-TensorFlow: Deep Learning for Supercomputers
• Simple, Distributed, and Accelerated Probabilistic Programming

Cited by in corpus (16)

• A Review of Uncertainty Quantification in Deep Learning: Techniques, Applications and Challenges
• Application of belief functions to medical image segmentation: A review
• Identifying Causal-Effect Inference Failure with Uncertainty-Aware Models
• Efficient and Scalable Bayesian Neural Nets with Rank-1 Factors
• Hyperparameter Ensembles for Robustness and Uncertainty Quantification
• Action and Perception as Divergence Minimization
• The k-tied Normal Distribution: A Compact Parameterization of Gaussian Mean Field Posteriors in Bayesian Neural Networks
• Uncertainty-Aware Machine Translation Evaluation
• A Tutorial on Sparse Gaussian Processes and Variational Inference
• Bayesian Neural Networks at Scale: A Performance Analysis and Pruning Study
• Neural Mixture Distributional Regression
• A Bayesian neural network predicts the dissolution of compact planetary systems
• Bayesian Transformer Language Models for Speech Recognition
• The Bayesian Method of Tensor Networks
• UNITE: Uncertainty-based Health Risk Prediction Leveraging Multi-sourced Data
• Estimating the Uncertainty of Neural Network Forecasts for Influenza Prevalence Using Web Search Activity