PaperlayerView exact revision on arXiv

Source-linked AI summary

Bayesian Deep Learning via Subnetwork Inference

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

arXiv:2010.14689v4cs.LGstat.ML

TL;DR

Deep neural networks can be poorly calibrated and overconfident, while scalable Bayesian inference over large weight spaces often requires restrictive approximations. The paper infers a posterior over a selected weight subnetwork and keeps other weights fixed, reporting competitive uncertainty estimates against ensembles and less expressive methods. Its scope is constrained by the memory cost of the subnetwork Hessian.

  • Problem

    Deep neural networks can be poorly calibrated and overconfident, while exact Bayesian inference is intractable and scalable methods often rely on restrictive approximations.

    Source: §1 · p. 1

  • Method

    The method obtains a full-network MAP estimate, selects a small weight subnetwork for expressive posterior inference, and keeps the remaining weights fixed.

    Source: §1 · p. 1, §4 · p. 4

  • Results

    Subnetwork linearized Laplace is competitive with state-of-the-art uncertainty methods such as deep ensembles and retains more uncertainty than crude full-network approximations.

    Source: §9 · p. 9

  • Takeaways & Limitations

    Reliable predictive posteriors need not require inference over all weights, and modelling weight correlations is crucial to obtaining them.

    Source: §9 · p. 9

    Hessian storage constrains subnetwork size; for very large models such as modern transformers, tractable subnetworks may be a vanishingly small proportion of the weights.

    Source: §7 · p. 9

Abstract

from arXiv · show

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.

1. Introduction

Deep neural networks can be poorly calibrated and overconfident, while scalable Bayesian inference often relies on restrictive approximations. The paper proposes inference over a small weight subset and reports that this approach can outperform less expressive full-network methods and deep ensembles.

1. Introduction

  • Poor calibration and overconfidence are especially concerning under train–test distribution shift and for safety-critical decisions.The introduction names healthcare and autonomous driving as examples.Source: §1 · p. 1
  • The work asks whether a small subnetwork can preserve full-network model uncertainty and reports that its posterior predictive distribution can represent the full network's.It evaluates uncertainty calibration and robustness to distribution shift.Source: §1 · pp. 1–2
  • The framework infers a posterior over a small subset of weights while keeping all remaining weights deterministic, allowing more expressive posterior approximations.Its linearized Laplace instantiation fits a full-network MAP estimate before inferring a full-covariance Gaussian over the subnetwork.Source: §1 · p. 1
  • Approximation during subnetwork selection is empirically less harmful to the posterior predictive than approximation during inference.The selection strategy uses a diagonal approximation to scale the Wasserstein-distance criterion.Source: §1 · pp. 1–2
  • Subnetwork inference outperforms popular Bayesian methods with less expressive full-network inference and deep ensembles on uncertainty calibration and distribution-shift robustness benchmarks.The experiments evaluate benchmarks for both properties.Source: §1 · p. 2

2. Subnetwork Posterior Approximation

Bayesian neural networks represent model uncertainty with a posterior over weights, but exact inference is intractable and faithful approximation is computationally challenging. The paper approximates that posterior by inferring over a small subset of weights and fixing the rest at point values.

2. Subnetwork Posterior Approximation

  • A weight posterior induces predictions by marginalizing the likelihood over plausible weights, but exact neural-network posterior inference is intractable.High-dimensional weights also make faithful approximation computationally challenging.Source: §2 · p. 2
  • A fully factorized posterior is a commonly used crude approximation, but it has been reported to suffer from severe pathologies.This motivates seeking richer approximations without inferring over every weight.Source: §2 · p. 2
  • Rather than assume an expressive posterior must cover all D weights, the method performs inference over S ≪ D weights.The motivation includes overparameterization and evidence that inference can work in low-dimensional projections or a network's last layer.Source: §2 · p. 2
  • Subnetwork inference keeps selected weights probabilistic while fixing the remaining weights at their point values, making richer posterior approximations feasible on the smaller subset.The framework approximates the subnetwork posterior with q(wS).Source: §2 · p. 2
  • Subnetwork inference removes selected weights' variances while retaining their means, unlike pruning methods that set weights themselves to zero.The paper frames these approaches as complementary: subnetwork inference targets predictive uncertainty, while weight pruning targets predictive mean and accuracy.Source: §2 · p. 2

3. Background: Linearized Laplace

Linearized Laplace approximates a neural network's weight posterior around a point estimate with a full-covariance Gaussian. Linearizing the network connects this approximation to a generalized linear model, while full covariance over modern networks is computationally intractable.

