# Bayesian Layers: A Module for Neural Network Uncertainty

arXiv: https://arxiv.org/abs/1812.03973 · HTML: https://mlsanity.com/paper/1812.03973/

## Key facts

- Cited by 57 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 10 Dec 2018.

## Built on by (most cited first, 30 of 57)

- [Bayesian Low-rank Adaptation for Large Language Models](https://mlsanity.com/paper/2308.13111/) (2023)
- [Hyperbolic Image Segmentation](https://mlsanity.com/paper/2203.05898/) (2022)
- [Priors in Bayesian Deep Learning: A Review](https://mlsanity.com/paper/2105.06868/) (2021)
- [Towards a Science of Human-AI Decision Making: A Survey of Empirical Studies](https://mlsanity.com/paper/2112.11471/) (2021)
- [Stochastic Neural Radiance Fields: Quantifying Uncertainty in Implicit 3D Representations](https://mlsanity.com/paper/2109.02123/) (2021)
- [BLoB: Bayesian Low-Rank Adaptation by Backpropagation for Large Language Models](https://mlsanity.com/paper/2406.11675/) (2024)
- [A review of uncertainty quantification in medical image analysis: probabilistic and non-probabilistic methods](https://mlsanity.com/paper/2310.06873/) (2023)
- [Uncertainty quantification in fine-tuned LLMs using LoRA ensembles](https://mlsanity.com/paper/2402.12264/) (2024)
- [Uncertainty-Aware Machine Translation Evaluation](https://mlsanity.com/paper/2109.06352/) (2021)
- [BlackJAX: Composable Bayesian inference in JAX](https://mlsanity.com/paper/2402.10797/) (2024)
- [Bayesian Attention Belief Networks](https://mlsanity.com/paper/2106.05251/) (2021)
- [Bayesian Transformer Language Models for Speech Recognition](https://mlsanity.com/paper/2102.04754/) (2021)
- [RecSim NG: Toward Principled Uncertainty Modeling for Recommender Ecosystems](https://mlsanity.com/paper/2103.08057/) (2021)
- [Calibrating Transformers via Sparse Gaussian Processes](https://mlsanity.com/paper/2303.02444/) (2023)
- [Gaussian Stochastic Weight Averaging for Bayesian Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2405.03425/) (2024)
- [Evaluating Uncertainty and Quality of Vision-Language-Action-enabled Robots](https://mlsanity.com/paper/2507.17049/) (2025)
- [A Bayesian neural network predicts the dissolution of compact planetary systems](https://mlsanity.com/paper/2101.04117/) (2021)
- [Fortuna: A Library for Uncertainty Quantification in Deep Learning](https://mlsanity.com/paper/2302.04019/) (2023)
- [deepregression: a Flexible Neural Network Framework for Semi-Structured Deep Distributional Regression](https://mlsanity.com/paper/2104.02705/) (2021)
- [Pathologies in priors and inference for Bayesian transformers](https://mlsanity.com/paper/2110.04020/) (2021)
- [Lymphoma segmentation from 3D PET-CT images using a deep evidential network](https://mlsanity.com/paper/2201.13078/) (2022)
- [Revisiting Kernel Attention with Correlated Gaussian Process Representation](https://mlsanity.com/paper/2502.20525/) (2025)
- [Applications of the Free Energy Principle to Machine Learning and Neuroscience](https://mlsanity.com/paper/2107.00140/) (2021)
- [Confidence estimation of classification based on the distribution of the neural network output layer](https://mlsanity.com/paper/2210.07745/) (2022)
- [Uncertainty Quantification for Forward and Inverse Problems of PDEs via Latent Global Evolution](https://mlsanity.com/paper/2402.08383/) (2024)
- [TyXe: Pyro-based Bayesian neural nets for Pytorch](https://mlsanity.com/paper/2110.00276/) (2021)
- [Block-local learning with probabilistic latent representations](https://mlsanity.com/paper/2305.14974/) (2023)
- [Flexible Visual Recognition by Evidential Modeling of Confusion and Ignorance](https://mlsanity.com/paper/2309.07403/) (2023)
- [Boundary Exploration for Bayesian Optimization With Unknown Physical Constraints](https://mlsanity.com/paper/2402.07692/) (2024)
- [Improving Interpretability of Deep Active Learning for Flood Inundation Mapping Through Class Ambiguity Indices Using Multi-spectral Satellite Imagery](https://mlsanity.com/paper/2404.19043/) (2024)

## Abstract (arXiv)

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.