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

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.