Source-linked AI summary
Bayesian Layers: A Module for Neural Network Uncertainty
Dustin Tran, Michael W. Dusenberry, Mark van der Wilk, Danijar Hafner
TL;DR
Bayesian Layers address software limitations that make uncertainty-aware functions difficult to compose and scale. They extend neural-network layers to distributions over functions, support modular inference and reversible computation, and demonstrate applications ranging from a 5-billion parameter Bayesian Transformer to model-based planning. The framework’s scope is limited because its per-layer inference modularization makes full-parameter methods such as Markov chain Monte Carlo difficult to fit.
Problem
Existing uncertainty-aware software is often monolithic or requires explicit variable management, leaving no unified layer-based design for flexible, scalable research.
Method
Bayesian Layers extend neural-network libraries with composable layers representing distributions over weights, units, activations, or functions, including modular estimators and reversible operations.
Results
The framework demonstrates a 5-billion parameter Bayesian Transformer on 512 TPUv2 cores and a Bayesian dynamics model that matches the original agent’s 614 cheetah score.
Takeaways & Limitations
Bayesian Layers make uncertainty-aware models easier to compose within scalable deep-learning ecosystems and support applications in translation and model-based planning.
Takeaways & Limitations
Per-layer inference modularization makes full-parameter methods such as Markov chain Monte Carlo difficult to fit within the framework.
Abstract
from arXiv · showhide
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.
1. Introduction
Bayesian Layers address software limitations in uncertainty-aware neural networks by extending familiar layer abstractions to distributions over functions. The framework supports scalable composition across Bayesian neural networks, Gaussian processes, reversible layers, and probabilistic programs.
- Existing uncertainty-aware software is often monolithic, limiting research flexibility and communication primitives needed to scale models to billions of parameters.
- Bayesian Layers replace deterministic-only layers with distributions over functions while preserving neural-network composition semantics.Uncertainty may reside in weights, pre-activation units, activations, or the function itself.
- The framework provides drop-in layer replacements and unified abstractions for Bayesian LSTMs, deep GPs, flow-based models, and Edward2 probabilistic programs.The supplied examples include Bayesian recurrent models, Gaussian processes, flow-based models, and stochastic processes.
2. Bayesian Layers
Bayesian Layers extend ordinary neural-network layers into modular distributions over weights, activations, functions, and reversible transformations while preserving familiar layer composition. The framework supports Bayesian neural networks, Gaussian processes, stochastic outputs, and flow-based models, but its per-layer inference modularity limits compatibility with full-parameter inference methods.
- 2. Bayesian Layers: Bayesian Layers represent uncertainty over weights, pre-activation units, activations, or entire functions, with each sample instantiating a different function.The framework extends layers from deterministic functions to distributions over functions.
- 2. Bayesian Layers: Each estimator is implemented as its own Layer, allowing the same Bayesian neural network to use different computational graphs for sampling, reparameterization, or quadrature.This design separates integral estimation from the Bayesian layer specification.
- 2.1. Bayesian Neural Network Layers: Bayesian layers preserve deterministic-layer constructor arguments and tensor input/output signatures, enabling drop-in replacement while changing parameter initialization and regularization defaults.The implementation overloads initializer and regularizer arguments to represent distributions and KL penalties.
- 2.3. Stochastic Output Layers: Stochastic output layers expose tractable output distributions for likelihood-based models and stochastic autoencoder encoders or decoders.Examples include discretized logistic mixtures and variational autoencoders using log-probability or KL divergence.
- 2.4. Reversible Layers: Reversible layers add inverse computation and optional log-determinant Jacobians, allowing transformed RandomVariables and flow-based models such as RealNVP.The same mechanism can also support reversible networks that avoid storing intermediate activations during forward passes.
- 2.5. Layers for Probabilistic Programming: The framework typically ties model specification to variational inference, making methods that depend on the full parameter space, such as Markov chain Monte Carlo, difficult to fit.This is the paper’s stated scope boundary for modular per-layer inference.
3. Experiments
Experiments demonstrate that Bayesian Layers supports scalable Bayesian models and practical uncertainty-aware applications, including language modeling, translation, and model-based planning.
- Experiments: The experiments present Bayesian Layers as efficient enough to enable model classes at new scales and flexibility, with benefits including model-based planning.The paper frames these experiments as demonstrations of both system efficiency and application value.
- Experiments: Bayesian Layers also supports probabilistic-programming constructions such as a Cox process with a deep GP prior and sparse GP posterior approximation.This separates modeling from inference and permits more flexible inference algorithms.
- Model-Parallel Bayesian Transformer for Machine Translation: Over 5-billion-parameter Bayesian Transformer models were fit on 512 TPUv2 cores, using up to 2500 TFLOPs.The Bayesian Transformer augmented attention projection matrices and pointwise feedforward layers with Bayesian counterparts.
- Model-Parallel Bayesian Transformer for Machine Translation: State-of-the-art perplexity was maintained at 23.1 on the One-Billion-Word Language Modeling Benchmark while achieving higher predictive variance.The authors identify challenges in scaling Bayesian neural networks and understanding their text applications.
- Bayesian Dynamics Model for Model-Based Reinforcement Learning: Bayesian PlaNet reached a score of 614 on cheetah, matching the original agent after replacing transition layers with VariationalDense layers.The weight-posterior KL divergence indicated a non-trivial learned belief.
4. Discussion
The discussion presents Bayesian Layers as a general module for uncertainty-aware modeling and illustrates its use in Bayesian planning, where performance is retained while uncertainty is modeled.
- Discussion: Bayesian Layers is designed for fast experimentation with neural network uncertainty.The module extends neural network libraries with uncertainty-aware layer abstractions.
- Discussion: Bayesian PlaNet reaches the same task performance as the original method while learning a non-trivial weight posterior and predicting 50 future time steps.The score summarizes median performance over five seeds and ten episodes, with 5th–95th percentile shading.
A. Bayesian ResNet-50
The Bayesian ResNet-50 appendix provides an implementation example using VariationalConv2D layers within a convolutional block.
- Bayesian ResNet-50: The Bayesian ResNet-50 convolutional block replaces convolutional operations with VariationalConv2D layers.The block applies specified filters, kernel sizes, and strides before batch normalization and ReLU activation.
B. Bayesian PlaNet
Bayesian PlaNet’s latent states support more precise predictions of reinforcement-learning environment velocities than the referenced prior result.
- Bayesian PlaNet: Bayesian PlaNet appears to capture more information in latent codes, producing more precise velocity predictions.The paper characterizes this additional latent information as “world knowledge.”