Expressive yet Tractable Bayesian Deep Learning via Subnetwork Inference

1 INTRODUCTION . Deep neural networks ( DNNs ) still suffer from critical shortcomings that make them unfit for important applications . For instance , DNNs tend to be poorly calibrated and overconfident in their predictions , especially when there is a shift in the train and test distributions ( Nguyen et al. , 2015 ; Guo et al. , 2017 ) . To reliably inform decision making , DNNs must be able to robustly quantify the uncertainty in their predictions , which is particularly important in safety-critical areas such as healthcare or autonomous driving ( Amodei et al. , 2016 ; Filos et al. , 2019a ; Fridman et al. , 2019 ) . Bayesian modeling ( Ghahramani , 2015 ; Gal , 2016 ) presents a principled way to capture predictive uncertainty via the posterior distribution over model parameters . Unfortunately , due to their nonlinearities , exact posterior inference is intractable in DNNs . Despite recent successes in the field of Bayesian deep learning ( Blundell et al. , 2015 ; Gal & Ghahramani , 2016 ; Osawa et al. , 2019 ; Maddox et al. , 2019 ; Dusenberry et al. , 2020 ) , existing methods are only made scalable to modern DNNs with large numbers of parameters by invoking unrealistic assumptions . This severely limits the expressiveness of the inferred posterior and thus deteriorates the quality of the induced uncertainty estimates ( Ovadia et al. , 2019 ; Fort et al. , 2019 ; Foong et al. , 2019a ; Ashukha et al. , 2020a ) . Due to the heavy overparameterization of DNNs , their accuracy is well-preserved by a small subnetwork ( Cheng et al. , 2017 ) . Additionally , recent work by Izmailov et al . ( 2019 ) has shown how performing inference over a low dimensional subspace of the weights can result in accurate uncertainty quantification . These observations prompt the following question for a DNN ’ s uncertainty : Can a full DNN ’ s model uncertainty be well-preserved by a small subnetwork ’ s model uncertainty ? We answer this question in the affirmative . We show both theoretically and empirically that the full network posterior can be well represented by a subnetwork ’ s posterior . As a result , we can use more expensive but faithful posterior approximations over just that subnetwork . We show that this achieves better uncertainty quantification than if we use cheaper , but more crude , posterior approximations over the full network . The contributions of this paper are as follows : 1 . We propose a new Bayesian deep learning approach that performs Bayesian inference over only a small subset of the model weights and keeps all other weights deterministic . This allows us to use expressive posterior approximations that are typically intractable in DNNs . 2 . As a concrete instantiation of this framework , we develop a practical and scalable Bayesian deep learning method that uses the linearized Laplace approximation to infer a fullcovariance Gaussian posterior over a subnetwork within a point-estimated neural network . 3 . We formally characterize the discrepancy between the posterior distributions over a subnetwork and the full network ( in terms of their Wasserstein distance ) in the linearized model , and derive a theoretically motivated strategy to select a subnetwork that minimizes this discrepancy under certain assumptions . 4 . We empirically show , on various benchmarks , that our method compares favourably against point-estimated networks and other Bayesian deep learning methods , experimentally confirming that expressive subnetwork inference is superior to crude inference over full networks . 2 SUBNETWORK POSTERIOR APPROXIMATION . Bayesian neural networks ( BNNs ) aim to capture model uncertainty , i.e. , uncertainty about the choice of weightsW which arises due to multiple plausible explanations of the training data { y , X } . Here , y is the dependent variable ( e.g . classification label ) andX is the feature matrix . A prior distribution p ( W ) is specified over the BNN ’ s weights . We wish to infer their full posterior distribution p ( W|y , X ) ∝ p ( y|X , W ) p ( W ) . ( 1 ) To make predictions , we then estimate the posterior predictive distribution that averages the network ’ s predictions across all possible settings of the weights , weighted by their posterior probability , i.e . p ( y∗|X∗ , y , X ) = ∫ W p ( y∗|X∗ , W ) p ( W|y , X ) dW . ( 2 ) Unfortunately , due to the size of modern deep neural networks , it is not only intractable to infer the exact posterior distribution p ( W|y , X ) in Eq . ( 1 ) , but it is even computationally challenging to properly approximate it . As a consequence , crude posterior approximations such as complete factorization are commonly employed ( Blundell et al. , 2015 ; Hernández-Lobato & Adams , 2015 ; Kingma et al. , 2015 ; Khan et al. , 2018 ; Osawa et al. , 2019 ) , i.e . p ( W|y , X ) ≈∏Dd=1 q ( wd ) where wd denotes the d-th weight in the D-dimensional neural network weight vector W ∈ RD ( the concatenation and flattening of all layers ’ weight matrices ) . Clearly , this is a very wishful assumption ; In practise , it suffers from severe pathologies ( Foong et al. , 2019a ; b ) . In this work , we question the implicit assumption that a good posterior approximation needs to include all BNN parameters . Instead , we aim to perform inference only over a small subset of the weights . This approach is well-motivated for two reasons : 1 . Overparameterization : Maddox et al . ( 2020 ) have shown that , in the neighborhood of local optima , there are many directions that leave the NN ’ s predictions unchanged . Moreover , NNs can be heavily pruned without sacrificing test-set accuracy ( Frankle & Carbin , 2019 ) . Thus , the majority of a NN ’ s predictive power might be isolated to a small subnetwork . 2 . Inference over submodels : Previous work1 has provided evidence that inference can be effective even when not done on the full parameter space . Izmailov et al . ( 2019 ) performed inference over a low-dimensional projection of the weights . Neural-linear models , which give a Bayesian treatment to only the last layer of a DNN , have shown to be competitive with full-network approaches ( Riquelme et al. , 2018 ; Ober & Rasmussen , 2019 ) . We thus combine these ideas , making the following two-step approximation of the posterior in Eq . ( 1 ) : p ( W|y , X ) ≈ p ( WS |y , X ) ∏ r δ ( wr − w∗r ) ≈ q ( WS ) ∏ r δ ( wr − w∗r ) . ( 3 ) The first approximation decomposes the full neural network posterior p ( W|y , X ) into a posterior p ( WS |y , X ) over the subnetwork WS and delta functions δ ( wr − w∗r ) over all remaining weights { wr } r , keeping them at fixed values w∗r ∈ R. This can be viewed as pruning the variances of the weights { wr } r to zero , which is in contrast to ordinary weight pruning methods that set the weights { wr } r themselves to zero . The second approximation is a result of posterior inference over the subnetwork still being intractable . In turn , we introduce the approximate distribution q ( WS ) . Yet , as the subnetwork is much smaller than the full network , we can afford to make q ( WS ) expressive and able to capture rich dependencies across the weights within the subnetwork . 3 SUBNETWORK INFERENCE VIA LAPLACE APPROXIMATION . To obtain a method that is as practical as possible , we propose to use inference techniques that can estimate a posterior distribution post-hoc from a point-estimated network . The Laplace approximation ( MacKay , 1992 ) is well-suited to this task as it derives the approximate posterior from the local optimization landscape . Other inference procedures , such as SWAG ( Maddox et al. , 2019 ) , could also be used . Nevertheless , we focus on Laplace due to it being a well-studied , fundamental technique . Step # 1 : Point Estimation . The first step of the procedure is to train a neural network to obtain a point estimate of the weights , denoted WMAP . This estimate should respect the Bayesian model given in Eq . ( 1 ) , and therefore we optimize the maximum a-posteriori ( MAP ) objective : WMAP = argmaxW [ log p ( y|X , W ) + log p ( W ) ] . ( 4 ) This can be done using standard stochastic gradient-based optimization methods commonly-used in modern deep learning ( Goodfellow et al. , 2016 ) . This step is illustrated in Fig . 1 ( a ) . Step # 2 : Subnetwork Selection . The second step is to identify a small subnetworkWS . Ideally , we would like to identify the subnetwork whose posterior is ‘ closest ’ to the full-network posterior . We formalize this argument in Section 4 and describe a principled strategy that , under certain conditions , minimizes the 2-Wasserstein distance between the sub- and full-network posteriors . All other weights not belonging to that subnetwork are then assigned fixed values : the MAP estimates obtained in Step # 1 . See Fig . 1 ( b ) for an illustration of this step . Step # 3 : Bayesian Inference . Given the subnetwork point estimateW SMAP , we use the Laplace approximation to infer a full-covariance Gaussian posterior distribution over the subnetwork WS : p ( WS |y , X ) ≈ q ( WS ) = N ( WS ; W S MAP , H −1 ) ( 5 ) where the posterior covariance matrix H−1 ∈ RD×D corresponds to the inverse of the average Hessian of the negative log posterior , i.e . H = NE [ −∂2 log p ( y|X , W ) /∂W 2 ] + λI . Here , the expectation is w.r.t . the data generating distribution and λ is the precision of a zero-mean factorized Gaussian prior p ( W ) = N ( W ; 0 , λ−1I ) . In practice , we approximate the Hessian H with the generalized Gauss-Newton ( GGN ) matrix H̃ ( Schraudolph , 2002 ) , i.e . H̃ = N∑ n=1 J > n HnJn + λI , with Jn = ∂f ( xn , W ) ∂W and Hn = ∂2L ( yn , f ( xn , W ) ) ∂2f ( xn , W ) ( 6 ) 1See Section 6 for a more thorough discussion of related work . where Jn ∈ RO×D is the Jacobian of the neural network features f ( xn , W ) ∈ RO w.r.t . the weights W , andHn ∈ RO×O is the Hessian of the loss L ( yn , f ( xn , W ) ) w.r.t . the features f ( xn , W ) .The GGN H̃ has clear practical advantages over the Hessian H ; see Martens & Sutskever ( 2011 ) and Martens ( 2016 ) . Using the Laplace approximation with the GGN Hessian can be viewed as an implicit local linearization of the underlying neural network f ( x , W ) at its MAP estimateWMAP , fMAPlin ( x , W ) = f ( x , WMAP ) + JWMAP ( x ) ( W −WMAP ) ( 7 ) where JWMAP ( x ) = ∂f ( x , WMAP ) /∂WMAP ∈ RO×D ( Immer et al. , 2020 ) . Note that the model in Eq . ( 7 ) is linear in W , as only the term JWMAP ( x ) W depends linearly onW , while the other terms are constant w.r.t.W and can thus be subsumed into an additive bias term ( Khan et al. , 2019 ) . The GGN approximation thus locally turns the underlying probabilistic model from a Bayesian neural network into a ( generalized ) linear model , with basis function expansion JWMAP ( x ) of covariate x ( Immer et al. , 2020 ) . Put differently , linearized Laplace in the neural network f ( x , W ) is equivalent to ordinary Laplace in the linear model fMAPlin ( x , W ) in Eq . ( 7 ) , as the GGN H̃ corresponding to f ( x , W ) in Eq . ( 6 ) is equivalent to the Hessian H corresponding to fMAPlin ( x , W ) in Eq . ( 7 ) ( Khan et al. , 2019 ) . This is a useful property that will allow us to derive a principled subnetwork selection strategy in Section 4 . This step is illustrated in Fig . 1 ( c ) . We emphasize that this whole procedure ( i.e . Steps # 1- # 3 ) is a perfectly valid mixed inference strategy , performing full Laplace inference over the selected subnetwork and MAP inference over all remaining weights . Step # 4 : Prediction . Given the linearized Laplace approximation over the subnetwork WS in Eqs . ( 5 ) and ( 6 ) , i.e . q ( WS ) = N ( WS ; W SMAP , H̃−1 ) , we can then compute the posterior predictive distribution . While , traditionally , one would compute the predictive distribution using the original Bayesian neural network likelihood , i.e . p ( y|X , W ) = p ( y|f ( x , W ) ) , Immer et al . ( 2020 ) recently suggested that , since inference was ( implicitly ) done in the GGN-linearized model , it is more principled to instead predict using the linearized likelihood Eq . ( 7 ) , i.e . p ( y|X , W ) = p ( y|fMAPlin ( x , W ) ) . This provides a formal justification for the empirical superiority of this approach observed previously ( Lawrence , 2001 ; Foong et al. , 2019b ) . We thus obtain the linearized predictive distribution p ( y∗|X∗ , y , X ) ≈ ∫ W p ( y∗|fMAPlin ( X∗ , W ) ) N ( WS ; W SMAP , H̃−1 ) ∏ r δ ( wr−w∗r ) dW . ( 8 ) There are two ways to compute Eq . ( 8 ) : Firstly , via a Monte Carlo approximation p ( y∗|X∗ , y , X ) ' 1 M ∑M m=1 p ( y ∗|fMAPlin ( X∗ , W̃m ) ) by sampling W̃m from N ( W SMAP , H̃−1 ) and ∏ r δ ( wr − w∗r ) , the latter of which is trivial . Secondly , due to linearity of p ( y∗|fMAPlin ( X∗ , W ) ) , there are closedform expressions which are exact for Gaussian likelihoods ( i.e . regression ) and approximate for categorical ones ( i.e . classification ) ( Bishop , 2006 ; Gibbs , 1998 ) . This step is illustrated in Fig . 1 ( d ) .

Summary: The authors focus on the important problem of scalable approximate inference in Bayesian NNs. More specifically, they propose a method for scalable BNNs via a (full-covariance Gaussian) Laplace approximation on a (Wasserstein-based) pruned subnetwork within a deterministically-trained model. They include a theoretical analysis for a simple generalized linear model, and experiments on 1D regression, tabular regression, and larger-scale image classification with CIFAR-10 (using the dataset shift setup from Ovadia et al., (2019)). From the experiments, they show that their method generally outperforms comparable methods (including deep ensembles) on metric performance and on the ability to capture in-between uncertainty.