3. Background: Linearized Laplace

  • Linearized Laplace approximates the weight posterior using a local linearization, with the GGN serving as a practical replacement for the Hessian.For Gaussian likelihood regression, the approximation is exact for the linearized network; classification uses a probit approximation for the predictive integral.Source: §3 · pp. 2–3
  • A second-order approximation around the MAP estimate yields a full-covariance Gaussian posterior with covariance H^-1.The GGN is commonly substituted for the Hessian in practice.Source: §3 · p. 3
  • The local linearized network acts as a generalized linear model whose Jacobian supplies basis functions for prediction.The paper reports that this GLM can outperform the corresponding Bayesian neural network with a GGN-Laplace posterior.Source: §3 · p. 3
  • For Gaussian-likelihood regression, the predictive mean and classification boundaries match those of the MAP-estimated network.The paper gives this property for the closed-form predictive expressions.Source: §3 · p. 3
  • Storing full D×D covariance for modern networks is computationally intractable, and diagonal or block-diagonal alternatives reduce predictive performance.The passage names diagonal and block-diagonal storage as cheaper approximations.Source: §3 · p. 4

4. Linearized Laplace Subnetwork Inference

Subnetwork linearized Laplace first obtains a point estimate of all weights, then selects a small subset for full-covariance Gaussian inference while keeping other weights fixed. This reduces the storage and inversion burden of Laplace inference for large networks.

4. Linearized Laplace Subnetwork Inference

  • The procedure trains or reuses a MAP estimate, selects a small weight subset, then performs full-covariance GGN-Laplace inference over that subset while fixing other weights at MAP values.The framework combines full Laplace inference on selected weights with MAP inference on the rest.Source: §4 · p. 4
  • The prior precision is scaled as λ_S = λ·S/D to preserve predictive-variance magnitude.The selected subnetwork weights are those identified in the preceding selection step.Source: §4 · p. 4
  • O(D²) storage and O(D³) inversion for the full GGN become O(S²) and O(S³) for the subnetwork GGN; subnetworks contain less than 1% of weights in experiments.Quadratic and cubic scaling in S remains necessary to capture weight correlations.Source: §4 · p. 4
  • Prediction locally linearizes the network while fixing weights outside the selected subnetwork at their MAP values.Source: §4 · p. 4

5. Subnetwork Selection

Subnetwork selection approximates the full-network posterior in weight space using a tractable Wasserstein-based criterion. Experiments report that diagonal approximations during selection are less harmful than during inference, while Wasserstein-selected subnetworks preserve richer uncertainty with fewer parameters.

5. Subnetwork Selection

  • The Gaussian posterior is a crude approximation when predictions are not made with the linearized model.Source: §5 · p. 4
  • Because function-space reasoning is difficult and common distances are ill-defined for Dirac-delta posteriors, selection minimizes squared 2-Wasserstein distance between weight posteriors.The exact full-network posterior is also intractable.Source: §5 · pp. 4–5
  • An independence assumption simplifies combinatorial subset selection, yielding a criterion that chooses the S weights with highest marginal variances rather than largest magnitudes.The variance-based criterion targets predictive uncertainty.Source: §5 · p. 5
  • Direct marginal-variance computation is intractable, so the method uses diagonal Laplace or diagonal SWAG approximations; mean-field inference has larger overhead.Source: §5 · p. 5
  • Diagonal approximation during selection is less harmful than during inference and outperforms random selection; Wasserstein selection retains richer predictive uncertainty at smaller parameter counts.Source: §5 · p. 5, Figure 2 · p. 5

6. Experiments

Across image-classification shifts, subnetwork inference yields stronger uncertainty calibration and robustness than other approaches, while experiments also show that selected subnetworks can preserve predictive uncertainty with fewer inferred weights. The experiments compare these results across synthetic regression, tabular regression, and image classification.

6. Experiments

  • The experiments assess uncertainty methods on toy regression, tabular regression, and ResNet-18 image classification, comparing selected subnetworks with full-network approximations and established baselines.The study includes full-covariance, diagonal, final-layer, and MAP comparisons in the synthetic regression setting.Source: §6 · p. 5, §6.1 · p. 6

6.1. How does Subnetwork Inference preserve Posterior Predictive Uncertainty?

  • On synthetic regression, subnetwork inference preserves more full-network uncertainty between data clusters than diagonal or final-layer inference, and Wasserstein selection preserves more with smaller subnetworks.Both random and Wasserstein subnetworks retain between-cluster uncertainty; the comparison holds while inferring over fewer weights.Source: §6.1 · p. 6

6.3. Image Classification under Distribution Shift

  • Subnetwork inference delivers better uncertainty calibration and robustness under image distribution shift than popular alternatives, while random subnetworks perform substantially worse than Wasserstein-selected ones.Figure 4 reports rotated-MNIST and corrupted-CIFAR results; on corrupted CIFAR, the method has the best log-likelihood at all corruption levels, and on rotated MNIST it exceeds almost all baselines, with VOGN the exception.Source: Figure 4 · p. 7, §6.3 · p. 8

