Subspace Inference for Bayesian Deep Learning

TL;DR

This paper introduces a subspace inference approach for Bayesian deep learning, enabling scalable Bayesian inference in neural networks by focusing on low-dimensional parameter subspaces derived from SGD trajectories.

Contribution

It proposes constructing low-dimensional subspaces of neural network parameters, such as principal components of SGD trajectories, to facilitate Bayesian inference methods like elliptical slice sampling and variational inference.

Findings

Bayesian model averaging in subspaces yields accurate predictions.

The method produces well-calibrated uncertainty estimates.

Applicable to both regression and image classification tasks.

Abstract

Bayesian inference was once a gold standard for learning with neural networks, providing accurate full predictive distributions and well calibrated uncertainty. However, scaling Bayesian inference techniques to deep neural networks is challenging due to the high dimensionality of the parameter space. In this paper, we construct low-dimensional subspaces of parameter space, such as the first principal components of the stochastic gradient descent (SGD) trajectory, which contain diverse sets of high performing models. In these subspaces, we are able to apply elliptical slice sampling and variational inference, which struggle in the full parameter space. We show that Bayesian model averaging over the induced posterior in these subspaces produces accurate predictions and well calibrated predictive uncertainty for both regression and image classification.