PaperlayerView exact revision on arXiv

Source-linked AI summary

Laplace Redux -- Effortless Bayesian Deep Learning

Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, Philipp Hennig

arXiv:2106.14806v3cs.LGstat.ML

TL;DR

Bayesian deep learning offers uncertainty and model-selection benefits, but the Laplace approximation has been underused amid concerns about cost, implementation, and performance. The paper reviews scalable variants, introduces the PyTorch library laplace, and evaluates them extensively. It finds the LA competitive with alternative posterior approximations at typically much lower computational cost, supporting wider practical adoption.

  • Problem

    Bayesian neural networks offer useful uncertainty and model-selection capabilities, but the Laplace approximation remains less popular amid concerns that it is difficult, expensive, or inferior.

  • Method

    The paper surveys scalable LA variants, introduces the easy-to-use PyTorch library laplace, and explains post-hoc construction from pretrained networks.

  • Results

    The LA is competitive with alternative posterior approximations while typically requiring much lower computational cost.

  • Takeaways & Limitations

    Simple LA variants can be applied post-hoc to pretrained models, including transformer models, to improve predictive uncertainty.

  • Takeaways & Limitations

    Naive Hessian computation is infeasible because the likelihood Hessian term scales quadratically with the number of network parameters.

Abstract

from arXiv · show

Bayesian formulations of deep learning have been shown to have compelling theoretical properties and offer practical functional benefits, such as improved predictive uncertainty quantification and model selection. The Laplace approximation (LA) is a classic, and arguably the simplest family of approximations for the intractable posteriors of deep neural networks. Yet, despite its simplicity, the LA is not as popular as alternatives like variational Bayes or deep ensembles. This may be due to assumptions that the LA is expensive due to the involved Hessian computation, that it is difficult to implement, or that it yields inferior results. In this work we show that these are misconceptions: we (i) review the range of variants of the LA including versions with minimal cost overhead; (ii) introduce "laplace", an easy-to-use software library for PyTorch offering user-friendly access to all major flavors of the LA; and (iii) demonstrate through extensive experiments that the LA is competitive with more popular alternatives in terms of performance, while excelling in terms of computational cost. We hope that this work will serve as a catalyst to a wider adoption of the LA in practical deep learning, including in domains where Bayesian approaches are not typically considered at the moment.

1 Introduction

The paper argues that misconceptions about the Laplace approximation’s implementation cost and performance have limited its adoption, despite Bayesian benefits for uncertainty and model selection. It surveys scalable variants, introduces the PyTorch library laplace, and reports competitive empirical performance at low cost.

  • Bayesian modeling addresses neural networks’ poor calibration, catastrophic forgetting, and difficult architecture or hyperparameter selection through uncertainty, continual learning, and model selection.
  • The Laplace approximation locally fits a Gaussian posterior around a MAP estimate using local curvature, enabling practical Bayesian inference from standard neural-network training.
  • The paper argues that concerns about the LA being hard to implement, expensive, or inferior because of its local nature are misconceptions.
  • The authors survey scalable LA variants and introduce laplace, an easy-to-use PyTorch library implementing a wide range of variants.
  • The empirical study finds the LA competitive with alternative approaches while emphasizing its simplicity and low computational cost.

2 The Laplace Approximation in Deep Learning

The LA converts standard MAP-trained neural networks into approximate Bayesian models by fitting a local Gaussian posterior and using it for prediction or evidence-based model selection. Its practical bottleneck is the naive Hessian, whose quadratic parameter scaling requires scalable approximations.

  • The LA supports probabilistic predictions through an approximate posterior and model selection through an approximate marginal likelihood or evidence.
  • The LA uses a second-order expansion around θMAP to construct a Gaussian approximation to the otherwise intractable posterior.
  • Standard regularized empirical risk minimization yields the MAP estimate because empirical losses correspond to negative log-likelihoods and regularizers to negative log-priors.
  • Under weight decay, the prior is assumed to be a zero-mean Gaussian with variance γ2I unless stated otherwise.
  • A naive Hessian implementation is infeasible because its likelihood term scales quadratically with networks containing millions or billions of parameters.

1 Inference over all Weights or Subsets of Weights

