Low Rank Based Subspace Inference for the Laplace Approximation of Bayesian Neural Networks

TL;DR

This paper introduces a low-rank subspace inference method for Bayesian neural networks using Laplace approximation, demonstrating that reduced covariance matrices can effectively approximate full models and improve scalability.

Contribution

It derives an optimal subspace model for Laplace approximation in Bayesian neural networks using low-rank techniques and provides a scalable approximation method that outperforms existing models.

Findings

Reduced covariance matrices closely match full Laplace approximation.

The proposed scalable approximation outperforms previous subspace models.

A new metric for comparing subspace approximation quality is introduced.

Abstract

Subspace inference for neural networks assumes that a subspace of their parameter space suffices to produce a reliable uncertainty quantification. In this work, we underpin the validity of this assumption by using low rank techniques. We derive an expression for a subspace model to a Bayesian inference scenario based on the Laplace approximation that is, in a certain sense, optimal given a specific dataset. We empirically show that a Laplace approximation constructed with a dimensionally reduced covariance matrix closely matches the full Laplace approximation obtained using the exact covariance matrix. Where feasible, this subspace model can serve as a baseline for benchmarking the performance of subspace models. In addition, we provide a scalable approximation of this subspace construction that is usable in practice and compare it to existing subspace models from the literature. In general, our approximation scheme outperforms previous work. Furthermore, we present a metric to qualitatively compare the approximation quality of different subspace models even if the exact Laplace approximation is unknown.