One Simple Trick to Fix Your Bayesian Neural Network

TL;DR

This paper demonstrates that using Leaky ReLU activations in Bayesian neural networks results in more Gaussian-like posteriors and improved calibration, addressing limitations of mean-field variational inference with ReLU activations.

Contribution

It provides a theoretical explanation and empirical evidence that Leaky ReLU activations improve posterior fitting and calibration in Bayesian neural networks.

Findings

Leaky ReLU leads to more Gaussian-like posteriors.

Leaky ReLU reduces expected calibration error compared to ReLU.

ReLU activations induce posteriors that are hard to fit with MFVI.

Abstract

One of the most popular estimation methods in Bayesian neural networks (BNN) is mean-field variational inference (MFVI). In this work, we show that neural networks with ReLU activation function induce posteriors, that are hard to fit with MFVI. We provide a theoretical justification for this phenomenon, study it empirically, and report the results of a series of experiments to investigate the effect of activation function on the calibration of BNNs. We find that using Leaky ReLU activations leads to more Gaussian-like weight posteriors and achieves a lower expected calibration error (ECE) than its ReLU-based counterpart.