The LA can scale by restricting probabilistic inference to a subset of weights, including only the last linear layer while keeping earlier features fixed at their MAP values. This provides a cost-effective option for large neural networks.

  • The LA can also treat all weights probabilistically when suitable Hessian approximations are used.
  • Per-layer or per-parameter weight decay corresponds to a more general Gaussian prior whose Hessian remains diagonal and constant.
  • The subnetwork LA applies probabilistic inference to only a subset of parameters and leaves the remaining parameters fixed at their MAP-estimated values.
  • A last-layer LA fixes the feature extractor and applies the approximation only to the final linear layer, offering a cost-effective approach for deep networks.

2 Hessian Approximations and Their Factorizations

Because full Fisher and generalized Gauss–Newton matrices remain quadratically large, practical deep LAs use factorization or low-rank assumptions. Options range from diagonal approximations to more expressive KFAC and low-rank variants.

  • Positive semi-definite Fisher and generalized Gauss–Newton approximations address potentially indefinite neural-network likelihood Hessians.
  • For common likelihoods, the Fisher and GGN are equivalent and have become the default curvature choices in deep LAs.
  • Further factorization is needed because the Fisher and GGN remain quadratically large.
  • Diagonal factorization is lightweight, whereas KFAC factorizes each within-layer Fisher into a Kronecker product of two smaller matrices.
  • KFAC can be refined with low-rank factor approximations, and recent work also enables low-rank Hessian or Fisher approximations.

3 Hyperparameter Tuning

The LA’s prior and likelihood hyperparameters affect performance, and marginal likelihood maximization offers a principled tuning alternative to cross-validation.

  • Marginal likelihood maximization provides a principled alternative to cross-validation for tuning LA hyperparameters.The prior variance γ2 is typically tuned because it can materially affect inference performance.

4 Approximate Predictive Distribution

The paper develops predictive approximations for LA posteriors, using linearization to obtain tractable output distributions and specialized approximations for classification.

  • The LA predictive integrates p(y | f(x*)) over the posterior p(θ | D), which is generally intractable.
  • Monte Carlo integration is general but can perform poorly with GGN or Fisher Hessians because the predictive may be inconsistent with the Hessian approximation.
  • For last-layer LA, the Hessian coincides with the GGN, making the linearized predictive exact; classification may still require an additional approximation.
  • Linearizing the network turns the LA posterior into a Gaussian distribution over outputs, reducing computation from parameter space to output space.
  • For Gaussian regression likelihoods, the predictive distribution has an analytic Gaussian form with covariance J(x*)⊺ΣJ(x*) + σ2I.

3 laplace: A Toolkit for Deep Laplace Approximations

The laplace library makes scalable Laplace approximations accessible in PyTorch by providing modular implementations, practical defaults, and configurable inference workflows.

  • laplace addresses the lack of a complete standardized implementation by offering an easy-to-use, extensible PyTorch library for scalable deep-neural-network LAs.
  • A typical workflow loads a pre-trained model, fits an LA, tunes prior precision by cross-validation, and predicts with a probit approximation.
  • The library supports posterior construction, marginal likelihood computation, and posterior prediction through efficient implementations and external second-order backends.
  • The recommended default is a post-hoc last-layer KFAC LA using a GGN Hessian approximation and empirical-Bayes hyperparameter tuning.
  • The library’s Hessian-factorization choices are constrained by the capabilities of its BackPACK and ASDL backends.

4 Experiments

Experiments compare LA design choices, Bayesian baselines, robustness, calibration, computational cost, and downstream applications across standard and real-world distribution shifts. Overall, LA variants are competitive or superior while adding minimal computational overhead.

  • 4 Experiments: The experiments evaluate LA variants across in-distribution, dataset-shift, and OOD settings, multiple architectures, data modalities, and downstream tasks.
  • 4.1 Choosing the Right Laplace Approximation: Post-hoc last-layer LA typically outperforms all-weight LA because it reduces underfitting and is significantly cheaper.
  • 4.1 Choosing the Right Laplace Approximation: More expressive covariance improves performance; KFAC balances expressiveness and speed, while diagonal approximations perform significantly worse.
  • 4.1 Choosing the Right Laplace Approximation: The probit predictive offers the best combined last-layer LA trade-off among OOD performance and cost, while MC can improve OOD detection at higher cost.
  • 4.2 Predictive Uncertainty Quantification: LA and LA* are competitive with or superior to Bayesian baselines for calibration and OOD detection, with LA* especially reducing MAP overconfidence.
  • 4.3 Realistic Distribution Shift: LA incurs negligible overhead over MAP, whereas DE, VB, and CSGHMC cost two to five times more and SWAG exceeds MAP substantially at prediction.
  • 4.3 Realistic Distribution Shift: LA is much better calibrated than MAP and competitive with temperature scaling and deep ensembles, especially on OOD splits of five WILDS datasets.
  • 4.4 Further Applications: After ten Permuted-MNIST tasks, KFAC LA reaches around 96% accuracy and outperforms other Bayesian continual-learning approaches.

