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 paper proposes to approximate the posterior distribution of a Bayesian neural network by an approximation that consists of a deterministic component. The authors select a sub network and infer approximate posterior distributions over the weights in the sub network. All other weights are estimated via MAP point estimation.  A sufficiently small sub-network allows high fidelity posterior approximations that do not make restrictive mean field assumptions to be tractable.