BNNpriors: A library for Bayesian neural network inference with different prior distributions

TL;DR

BNNpriors is a versatile library that facilitates Bayesian neural network inference with various prior distributions, improving uncertainty calibration and predictive performance.

Contribution

It introduces a modular library enabling state-of-the-art MCMC inference with diverse priors, including heavy-tailed and hierarchical types, advancing Bayesian neural network research.

Findings

Facilitated discoveries on the cold posterior effect

Supports a wide range of prior distributions

Eases design of custom priors

Abstract

Bayesian neural networks have shown great promise in many applications where calibrated uncertainty estimates are crucial and can often also lead to a higher predictive performance. However, it remains challenging to choose a good prior distribution over their weights. While isotropic Gaussian priors are often chosen in practice due to their simplicity, they do not reflect our true prior beliefs well and can lead to suboptimal performance. Our new library, BNNpriors, enables state-of-the-art Markov Chain Monte Carlo inference on Bayesian neural networks with a wide range of predefined priors, including heavy-tailed ones, hierarchical ones, and mixture priors. Moreover, it follows a modular approach that eases the design and implementation of new custom priors. It has facilitated foundational discoveries on the nature of the cold posterior effect in Bayesian neural networks and will hopefully catalyze future research as well as practical applications in this area.