Expressive yet Tractable Bayesian Deep Learning via Subnetwork Inference

Erik Daxberger, Eric Nalisnick, James Urquhart Allingham, Javier Antoran, José Miguel Hernández-Lobato

TL;DR: We propose a Bayesian deep learning method that does expressive inference over a carefully chosen _subnetwork_ within a neural network, and show that this works better than doing crude inference over the full network.

Abstract: Scaling Bayesian inference to the large parameter spaces of deep neural networks requires restrictive approximations. We propose performing inference over only a small subset of the model parameters while keeping all others as point estimates. This enables us to use expressive posterior approximations that are intractable in the full model. In particular, we develop a practical and scalable Bayesian deep learning method that first trains a point estimate, and then infers a full covariance Gaussian posterior approximation over a subnetwork. We propose a subnetwork selection procedure which aims to optimally preserve posterior uncertainty. Empirical studies demonstrate the effectiveness of our approach.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2010.14689/code)