6.2. Subnetwork Inference in Large Models vs Full Inference over Small Models

  • On tabular datasets, larger networks with Wasserstein-selected subnetworks perform best even when inference uses as many weights as smaller networks, except on Wine-gap and Protein-gap.Mean test log-likelihood measures both accuracy and uncertainty; the authors suggest that allocating inference to informative variances and covariances helps, and note that compute may favor subnetworks in larger networks over full inference in smaller ones.Source: §6.2 · p. 7

6.3. Image Classification under Distribution Shift

  • With ResNet-18, subnetworks as small as 100–40K weights retain strong calibration as performance degrades smoothly with size, using the same hyperparameters across sizes.The figure compares log-likelihoods on rotated MNIST and corrupted CIFAR10 against ensembles and Diagonal Laplace, alongside covariance-memory footprints.Source: Figure 5 · p. 8
  • Subnetwork inference reduces covariance storage to about 4.0 MB for a 1K-sized subnetwork while remaining competitive with deep ensembles, versus about 500 TB for full ResNet-18 covariance.The comparison supports applying otherwise intractable inference methods to larger networks.Source: §6.3 · p. 8

7. Scope and Limitations

The method has computational and tuning boundaries: Jacobian computation is challenging for multi-output models, Hessian storage limits subnetwork size, and prior precision selection depends on a sensible full-network setting.

7. Scope and Limitations

  • Linearized Laplace requires as many backward passes as model outputs, limiting its use for semantic segmentation and classification with many classes.Other inference methods without this limitation can still be used within the subnetwork framework.Source: §7 · p. 8
  • Prior precision strongly affects Laplace performance, but updating it for subnetworks relies on a sensible full-network setting, which is often difficult to obtain.The authors identify cross-validation with the subnetwork approximation as the current best approach.Source: §7 · p. 8
  • Hessian storage limits subnetwork size; a 40K-weight Hessian takes about 6.4 GB, making tractable subnetworks a vanishingly small share of very large models such as transformers.Finding better subnetwork selection strategies remains future work.Source: §7 · p. 9

8. Related Work

The related work situates subnetwork inference among Bayesian neural-network approximations, neural linear methods, pruning, and inference over low-dimensional weight spaces. The paper distinguishes its focus on predictive uncertainty and axis-aligned weight subsets from approaches targeting predictive accuracy or implicit subspaces.

Bayesian Deep Learning

  • Deep Bayesian methods use surrogate posteriors, often Gaussians, because sampling methods are difficult to scale; mean-field approximations scale but have limited expressivity and empirical under-performance.The paper reports that scaling correlation-aware approximations by lowering weight-space dimensionality outperforms diagonal approximations, contrasting with Farquhar et al. (2020).Source: §8 · p. 9

Neural Linear Methods

  • Neural linear methods perform inference only over the final network layer while fixing earlier layers, making them a special case of subnetwork inference.The cited examples include Snoek et al. (2015) and Riquelme et al. (2018).Source: §8 · p. 9

Inference over Subspaces

  • Unlike pruning, which removes weights to preserve predictive means, subnetwork inference retains weights but fixes some variances to preserve predictive uncertainty.The paper identifies neural-network pruning with selecting weights needed for accurate predictions and contrasts it with selecting weights relevant to uncertainty.Source: §8 · p. 9
  • Prior subspace approaches use SGD-trajectory principal components or low-rank posterior parameterizations; this work instead identifies axis-aligned subsets relevant to predictive uncertainty.The cited subspace and low-rank approaches include Izmailov et al. (2019), Rossi et al. (2020), Swiatkowski et al. (2020), and Dusenberry et al. (2020).Source: §8 · p. 9

9. Conclusion

The conclusion argues that subnetwork linearized Laplace is competitive with state-of-the-art uncertainty methods, while inference over a subset of weights can retain reliable predictive posteriors. Additional image experiments examine its memory efficiency, scalability, and OOD rejection.

9. Conclusion

  • The conclusion identifies weight correlations as crucial, unimodal approximations as potentially competitive with multi-mode approximations, and inference over all weights as unnecessary for reliable predictive posteriors.The proposed framework keeps non-subnetwork weights deterministic and decouples computational cost from total model size, allowing a trade-off between cost and approximation quality.Source: §9 · p. 9
  • The method is competitive with state-of-the-art uncertainty methods, retains more uncertainty than crude full-network approximations, and can be applied post-hoc to pre-trained models.The authors also report that it allows larger networks without sacrificing uncertainty-estimate quality.Source: §9 · p. 9

