A Framework for Variational Inference of Lightweight Bayesian Neural Networks with Heteroscedastic Uncertainties

TL;DR

This paper introduces a simple, sampling-free variational inference framework for lightweight Bayesian Neural Networks that effectively models heteroscedastic uncertainties by embedding them into parameter variances.

Contribution

It proposes a novel approach to model heteroscedastic uncertainties within BNNs by embedding them into parameter variances, enhancing lightweight network performance.

Findings

Improved predictive performance for lightweight BNNs.

Effective modeling of heteroscedastic aleatoric and epistemic uncertainties.

Sampling-free variational inference framework.

Abstract

Obtaining heteroscedastic predictive uncertainties from a Bayesian Neural Network (BNN) is vital to many applications. Often, heteroscedastic aleatoric uncertainties are learned as outputs of the BNN in addition to the predictive means, however doing so may necessitate adding more learnable parameters to the network. In this work, we demonstrate that both the heteroscedastic aleatoric and epistemic variance can be embedded into the variances of learned BNN parameters, improving predictive performance for lightweight networks. By complementing this approach with a moment propagation approach to inference, we introduce a relatively simple framework for sampling-free variational inference suitable for lightweight BNNs.