TyXe: Pyro-based Bayesian neural nets for Pytorch

2021/10/01 by Hippolyt Ritter, Ritter, Hippolyt, Theofanis Karaletsos +1 · 1 citation
Computer Science · Mathematics · #Adversarial Robustness in Machine Learning #FOS: Computer and information sciences #Gaussian Processes and Bayesian Inference #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Machine Learning and Data Classification #cs.LG #stat.ML

paper · pdf · doi:10.48550/arxiv.2110.00276

Previously presented at PROBPROG 2020

arxiv created 2021/10/01 · openalex publication_date 2021/10/01 · arxiv updated 2021/10/04 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

We introduce TyXe, a Bayesian neural network library built on top of Pytorch and Pyro. Our leading design principle is to cleanly separate architecture, prior, inference and likelihood specification, allowing for a flexible workflow where users can quickly iterate over combinations of these components. In contrast to existing packages TyXe does not implement any layer classes, and instead relies on architectures defined in generic Pytorch code. TyXe then provides modular choices for canonical priors, variational guides, inference techniques, and layer selections for a Bayesian treatment of the specified architecture. Sampling tricks for variance reduction, such as local reparameterization or flipout, are implemented as effect handlers, which can be applied independently of other specifications. We showcase the ease of use of TyXe to explore Bayesian versions of popular models from various libraries: toy regression with a pure Pytorch neural network; large-scale image classification with torchvision ResNets; graph neural networks based on DGL; and Neural Radiance Fields built on top of Pytorch3D. Finally, we provide convenient abstractions for variational continual learning. In all cases the change from a deterministic to a Bayesian neural network comes with minimal modifications to existing code, offering a broad range of researchers and practitioners alike practical access to uncertainty estimation techniques. The library is available at https://github.com/TyXe-BDL/TyXe.

Citations

• Probabilistic Backpropagation for Scalable Learning of Bayesian Neural Networks
• Automated Variational Inference in Probabilistic Programming
• Variational Dropout and the Local Reparameterization Trick
• A Practical Bayesian Framework for Backpropagation Networks
• The No-U-Turn Sampler: Adaptively Setting Path Lengths in Hamiltonian Monte Carlo
• Probabilistic Meta-Representations Of Neural Networks
• Weight Uncertainty in Neural Networks
• TensorFlow: A system for large-scale machine learning
• Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• The k-tied Normal Distribution: A Compact Parameterization of Gaussian Mean Field Posteriors in Bayesian Neural Networks
• NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis
• Sparse Uncertainty Representation in Deep Learning with Inducing Weights

Cited by

• Bayesian Neural Networks versus deep ensembles for uncertainty quantification in machine learning interatomic potentials

Related

• BayesDLL: Bayesian Deep Learning Library
• Pyro: Deep Universal Probabilistic Programming
• py/cuTAGI: An Open-Source Library for Tractable Approximate Gaussian Inference in Bayesian Neural Networks
• Bayesian Neural Networks: Essentials
• Torchbearer: A Model Fitting Library for PyTorch