Riemannian Laplace approximations for Bayesian neural networks
Federico Bergamin Pablo Moreno-Muñoz Søren Hauberg Georgios Arvanitidis
Bayesian neural networks often approximate the weight-posterior with a
Gaussian distribution. However, practical posteriors are often, even locally,
highly non-Gaussian, and empirical performance deteriorates. We propose a
simple parametric approximate posterior that adapts to the shape of the true
posterior through a Riemannian metric that is determined by the log-posterior
gradient. We develop a Riemannian Laplace approximation where samples naturally
fall into weight-regions with low negative log-posterior. We show that these
samples can be drawn by solving a system of ordinary differential equations,
which can be done efficiently by leveraging the structure of the Riemannian
metric and automatic differentiation. Empirically, we demonstrate that our
approach consistently improves over the conventional Laplace approximation
across tasks. We further show that, unlike the conventional Laplace
approximation, our method is not overly sensitive to the choice of prior, which
alleviates a practical pitfall of current approaches.