Riemannian Laplace approximations for Bayesian neural networks

Abstract

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.

arXiv comment28 pages, 12 figures. Under submission

References in corpus (15)

• Adam: A Method for Stochastic Optimization
• SciPy 1.0--Fundamental Algorithms for Scientific Computing in Python
• Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning
• Visualizing the Loss Landscape of Neural Nets
• Optimizing Neural Networks with Kronecker-factored Approximate Curvature
• Eigenvalues of the Hessian in Deep Learning: Singularity and Beyond
• Variational Boosting: Iteratively Refining Posterior Approximations
• Practical Gauss-Newton Optimisation for Deep Learning
• Metrics for Probabilistic Geometries
• Fast and Robust Shortest Paths on Manifolds Learned from Data
• Probabilistic Solutions to Differential Equations and their Application to Riemannian Statistics
• Do Bayesian Neural Networks Need To Be Fully Stochastic?
• Adapting the Linearised Laplace Model Evidence for Modern Deep Learning
• Posterior Refinement Improves Sample Efficiency in Bayesian Neural Networks
• The Geometry of Neural Nets' Parameter Spaces Under Reparametrization