Bayesian Layers: A Module for Neural Network Uncertainty

By Dustin Tran, Michael W. Dusenberry, Mark van der Wilk, Danijar Hafner. 2018

150 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Statistical and Nonlinear Physics

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 Statistical and Nonlinear Physics

Citations in the corpus

Towards a Science of Human-AI Decision Making: A Survey of Empirical Studies (2021, background): “Recent work has experimented with deep probabilistic models to give more reliable uncertainty estimations, including Bayesian neural networks [136, 154], deep neural networks that integrate dropout [47], or ensemble methods to approximate Bayesian inference [83].”
Noise Contrastive Priors for Functional Uncertainty (2018, methodology): “An implementation of NCP is also available in Aboleth (Aboleth Developers, 2017) and Bayesian Layers (Tran et al., 2018).”
Hyperparameter Ensembles for Robustness and Uncertainty Quantification (2020, background): “• From an implementation perspective, (7) enables direct reuse of existing code, e.g., DenseBatchEnsemble and Conv2DBatchEnsemble from [55].”
The k-tied Normal Distribution: A Compact Parameterization of Gaussian Mean Field Posteriors in Bayesian Neural Networks (2020, unknown): “To implement the MLP and CNN models we use the tfp.layers module from the Tensorflow Probability, while to implement the LSTM model we use the LSTMCellReparameterization5 class from the Edward2 Layers module (Tran et al., 2019).”
Action and Perception as Divergence Minimization (2020, methodology): “BNNs have been leveraged for RL for robustness (Okada et al., 2020; Tran et al., 2019) and exploration (Houthooft et al., 2016; Azizzadenesheli et al., 2018).”