HybridBNN: Joint Training of Deterministic and Stochastic Layers in Bayesian Neural Nets

Amin Nejatbakhsh, Julien Boussard

Keywords: Bayesian Neural Networks, Calibration, Uncertainty Estimation

TL;DR: We developed a model for joint training of deterministic and stochastic nodes in a neural network and showed it possesses desirable calibration properties.

Abstract:

Bayesian Neural Nets are proposed as flexible models that can provide calibrated uncertainty estimates for out-of-distribution data.
Due to the high dimensionality of BNN posteriors and the intractability of exact inference, numerous approximate inference techniques have been proposed. However, issues persist. Some approaches lack a proper Bayesian formulation while others result in inexpressive or uncalibrated posteriors, defeating the primary purpose of BNNs. Recently, subspace inference has been proposed to overcome these challenges by running the inference on a lower-dimensional subspace of network parameters. While achieving promising results, these methods are mathematically involved and therefore extending them to general architectures and problems is challenging. Here, we propose a new subspace inference method---called HybridBNN---that divides the network weights into deterministic and stochastic subsets before training. We develop an expectation-maximization algorithm for the joint inference of the posterior over the stochastic weights as well as the optimization of the deterministic ones. HybridBNN achieves competitive prediction and calibration performance on two regression and classification toy datasets and a benchmark dataset for in and out-of-domain distributions. The simplicity and flexibility of HybridBNN make it a favorable candidate for developing generic calibrated models.

Submission Number: 12