A.1. Comparing the Parameter Efficiency of Subnetwork Linearized Laplace with Deep Ensembles

  • At large rotations and corruptions, subnetwork linearized Laplace achieves greater log-likelihood than the asymptotic value of deep ensembles, whose performance plateaus around 15 members.The comparison uses memory equivalent to around 163 ensemble elements; the authors note that large ensembles require a forward pass per element, whereas linearized Laplace requires one backward pass per output dimension and one forward pass.Source: §A.1 · p. 13

A.2. Scalability of Subnetwork Linearised Laplace in the number of Weights

  • With 0.167% of ResNet-50’s parameters, subnetwork inference improves on MAP in log-likelihood and calibration metrics.The subnetwork contains 39,190 of 23,466,560 parameters.Source: §A.2 · p. 13, Figure 7 · p. 14

A.3. Out-of-Distribution Rejection

  • For OOD detection, methods perform similarly on CIFAR-SVHN except that ensembles do best; on MNIST-Fashion, SWAG performs best, followed by subnetwork linearized Laplace and ensembles.In the rejection-classification setting, ensembles lead on CIFAR-SVHN and SWAG performs worse; on MNIST-Fashion, SWAG leads, followed by subnetwork linearized Laplace, while other methods fail to distinguish very uncertain in-distribution data from low-uncertainty OOD points.Source: §A.3 · pp. 13–14

A.4. Additional Rotation and Corruption Results

The appendix adds expected calibration error and Brier score to the image-classification results and lists results by rotation and corruption condition.

A.4. Additional Rotation and Corruption Results

B. Derivations for the Wasserstein Pruning Objective

The appendix derives the Wasserstein pruning objective and the subnetwork prior-precision adjustment, then describes toy, UCI, and image-experiment procedures.

B.1. Derivation for (22)

  • For the linearized model, the posterior is Gaussian or approximately Gaussian, and the subnetwork posterior can be treated as a degenerate Gaussian with zero covariance for point-estimated weights.The derivation expresses the pruning objective using the squared 2-Wasserstein distance between Gaussian distributions.Source: §B.1 · p. 19

B.2. Derivation for (23)

  • The diagonal pruning rule represents each weight’s inclusion with a mask value of 1 or 0.The derivation identifies the full-network and subnetwork covariance matrices and simplifies the Wasserstein objective under the diagonal approximation.Source: §B.1 · p. 19, §B.2 · p. 19

C. Updating the prior precision for uncertainty estimation with subnetworks

  • The subnetwork prior precision is scaled by S/D to keep the magnitude of its Jacobian-based kernel similar to the full-network kernel.The subnetwork feature expansion removes D − S terms from the full-network sum.Source: C. Updating the prior precision for… · pp. 19–20

D.1. Toy Experiments

  • Toy experiments compare uncertainty estimates after MAP training, with approaches sharing the same predictive mean and prior precision scaled as λS = λ · S/D.The toy setup uses a two-hidden-layer network with 50 ReLU units per layer; inference uses GGN Laplace over weights, and final-layer inference retains only final-layer covariance entries.Source: §D.1 · p. 20

D.2. UCI Experiments

  • UCI experiments obtain MAP estimates, select prior precision on validation data, and compare full-network GGN Laplace with diagonal-Hessian Wasserstein pruning at matched variance counts.The validation-set log-likelihood selects prior precision from a grid, and the linearized predictive mean is shared by networks with the same number of weights.Source: §D.2 · p. 20

D.3. Image Experiments

  • Image experiments apply the compared inference methods to trained ResNet-18 and ResNet-50 models, with prior precision tuned on a validation set for Ours, Diag-Laplace, and SWAG.The appendix also specifies the image-model training procedure and architecture modifications, including replacing the first 7 × 7 convolution with a 3 × 3 convolution and removing the first max-pooling layer.Source: §D.3 · p. 21

D.4. Datasets

The appendix identifies the datasets used for toy, tabular-regression, and image experiments and gives the rationale and split procedure for the tabular datasets.

D.4. Datasets

  • The 1D toy dataset comes from Antorán et al. (2020), while Wine and Protein come from the UCI repository and Kin8nm from OpenML.Standard splits use 90% of the data for training and 10% for validation; gap splits remove the middle 33% along each input dimension for validation.Source: §D.4 · pp. 21–22
  • Wine, Kin8nm, Protein, and their gap variants were selected because performance depended more on uncertainty quality; on most other common UCI regression datasets, MAP predictions outperformed Bayesian ensembles in log-likelihood.The gap splits are designed to test out-of-distribution uncertainty.Source: §D.4 · pp. 21–22
  • The appendix provides summary tables for the tabular-regression datasets and image datasets used in the experiments.The image-dataset table caption specifies that test and train set sizes are shown in brackets.Source: Table 22 · p. 22, Table 23 · p. 22

Related papers

Loading 2010.14689v4…