5 Related Work

The LA is a local Gaussian approximation, sharing mode-local limitations with other Gaussian methods while avoiding some of their computational burdens. In continual learning, Bayesian approaches outperform MAP, with LA-KFAC leading closely followed by VOGN.

  • Continual learning: MAP fails catastrophically as more Permuted-MNIST tasks are added, whereas Bayesian approaches substantially outperform MAP.The comparison concerns continual-learning performance as the number of tasks increases.
  • Continual learning: LA-KFAC performs best on the Permuted-MNIST continual-learning benchmark, closely followed by VOGN.These results are reported for Figure 7's comparison of continual-learning methods.
  • Local approximations: LA, mean-field variational inference, and SWAG each capture only local information by covering a single posterior mode.SWAG forms a diagonal-plus-low-rank Gaussian from SGD iterates and requires storing many network copies plus a costly batch-normalization heuristic at test time.
  • Local approximations: MCMC methods and deep ensembles can explore several posterior modes, but prior work suggests single-mode approximations may be less limiting in practice than expected.The cited explanation attributes this possibility to the complex nonlinear relationship between neural-network parameters and outputs.

6 Conclusion

The paper presents the Laplace approximation as a simple, competitive, and versatile approach to Bayesian deep learning. It combines efficient variants and a PyTorch library to provide drop-in Bayesian functionality across neural-network settings.

  • Conclusion: The paper argues that the Laplace approximation is simple, competitive, versatile, and deserving of wider adoption in Bayesian deep learning.This conclusion is based on the paper's review and evaluation of LA variants.
  • Conclusion: The review covers LA variants with minimal cost overhead that can be applied post-hoc to pre-trained off-the-shelf models.These variants extend Bayesian inference without requiring every model to be trained from scratch.
  • Conclusion: A comprehensive evaluation finds the LA on par with other posterior-approximation approaches at typically much lower computational cost.The conclusion reports this as the paper's broad comparative result.
  • Conclusion: The laplace library is a modular, extensible PyTorch implementation providing user-friendly access to major LA variants.The library is introduced because efficient LA implementation is not straightforward.
  • Conclusion: A probabilistic treatment of only some weights can improve predictive uncertainty in pre-trained transformer models.The conclusion identifies this as an application of a simple LA variant.

Appendix A Derivation

The appendix derives the Laplace posterior and develops tractable subnetwork, last-layer, diagonal, KFAC, and truncated-diagonal curvature approximations. These factorizations trade expressiveness for reduced storage and computation, enabling practical post-hoc or online use.

  • Laplace derivation: The Laplace approximation expands the log posterior around a MAP estimate and identifies the resulting Gaussian covariance as the inverse negative Hessian.The approximate posterior is N(θ | θMAP, Σ), with Σ := Λ−1 and Λ the negative Hessian evaluated at θMAP.
  • Subnetwork Laplace: Full covariance storage is intractable for modern networks, motivating inference over a smaller parameter subnetwork while fixing remaining weights at their MAP values.The subnetwork posterior combines a Laplace posterior over θS with deterministic values for the remaining parameters.
  • Subnetwork Laplace: The subnetwork size S is user-controlled and can be chosen so a full S × S covariance is tractable, preserving dependencies within the selected weights.Less expressive Hessian/Fisher factorizations can also be used for the subnetwork.
  • Subnetwork Laplace: Under an independence assumption, selecting the S weights with the largest marginal variances minimizes the simplified subnetwork objective.The objective uses a binary mask indicating subnetwork membership and marginal variances from the covariance matrix.
  • Last-layer Laplace: The last-layer LA treats the final weight matrix probabilistically while keeping earlier layers at MAP values, reducing computation because the matrix is small.It is a special case of subnetwork inference using the last layer as θS.
  • Hessian factorization: The full Fisher is often infeasible, whereas diagonal factorization stores a length-D vector and reduces inversion cost from O(D3) to O(D).The diagonal approximation treats the negative-log-posterior Hessian as diagonal.
  • Hessian factorization: KFAC models within-layer weight correlations while assuming independence between weights in different layers, positioning it between diagonal and full-Fisher approximations.Its layer-wise factors arise by assuming independence between activations and log-likelihood gradients.

