Linearised Laplace Inference in Networks with Normalisation Layers and the Neural g-Prior

Javier Antoran, James Urquhart Allingham, David Janz, Erik Daxberger, Eric Nalisnick, José Miguel Hernández-Lobato

Keywords: Linearised, Laplace, Evidence, Marginal Likelihood, Uncertainty, batch normalisation, layer normalisation, group normalisation, scale invariance, BNN, Bayesian Neural Network

TL;DR: We show how the Laplace marginal likelihood is broken in networks with normalisation layers and how to fix it. We also introduce a scale-invariant prior for NNs.

Abstract: We show that for neural networks (NN) with normalisation layers, i.e. batch norm, layer norm, or group norm, the Laplace model evidence does not approximate the volume of a posterior mode and is thus unsuitable for model selection. We instead propose to use the Laplace evidence of the linearized network, which is robust to the presence of these layers. We also identify heterogeneity in the scale of Jacobian entries corresponding to different weights. We ameliorate this issue by extending the scale-invariant g-prior to NNs. We demonstrate these methods on toy regression, and image classification with a CNN.