[Paper Review] Bayesian Deep Learning via Subnetwork Inference

Erik Daxberger, Eric Nalisnick+3 more|UvA-DARE (University of Amsterdam)|Oct 28, 2020

Gaussian Processes and Bayesian Inference79 references11 citations

TL;DR

This paper proposes subnetwork inference, a scalable Bayesian deep learning method that performs full-covariance Gaussian posterior approximation only over a small, strategically selected subset of neural network weights, while keeping the rest as point estimates. By using linearized Laplace approximation on the subnetwork and a Wasserstein-based selection strategy, the method achieves accurate uncertainty quantification comparable to deep ensembles, outperforming full-network Bayesian methods with less expressive posteriors.

ABSTRACT

The Bayesian paradigm has the potential to solve core issues of deep neural networks such as poor calibration and data inefficiency. Alas, scaling Bayesian inference to large weight spaces often requires restrictive approximations. In this work, we show that it suffices to perform inference over a small subset of model weights in order to obtain accurate predictive posteriors. The other weights are kept as point estimates. This subnetwork inference framework enables us to use expressive, otherwise intractable, posterior approximations over such subsets. In particular, we implement subnetwork linearized Laplace as a simple, scalable Bayesian deep learning method: We first obtain a MAP estimate of all weights and then infer a full-covariance Gaussian posterior over a subnetwork using the linearized Laplace approximation. We propose a subnetwork selection strategy that aims to maximally preserve the model's predictive uncertainty. Empirically, our approach compares favorably to ensembles and less expressive posterior approximations over full networks. Our proposed subnetwork (linearized) Laplace method is implemented within the laplace PyTorch library at https://github.com/AlexImmer/Laplace.

Motivation & Objective

• To address the intractability of Bayesian inference in large deep neural networks due to high-dimensional weight spaces.
• To improve uncertainty calibration and robustness to distribution shift in deep learning models.
• To reduce computational cost while preserving predictive uncertainty by focusing Bayesian inference on a small, critical subset of weights.
• To demonstrate that expressive posterior approximations over subnetworks can match or exceed the performance of full-network Bayesian methods and deep ensembles.

Proposed method

• The method first computes a maximum a posteriori (MAP) estimate for all weights in the full neural network.
• It then selects a small, representative subnetwork of weights using a Wasserstein distance-based criterion to preserve predictive uncertainty.
• A full-covariance Gaussian posterior is inferred over the selected subnetwork using the linearized Laplace approximation.
• The remaining weights are kept as fixed point estimates, enabling efficient prediction via marginalization over the subnetwork posterior.
• During subnetwork selection, a diagonal approximation is used for scalability, minimizing impact on posterior quality.
• The final prediction combines the full network’s architecture with a Bayesian posterior over the subnetwork, enabling uncertainty-aware inference.

Experimental results

Research questions

• RQ1 Can Bayesian inference over a small subset of weights preserve the predictive uncertainty of a full neural network?
• RQ2 Does subnetwork inference with expressive posterior approximations outperform less expressive full-network Bayesian methods?
• RQ3 Can a subnetwork-based approach match or exceed the uncertainty calibration of deep ensembles?
• RQ4 How does the choice of subnetwork selection strategy affect posterior quality and uncertainty estimation?
• RQ5 Is it more efficient and effective to make approximations during subnetwork selection rather than during posterior inference?

Key findings

• The proposed subnetwork (linearized) Laplace method achieves uncertainty calibration and robustness to distribution shift comparable to deep ensembles, despite using only a small fraction of the weights for Bayesian inference.
• The method outperforms full-network Bayesian methods with diagonal or factorized posteriors, demonstrating that expressive posterior approximations on subnetworks are more effective than less expressive ones on the full network.
• Subnetwork selection based on Wasserstein distance between full and subnetwork posteriors significantly improves uncertainty preservation, even when using a diagonal approximation during selection.
• The approach maintains high predictive accuracy while enabling full-covariance posterior inference over a low-dimensional subspace, making it computationally feasible for large models.
• Empirical results on MNIST, CIFAR10, Fashion-MNIST, SVHN, and tabular UCI datasets confirm superior calibration and robustness compared to baseline Bayesian and ensemble methods.
• The method is efficiently implemented in the Laplace PyTorch library, enabling reproducibility and integration into existing deep learning workflows.