Laplace Redux - Effortless Bayesian Deep Learning

laplace library : https : //github.com/AlexImmer/Laplace Experiments : https : //github.com/runame/laplace-redux 1 Introduction . Despite their successes , modern neural networks ( NNs ) still suffer from several shortcomings that limit their applicability in some settings . These include ( i ) poor calibration and overconfidence , especially when the data distribution shifts between training and testing [ 1 ] , ( ii ) catastrophic forgetting of previously learned tasks when continuously trained on new tasks [ 2 ] , and ( iii ) the difficulty of selecting suitable NN architectures and hyperparameters [ 3 ] . Bayesian modeling [ 4 , 5 ] provides a principled and unified approach to tackle these issues by ( i ) equipping models with robust uncertainty estimates [ 6 ] , ( ii ) enabling models to learn continually by capturing past information [ 7 ] , and ( iii ) allowing for automated model selection by optimally trading off data fit and model complexity [ 8 ] . Even though this provides compelling motivation for using Bayesian neural networks ( BNNs ) [ 9 ] , they have not gained much traction in practice . Common criticisms include that BNNs are difficult ⇤Equal contributors ; author ordering sampled uniformly at random . Correspondence to : ead54 @ cam.ac.uk , agustinus.kristiadi @ uni-tuebingen.de , alexander.immer @ inf.ethz.ch , runa.eschenhagen @ student.uni-tuebingen.de . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . to implement , finicky to tune , expensive to train , and hard to scale to modern models and datasets . For instance , popular variational Bayesian methods [ 10–12 , etc . ] require considerable changes to the training procedure and model architecture . Also , their optimization process is slower and typically more unstable unless carefully tuned [ 13 ] . Other methods , such as deep ensembles [ 14 ] , Monte Carlo dropout [ 6 ] , and SWAG [ 15 ] promise to bring uncertainty quantification to standard NNs in simple manners . But these methods either require a significant cost increase compared to a single network , have limited empirical performance , or an unsatisfying Bayesian interpretation . In this paper , we argue that the Laplace approximation ( LA ) is a simple and cost-efficient , yet competitive approximation method for inference in Bayesian deep learning . First proposed in this context by MacKay [ 16 ] , the LA dates back to the 18th century [ 17 ] . It locally approximates the posterior with a Gaussian distribution centered at a local maximum , with covariance matrix corresponding to the local curvature . Two key advantages of the LA are that the local maximum is readily available from standard maximum a posteriori ( MAP ) training of NNs , and that curvature estimates can be easily and efficiently obtained thanks to recent advances in second-order optimization , both in terms of more efficient approximations to the Hessian [ 18–20 ] and easy-to-use software libraries [ 21 ] . Together , they make the LA practical and readily applicable to many already-trained NNs—the LA essentially enables practitioners to turn their high-performing point-estimate NNs into BNNs easily and quickly , without loss of predictive performance . Furthermore , the LA to the marginal likelihood may even be used for Bayesian model selection or NN training [ 8 , 22 ] . Figure 1 provides an intuition of the LA—we first fit a point estimate of the model and then estimate a Gaussian distribution around that . Yet , despite recent progress in scaling and improving the LA for deep learning [ 23–29 ] , it is far less widespread than other methods . This is likely due to misconceptions , like that the LA is hard to implement due to the Hessian computation , that it must necessarily perform worse than the competitors due to its local nature , or quite simply that it is old and too simple . Here , we show that these are indeed misconceptions . Moreover , we argue that the LA deserves a wider adoption in both practical and research-oriented deep learning . To this end , our work makes the following contributions : 1 . We first survey recent advances and present the key components of scalable and practical Laplace approximations in deep learning ( Section 2 ) . 2 . We then introduce laplace , an easy-to-use PyTorch-based library for “ turning a NN into a BNN ” via the LA ( Section 3 ) . laplace implements a wide range of different LA variants . 3 . Lastly , using laplace , we show in an extensive empirical study that the LA is competitive to alternative approaches , especially considering how simple and cheap it is ( Section 4 ) . 2 The Laplace Approximation in Deep Learning . The LA can be used in two different ways to benefit deep learning : Firstly , we can use the LA to approximate the model ’ s posterior distribution ( see Eq . ( 5 ) below ) to enable probabilistic predictions ( as also illustrated in Fig . 1 ) . Secondly , we can use the LA to approximate the model evidence ( see Eq . ( 6 ) ) to enable model selection ( e.g . hyperparameter tuning ) . The canonical form of ( supervised ) deep learning is that of empirical risk minimization . Given , e.g. , an i.i.d . classification dataset D : = { ( xn 2 RM , yn 2 RC ) } Nn=1 , the weights ✓ 2 RD of an L-layer NN f✓ : RM ! RC are trained to minimize the ( regularized ) empirical risk , which typically decomposes into a sum over empirical loss terms ` ( xn , yn ; ✓ ) and a regularizer r ( ✓ ) , ✓MAP = arg min✓2RD L ( D ; ✓ ) = arg min✓2RD ⇣ r ( ✓ ) + PN n=1 ` ( xn , yn ; ✓ ) ⌘ . ( 1 ) From the Bayesian viewpoint , these terms can be identified with i.i.d . log- likelihoods and a log-prior , respectively and , thus , ✓MAP is indeed a maximum a-posteriori ( MAP ) estimate : ` ( xn , yn ; ✓ ) = log p ( yn | f✓ ( xn ) ) and r ( ✓ ) = log p ( ✓ ) ( 2 ) For example , the widely used weight regularizer r ( ✓ ) = 12 2 k✓k2 ( a.k.a . weight decay ) corresponds to a centered Gaussian prior p ( ✓ ) = N ( ✓ ; 0 , 2I ) , and the cross-entropy loss amounts to a categorical likelihood . Hence , the exponential of the negative training loss exp ( L ( D ; ✓ ) ) amounts to an unnormalized posterior . By normalizing it , we obtain p ( ✓ | D ) = 1Z p ( D | ✓ ) p ( ✓ ) = 1 Z exp ( L ( D ; ✓ ) ) , Z : = R p ( D | ✓ ) p ( ✓ ) d✓ ( 3 ) with an intractable normalizing constant Z. Laplace approximations [ 17 ] use a second-order expansion of L around ✓MAP to construct a Gaussian approximation to p ( ✓ | D ) . I.e . we consider : L ( D ; ✓ ) ⇡ L ( D ; ✓MAP ) + 1 2 ( ✓ ✓MAP ) | r 2 ✓L ( D ; ✓ ) |✓MAP ( ✓ ✓MAP ) , ( 4 ) where the first-order term vanishes at ✓MAP . Then we can identify the Laplace approximation as Laplace posterior approximation p ( ✓ | D ) ⇡ N ( ✓ ; ✓MAP , ⌃ ) with ⌃ : = r 2 ✓L ( D ; ✓ ) |✓MAP 1 . ( 5 ) The normalizing constant Z ( which is typically referred to as the marginal likelihood or evidence ) is useful for model selection and can also be approximated as Laplace approximation of the evidence Z ⇡ exp ( L ( D ; ✓MAP ) ) ( 2⇡ ) D/2 ( det ⌃ ) 1/2 . ( 6 ) See Appendix A for more details . Thus , to obtain the approximate posterior , we first need to find the argmax ✓MAP of the log-posterior function , i.e . do “ standard ” deep learning with regularized empirical risk minimization . The only additional step is to compute the inverse of the Hessian matrix at ✓MAP ( see Figure 1 ( b ) ) . The LA can therefore be constructed post-hoc to a pre-trained network , even one downloaded off-the-shelf . As we discuss below , the Hessian computation can be offloaded to recently advanced automatic differentiation libraries [ 21 ] . LAs are widely used to approximate the posterior distribution in logistic regression [ 31 ] , Gaussian process classification [ 32 , 33 ] , and also for Bayesian neural networks ( BNNs ) , both shallow [ 34 ] and deep [ 23 ] . The latter is the focus of this work . Generally , any prior with twice differentiable log-density can be used . Due to the popularity of the weight decay regularizer , we assume that the prior is a zero-mean Gaussian p ( ✓ ) = N ( ✓ ; 0 , 2I ) unless stated otherwise.2 The Hessian r2✓L ( D ; ✓ ) |✓MAP then depends both on the ( simple ) log-prior / regularizer and the ( complicated ) log-likelihood / empirical risk : r 2 ✓L ( D ; ✓ ) |✓MAP = 2I PN n=1 r 2 ✓ log p ( yn | f✓ ( xn ) ) |✓MAP . ( 7 ) A naive implementation of the Hessian is infeasible because the second term in Eq . ( 7 ) scales quadratically with the number of network parameters , which can be in the millions or even billions [ 35 , 36 ] . In recent years , several works have addressed scalability , as well as other factors that affect approximation quality and predictive performance of the LA . In the following , we identify , review , and discuss four key components that allow LAs to scale and perform well on modern deep architectures . See Fig . 2 for an overview and Appendix B for a more detailed version of the review and discussion . Four Components of Scalable Laplace Approximations for Deep Neural Networks 1 Inference over all Weights or Subsets of Weights In most cases , it is possible to treat all weights probabilistically when using appropriate approximations of the Hessian , as we discuss below in 2 . Another simple way to scale the LA to large NNs 2One can also consider a per-layer or even per-parameter weight decay , which corresponds to a more general , but still comparably simple Gaussian prior . In particular , the Hessian of this prior is still diagonal and constant . initialized or pre-trained weights ✓ ) , with corresponding laplace code . 1 We first choose which part of the model we want to perform inference over with the LA . 2 We then select how to to approximate the Hessian . 3 We can then perform model selection using the evidence : ( a ) If we started with an untrained model f✓ , we can jointly train the model and use the evidence to tune hyperparameters online . ( b ) If we started with a pre-trained model , we can use the evidence to tune the hyperparameters post-hoc . Here , shades represent the loss landscape , while contours represent LA log-posteriors—faded contours represent intermediate iterates during hyperparameter tuning to obtain the final log-posterior ( thick yellow contours ) . 4 Finally , to make predictions for a new input x⇤ , we have several options for computing/approximating the predictive distribution p ( y|f✓ ( x⇤ ) , D ) . ( without Hessian approximations ) is the subnetwork LA [ 27 ] , which only treats a subset of the model parameters probabilistically with the LA and leaves the remaining parameters at their MAP-estimated values . An important special case of this applies the LA to only the last linear layer of an L-layer NN , while fixing the feature extractor defined by the first L 1 layers at its MAP estimate [ 37 , 28 ] . This last-layer LA is cost-effective yet compelling both theoretically and in practice [ 28 ] . 2 Hessian Approximations and Their Factorizations One advance in second-order optimization that the LA can benefit from are positive semi-definite approximations to the ( potentially indefinite ) Hessian of the log-likelihoods of NNs in the second term of Eq . ( 7 ) [ 38 ] . The Fisher information matrix [ 39 ] , abbreviated as the Fisher and defined by F : = PN n=1 Eby⇠p ( y | f✓ ( xn ) ) [ ( r✓ log p ( by | f✓ ( xn ) ) |✓MAP ) ( r✓ log p ( by | f✓ ( xn ) ) |✓MAP ) | ] , ( 8 ) is one such choice.3 One can also use the generalized Gauss-Newton matrix ( GGN ) matrix [ 41 ] G : = PN n=1 J ( xn ) ⇣ r 2 f log p ( yn | f ) |f=f✓MAP ( xn ) ⌘ J ( xn ) | , ( 9 ) where J ( xn ) : = r✓f✓ ( xn ) |✓MAP is the NN ’ s Jacobian matrix . As the Fisher and GGN are equivalent for common log-likelihoods [ 38 ] , we will henceforth refer to them interchangeably . In deep LAs , they have emerged as the default choice [ 23 , 24 , 28 , 29 , 27 , 26 , etc. ] . 3If , instead of taking expectation in ( 8 ) , we use the training label yn , we call the matrix the empirical Fisher , which is distinct from the Fisher [ 38 , 40 ] . As F and G are still quadratically large , we typically need further factorization assumptions . The most lightweight is a diagonal factorization which ignores off-diagonal elements [ 42 , 43 ] . More expressive alternatives are block-diagonal factorizations such as Kronecker-factored approximate curvature ( KFAC ) [ 18–20 ] , which factorizes each within-layer Fisher4 as a Kronecker product of two smaller matrices . KFAC has been successfully applied to the LA [ 23 , 24 ] and can be improved by low-rank approximations of the KFAC factors [ 29 ] by leveraging their eigendecompositions [ 44 ] . Finally , recent work has studied/enabled low-rank approximations of the Hessian/Fisher [ 45–47 ] . 3 Hyperparameter Tuning As with all approximate inference methods , the performance of the LA depends on the ( hyper ) parameters of the prior and likelihood . For instance , it is typically beneficial to tune the prior variance 2 used for inference [ 23 , 28 , 27 , 26 , 22 ] . Commonly , this is done through cross-validation , e.g . by maximizing the validation log-likelihood [ 23 , 48 ] or , additionally , using out-of-distribution data [ 28 , 49 ] . When using the LA , however , marginal likelihood maximization ( a.k.a . empirical Bayes or the evidence framework [ 34 , 50 ] ) constitutes a more principled alternative to tune these hyperparameters , and requires no validation data . Immer et al . [ 22 ] showed that marginal likelihood maximization with LA can work in deep learning and even be performed in an online manner jointly with the MAP estimation . Note that such approach is not necessarily feasible for other approximate inference methods because most do not provide an estimate of the marginal likelihood . Other recent approaches for hyperparameter tuning for the LA include Bayesian optimization [ 51 ] or the addition of dedicated , trainable hidden units for the sole purpose of uncertainty tuning [ 49 ] . 4 Approximate Predictive Distribution To predict using a posterior ( approximation ) p ( ✓ | D ) , we need to compute p ( y | f ( x⇤ ) , D ) =R p ( y | f✓ ( x⇤ ) ) p ( ✓ | D ) d✓ for any test point x⇤ 2 Rn , which is intractable in general . The simplest but most general approximation to p ( y |x⇤ , D ) is Monte Carlo integration using S samples ( ✓s ) Ss=1 from p ( ✓ | D ) : p ( y | f ( x⇤ ) , D ) ⇡ S 1 PS s=1 p ( y | f✓s ( x⇤ ) ) . However , for LAs with GGN and Fisher Hessian approximations Monte Carlo integration can perform poorly [ 48 , 26 ] . Immer et al . [ 26 ] attribute this to the inconsistency between Hessian approximation and the predictive and suggest to use a linearized predictive instead , which can also be useful for theoretic analyses [ 28 ] . For the last-layer LA , the Hessian coincides with the GGN and the linearized predictive is exact . The predictive of a linearized neural network with a LA approximation to the posterior p ( ✓ | D ) ⇡ N ( ✓ ; ✓MAP , ⌃ ) results in a Gaussian distribution on neural network outputs f⇤ : = f ( x⇤ ) and therefore enables simple approximations or even a closed-form solution . The distribution on the outputs is given by p ( f⇤ |x⇤ , D ) ⇡ N ( f⇤ ; f✓MAP ( x⇤ ) , J ( x⇤ ) |⌃J ( x⇤ ) ) and is typically significantly lowerdimensional ( number of outputs C instead of parameters D ) . It can also be inferred entirely in function space as a Gaussian process [ 25 , 26 ] . Given the distribution on outputs f⇤ , the predictive distribution can be obtained by integration against the likelihood : p ( y |x⇤ , D ) = R p ( y | f⇤ ) p ( f⇤ |x⇤ , D ) d✓ . In the case of regression with a Gaussian likelihood with variance 2 , the solution can even be obtained analytically : p ( y |x⇤ , D ) ⇡ N ( y ; f✓MAP ( x⇤ ) , J ( x⇤ ) |⌃J ( x⇤ ) + 2I ) . For non-Gaussian likelihoods , e.g . in classification , a further approximation is needed . Again , the simplest approximation to this is Monte Carlo integration . In the binary case , we can employ the probit approximation [ 31 , 16 ] which approximates the logistic function with the probit function . In the multi-class case , we can use its generalization , the extended probit approximation [ 52 ] . Finally , first proposed for non-BNN applications [ 53 , 54 ] , the Laplace bridge approximates the softmax-Gaussian integral via a Dirichlet distribution [ 55 ] . The key advantage is that it yields a distribution of the integral solutions . 3 laplace : A Toolkit for Deep Laplace Approximations Implementing the LA is non-trivial , as it requires efficient computation and storage of the Hessian . While this is not fundamentally difficult , there exists no complete , easy-to-use , and standardized implementation of various LA flavors—instead , it is common for deep learning researchers to repeatedly re-implement the LA and Hessian computation with varying efficiency [ 56–58 , etc. ] . An efficient implementation typically requires hundreds of lines of code , making it hard to quickly prototype 4The elements F or G corresponding to the weight Wl ✓ ✓ of the l-th layer of the network . 1 from laplace import Laplace 2 3 # Load pre-trained model 4 model = load_map_model ( ) 5 6 # Define and fit LA variant with custom settings 7 la = Laplace ( model , 'classification ' , 8 subset_of_weights='all ' , 9 hessian_structure='diag ' ) 10 la.fit ( train_loader ) 11 la.optimize_prior_precision ( method='CV ' , 12 val_loader=val_loader ) 13 14 # Make prediction with custom predictive approx . 15 pred = la ( x , pred_type='glm ' , link_approx='probit ' ) Listing 1 : Fit diagonal LA over all weights of a pre-trained classification model , do post-hoc tuning of the prior precision hyperparameter using cross-validation , and make a prediction for input x with the probit approximation . 1 from laplace import Laplace 2 3 # Load un- or pre-trained model 4 model = load_map_model ( ) 5 6 # Fit default , recommended LA variant : 7 # Last-layer KFAC LA 8 la = Laplace ( model , 'regression ' ) 9 la.fit ( train_loader ) 10 11 # Differentiate marginal likelihood w.r.t . 12 # prior precision and observation noise 13 ml = la.marglik ( prior_precision=prior_prec , 14 sigma_noise=obs_noise ) 15 ml.backward ( ) Listing 2 : Fit KFAC LA over the last layer of a pre- or un-trained regression model and differentiate its marginal likelihood w.r.t . some hyperparameters for post-hoc hyperparameter tuning or online empirical Bayes ( see Immer et al . [ 22 ] ) . with the LA . To address this , we introduce laplace : a simple , easy-to-use , extensible library for scalable LAs of deep NNs in PyTorch [ 59 ] . laplace enables all sensible combinations of the four components discussed in Section 2—see Fig . 2 for details . Listings 1 and 2 show code examples . The core of laplace consists of efficient implementations of the LA ’ s key quantities : ( i ) posterior ( i.e . Hessian computation and storage ) , ( ii ) marginal likelihood , and ( iii ) posterior predictive . For ( i ) , to take advantage of advances in automatic differentiation , we outsource the Hessian computation to state-of-the-art , optimized second-order optimization libraries : BackPACK [ 21 ] and ASDL [ 60 ] . Moreover , we design laplace in a modular manner that makes it easy to add new backends and approximations in the future . For ( ii ) , we follow Immer et al . [ 22 ] in our implementation of the LA ’ s marginal likelihood—it is thus both efficient and differentiable and allows the user to implement both online and post-hoc marginal likelihood tuning , cf . Listing 2 . Note that laplace also supports standard cross-validation for hyperparameter tuning [ 23 , 28 ] , as shown in Listing 1 . Finally , for ( iii ) , laplace supports all approximations to the posterior predictive distribution discussed in Section 2—it thus provides the user with flexibility in making predictions , depending on the computational budget . Default behavior To abstract away from a large number of options available ( Section 2 ) , we provide the following default choices based on our extensive experiments ( Section 4 ) ; they should be applicable and perform decently in the majority of use cases : we assume a pre-trained network and treat only the last-layer weights probabilistically ( last-layer LA ) , use the KFAC factorization of the GGN and tune the hyperparameters post-hoc using empirical Bayes . To make predictions , we use the closed-form Gaussian predictive distribution for regression and the ( extended ) probit approximation for classification . Of course , the user can pick custom choices ( Listings 1 and 2 ) . Limitations Because laplace employs external libraries ( BackPACK [ 21 ] and ASDL [ 60 ] ) as backends , it inherits the available choices of Hessian factorizations from these libraries . For instance , the LA variant proposed by Lee et al . [ 29 ] can currently not be implemented via laplace , because neither backend supports eigenvalue-corrected KFAC [ 44 ] ( yet ) .

Summary: The authors bring to light the recent progress in approximate Bayesian inference via the Laplace approximation, and note that despite its simplicity, its adoption has not kept pace for modern deep learning methods. The work describes a consolidation of recent work into a user-friendly PyTorch-based software library called `laplace`. LA can be applied post-hoc to existing deep learning models. With this package in place, the authors provide practical recommendations for reasonable defaults when used in practice.