Optimal Subspace Inference for the Laplace Approximation of Bayesian Neural Networks
This work addresses the computational bottleneck of uncertainty estimation in Bayesian neural networks, offering a more efficient method for practitioners, though it is incremental as it builds on existing subspace inference approaches.
The paper tackles the problem of efficiently approximating Bayesian inference in neural networks by deriving the optimal subspace model for the Laplace approximation, showing that often less than 1% of parameters are sufficient for reliable uncertainty quantification.
Subspace inference for neural networks assumes that a subspace of their parameter space suffices to produce a reliable uncertainty quantification. In this work, we mathematically derive the optimal subspace model to a Bayesian inference scenario based on the Laplace approximation. We demonstrate empirically that, in the optimal case, often a fraction of parameters less than 1% is sufficient to obtain a reliable estimate of the full Laplace approximation. Since the optimal solution is derived, we can evaluate all other subspace models against a baseline. In addition, we give an approximation of our method that is applicable to larger problem settings, in which the optimal solution is not computable, 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 different subspace models even if the exact Laplace approximation is unknown.