B.4.6 Classification and Generalized Regression

For classification, the paper compares several approximations for the softmax-Gaussian predictive integral and evaluates Laplace variants on CIFAR-10 and MNIST distribution-shift experiments. The probit approximation is computationally efficient, while the Laplace bridge provides a richer distributional approximation.

  • Predictive approximations: The delta method uses a second-order Taylor expansion of softmax around the predictive mean to obtain an analytic approximation.Its expression is softmax(µ_∗) + 1/2 tr(BΣ_∗), where B is the softmax Hessian at µ_∗.
  • Predictive approximations: The probit approximation replaces the logistic-sigmoid with the standard Normal CDF, making the binary predictive integral analytically solvable.Its multiclass generalization depends only on the diagonal of the predictive covariance and therefore ignores logit correlations.
  • Predictive approximations: The Laplace bridge maps a Dirichlet distribution to a Gaussian in logit space and maps the Gaussian parameters back to Dirichlet parameters.This yields a full distribution over solutions of the softmax-Gaussian integral while remaining comparably simple to the probit approximation.
  • Empirical comparisons: The LA achieves AUROC↑ 91.9 on CIFAR-10 versus SVHN, compared with 81.9 for MC-dropout and 80.0 for VOGN.The cited comparison uses ResNet-18 models described as comparable to the paper’s model.
  • Computational cost: The default probit-based LA has prediction wall-clock time essentially equal to MAP because it requires neither MC sampling nor multiple forward passes.The implementation’s GGN/Fisher or last-layer variants can use deterministic predictive approximations, unlike the listed sampling-based baselines.
  • Computational cost: LA prediction with MC samples on logits is around 20% slower than deterministic probit prediction and remains significantly faster than the other methods.The paper reports that deterministic probit predictions perform consistently on par with or better than MC sampling.

C.3 WILDS Experiments

The WILDS experiments test Laplace methods across five real-world distribution shifts spanning image, text, and regression tasks. Laplace is reported as significantly better calibrated than MAP and competitive with temperature scaling and ensembles, while continual-learning experiments use recursive posterior updates with scalable Hessian approximations.

  • WILDS experiments: WILDS covers five datasets spanning hospital histopathology, satellite imagery, demographic text shifts, reviewer shifts, and cross-country wealth regression.The tasks use DenseNet-121, DistilBERT-base-uncased, and ResNet-18 architectures across classification and regression outputs.
  • WILDS experiments: Results report means and standard errors across the number of seeds supplied for each WILDS dataset.Figure 10 includes accuracy or mean squared error, confidence, negative log-likelihood, calibration metrics, and Brier score where applicable.
  • WILDS experiments: The last-layer Laplace method uses KFAC covariance for FMoW and full covariance for the other WILDS datasets, with 10,000-sample linearized MC prediction.The deep-ensemble baseline aggregates five pretrained neural networks from the original benchmark.
  • WILDS results: Laplace is significantly better calibrated than MAP and competitive with temperature scaling and ensembles, especially on OOD splits.Temperature scaling is reported only for classification because it does not apply to the PovertyMap regression task.
  • Continual learning: In continual learning, Bayesian inference updates the posterior recursively after each task by using the previous posterior as the next prior.The update is p(θ | D1, . . . , Dt) ∝ p(Dt | θ)p(θ | D1, . . . , Dt−1).
  • Continual learning: The Laplace approximation supports these continual-learning updates, and KFAC improves performance over diagonal Hessian approximations in prior work.The paper benchmarks this setting on Permuted-MNIST with ten consecutive pixel-permutation tasks.
  • Continual learning: Large neural-network Hessian sums are typically intractable, so continual-learning implementations use diagonal or KFAC approximations.The diagonal update is exact for adding Hessians and the log prior, whereas adding KFAC factors approximates the sum of Kronecker products.
  • Resource comparison: Table 5 compares theoretical memory complexity and actual memory footprint for Wide ResNet 16-4 on CIFAR-10.Its notation includes model parameters, last-layer neurons, outputs, SWAG snapshots, CSGHMC samples, and ensemble size.
Loading 2106.14806v3…