Bayesian Layers: A Module for Neural Network Uncertainty

TL;DR

Bayesian Layers provides a flexible module for integrating uncertainty into neural networks, supporting various stochastic components and enabling scalable experimentation with models like Bayesian Transformers and deep GPs.

Contribution

It introduces a unified, extensible module for neural network uncertainty, compatible with existing libraries and capable of scaling to large models and complex probabilistic architectures.

Findings

Successfully trained a 5-billion parameter Bayesian Transformer on TPUv2

Demonstrated Bayesian Layers in diverse architectures like LSTMs and GPs

Integrated Bayesian Layers with Edward2 for probabilistic programming

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.