Expressive yet Tractable Bayesian Deep Learning via Subnetwork Inference
Abstract
Scaling Bayesian inference to the large parameter spaces of deep neural networks requires restrictive approximations.
We propose performing inference over only a small subset of the model parameters while keeping all others as point estimates.
This enables us to use expressive posterior approximations that are intractable in the full model.
In particular, we develop a practical and scalable Bayesian deep learning method that ﬁrst trains a point estimate, and then infers a full covariance Gaussian posterior approximation over a subnetwork.
We propose a subnetwork selection procedure which aims to optimally preserve posterior uncertainty.
Empirical studies demonstrate the eﬀectiveness of our approach.
1. Subnetwork Posterior Approximation
Deep neural networks (DNNs) still suﬀer from critical shortcomings that make them unﬁt for important applications.
For instance, DNNs tend to be poorly calibrated and overconﬁdent 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).
Bayesian modeling (Ghahramani, 2015) presents a principled way to capture model uncertainty , i.e., uncertainty about the choice of weights W which arises due to multiple plausible explanations of the training data { y , X } .
Here, y is the target (e.g. classiﬁcation label) and X are the features.
A prior distribution p ( W ) is speciﬁed over the Bayesian neural network (BNN) weights.
We wish to infer their full posterior distribution p ( W | y , X ) ∝ p
( y | X , W ) p ( W ).
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 ) =
(cid:82) W p ( y ∗ | X ∗ , W ) p ( W | y , X ) d W . xxh 2 hhy 0.50. [Figure]
xxh 2 hhy 0.50.
[Figure]
xxh 2 hhy 0.50.5 0.10.
[Figure]
Bayes. Inference
xxh 2 hhy 0.50.5 0.10. [Figure] Prediction 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 ), 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; Gal and Ghahramani, 2016; Hern´andez-Lobato and Adams, 2015; Kingma et al., 2015; Khan et al., 2018; Osawa et al., 2019), i.e. p ( W | y , X ) ≈ (cid:81)
Dd =1 q (w d ) where w d denotes the d -th weight in the Ddimensional neural network weight vector W ∈ R D (the concatenation and ﬂattening of all layers’ weight matrices).
In practice, this severely limits the expressiveness of the inferred posterior and thus deteriorates the quality of the induced predictive uncertainty estimates (Ovadia et al., 2019; Fort et al., 2019; Foong et al., 2019a,b; Ashukha et al., 2020).
In this work, we question the implicit assumption that a good posterior approximation needs to include all model parameters.
Instead, we aim to perform inference only over a small subset of the model weights.
This approach is well-motivated for two reasons:
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 sacriﬁcing test set accuracy (Frankle and Carbin, 2019).
This suggests that the majority of a NN’s predictive power is isolated to a small subnetwork.
Inference over submodels: Previous work 1 has provided evidence that inference can be eﬀective even when not done over 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 empirically been shown to perform well (Riquelme et al., 2018;
Ober and Rasmussen, 2019; Kristiadi et al., 2020).
We thus combine these ideas, making the following two-step approximation of the posterior:
[Equation]
The ﬁrst approximation in Eq.
(1) decomposes the full neural network posterior p ( W | y , X ) into a posterior p ( W S
| y , X ) over the subnetwork W S and delta functions δ (w r − w ∗ r ) over all remaining weights { w r } r , keeping them deterministic at ﬁxed values w ∗ r ∈ R .
This can be viewed as pruning the variances of the weights { w r } r to zero, which is in contrast to ordinary weight pruning methods that set the weights { w r } r themselves to zero (Frankle and
Carbin, 2019).
The second approximation in Eq.
(1) introduces the approximate distribution q ( W S ), because exact posterior inference over the subnetwork W S is still intractable.
Yet, as the subnetwork is much smaller than the full network, we can aﬀord to make q ( W S ) expressive and able to capture rich dependencies across the weights within the subnetwork.
In this paper, we show both theoretically and empirically that the full neural network posterior can be well represented by a subnetwork’s posterior as deﬁned in Eq.
As a result, expressive posterior inference over a subnetwork results in better uncertainty calibration and robustness to distribution shift than crude inference over the full network.
2. Subnetwork Inference via Laplace Approximation
As a concrete instantiation of the described subnetwork inference framework, we now describe a practical and scalable Bayesian deep learning method.
To this end, we use the (linearized)
Laplace approximation (MacKay, 1992) for post-hoc inference of a full-covariance Gaussian posterior over a subnetwork within a pre-trained, point-estimated neural network.
Step #1: Point Estimation.
The ﬁrst step of the proposed procedure is to train a neural network to obtain a point estimate of the weights, denoted W MAP .
This point estimate should respect our Bayesian model, and therefore we optimize the maximum a-posteriori (MAP) objective, i.e. W MAP = arg max W [log p ( y | X , W ) + log p ( W )].
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 subnetwork
W S (illustrated in Fig.
1 (b)).
Ideally, we would like to ﬁnd the subnetwork whose posterior is ‘closest’ to the full-network posterior.
To this end, we ﬁrst formally characterize the discrepancy between the posterior distributions over a subnetwork and the full network in terms of their 2-Wasserstein distance.
We then derive a principled strategy that, under certain conditions, minimizes this discrepancy: 1) Estimate a fully-factorized posterior distribution over all weights in the full network (e.g. using a diagonal Laplace approximation) and 2) deﬁne the subnetwork to comprise of the weights with the largest marginal variances .
The full details of this procedure are described in Appendix B. All other weights not belonging to that subnetwork are then assigned ﬁxed values: the MAP estimates obtained in Step #1.
Step #3: Bayesian Inference.
Given the subnetwork point estimate W SMAP , we use the Laplace approximation to infer a full-covariance Gaussian posterior distribution over the subnetwork W S , i.e. p ( W S |
y , X )
≈ q ( W S )
=
N (cid:0) W S ; W SMAP , H − 1 (cid:1) .
Here, the posterior covariance matrix H − 1 ∈ R D
× D corresponds to the inverse of the average Hessian of the negative log posterior, i.e. H = N E (cid:2) − ∂ 2 log p ( y | X , W ) /∂
W 2 (cid:3)
+ λ I , where 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 , λ − 1 I ).
In practice, we approximate the Hessian
H with the generalized Gauss-Newton (GGN) matrix (cid:101) H (Schraudolph, 2002), i.e.
(cid:101) H =
N (cid:88) n =
J (cid:62) n H n J n + λ I , with J n = ∂ f ( x n , W ) ∂ W and H n = ∂ 2 L ( y n , f ( x n , W ))
∂ 2 f ( x n , W ) (2) where J n ∈ R
O
× D is the Jacobian of the neural network features f ( x n , W ) ∈
R O w.r.t.
the weights W , and H n ∈ R
O × O is the Hessian of the loss L ( y n , f ( x n , W ))
w.r.t.
the features f ( x n , W ).
The GGN (cid:101) H has clear practical advantages over the Hessian H ; see
Martens and Sutskever (2011) and Martens (2016).
Using the Laplace approximation with the GGN-approximation to the Hessian can be viewed as an implicit local linearization of the underlying neural network f
( x , W ) at its MAP estimate W MAP , i.e.
[Equation]
(3) where J W MAP ( x ) =
∂
f ( x , W MAP ) /∂
W MAP ∈ R O × D (Immer et al., 2020).
Note that the model in Eq.
(3) is linear in W , as only the term J W MAP ( x ) W depends linearly on W , 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 J W MAP ( x ) of covariate x
(Immer et al., 2020).
Put diﬀerently, linearized Laplace in the neural network f
( x , W ) is equivalent to ordinary Laplace in the linear model f MAPlin ( x , W ) in Eq.
(3), as the GGN (cid:101) H corresponding to f ( x , W ) in Eq.
(2) is equivalent to the Hessian H corresponding to f MAPlin ( x , W ) in Eq.
(3) (Khan et al., 2019).
This is a useful property that allows us to derive a principled subnetwork selection strategy in Appendix B.
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
W [Equation] , i.e.
[Equation]
[Equation] SMAP , [Equation]
H [Equation] 1 ), we can then compute the posterior predictive distribution to make predictions (illustrated in Fig.
1 (d)).
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.
(3), i.e. p ( y | X , W ) = p ( y | f MAPlin ( x , W )).
This provides a formal justiﬁcation for the empirical superiority of this approach observed previously (Lawrence, 2001;
Foong et al., 2019b).
We thus obtain the linearized predictive distribution p [Equation] ≈
[Equation]
There are two ways to compute Eq.
(4): Firstly, via a Monte Carlo approximation p ( y ∗ | X ∗ , y , X ) (cid:39) 1 M (cid:80) Mm =1 p ( y ∗ | f MAPlin ( X ∗ , (cid:102) W m )) by sampling (cid:102) W m from N ( W SMAP , (cid:101) H − 1 ) and (cid:81) r δ (w r − w ∗ r ), the latter of which is trivial.
Secondly, due to linearity of p ( y ∗ | f MAPlin ( X ∗ , W )), there are closed-form expressions which are exact for Gaussian likelihoods (i.e. regression) and approximate for categorical ones (i.e. classiﬁcation)
(Bishop, 2006; Gibbs, 1998).
[Figure]
3. Empirical Analysis
We now empirically assess the eﬀectiveness of subnetwork inference compared to pointestimated NNs and methods that do less expressive inference over the full network.
We consider two tasks: 1) small-scale synthetic regression and 2) large-scale image classiﬁcation.
In Appendix C, we report additional empirical results on medium-scale tabular regression.
3.1. How does subnetwork inference retain posterior predictive uncertainty?
We ﬁrst assess how the predictive distribution of a full-covariance Gaussian posterior over a selected subnetwork qualitatively compares to that obtained from 1) a full-covariance
Gaussian over the full network ( Full Cov ), 2) a factorised Gaussian posterior over the full network ( Diag ), 3) a full-covariance Gaussian over only the ( Final layer ) of the network (Kristiadi et al., 2020), and 4) a point estimate ( MAP ).
For subnetwork inference, we consider both Wasserstein ( Wass ) (as described in Appendix B) and uniform random subnetwork selection ( Rand ) to obtain subnetworks that comprise of only 50%, 3% and 1% of the model parameters.
Our NN consists of 2 ReLU hidden layers with 50 hidden units each.
We use a homoscedastic Gaussian likelihood with noise variance optimised via maximum likelihood.
Using the linearized predictive in Eq.
(4), all approaches share their predictive mean, allowing us to better compare their uncertainties.
All methods share a prior precision of λ = 3.
We use a synthetic 1D regression task with two separated clusters of inputs (Antor´an et al., 2020), allowing us to probe for ‘in-between’ uncertainty (Foong et al., 2019b).
Results are shown in Fig.
Subnetwork inference preserves more of the uncertainty of full network inference than diagonal Gaussian or ﬁnal layer inference while doing inference over fewer weights.
By capturing weight correlations, subnetwork inference retains uncertainty in between clusters of data.
This is true for both random and Wasserstein subnetwork [Figure] selection.
However, the latter preserves more uncertainty with smaller subnetworks.
This suggests that expressive inference over a carefully selected subnetwork retains more predictive uncertainty than crude approximations over the full network .
3.2. How robust is our approach to distribution shift in image classiﬁcation?
We assess the robustness of large CNNs with subnetwork inference to distribution shift on image classiﬁcation tasks compared to the following baselines: MAP-estimated networks,
Bayesian deep learning methods that do less expressive inference over the full network: MC
Dropout (Gal and Ghahramani, 2016), diagonal Laplace (both of which assume factorisation of the weight posterior), and SWAG (Maddox et al., 2019) (which assumes a diagonal plus low-rank posterior).
We also assess deep ensembles (Lakshminarayanan et al., 2017), which is considered state-of-the-art for uncertainty quantiﬁcation in deep learning (Ovadia et al., 2019; Ashukha et al., 2020).
For all approaches, we use a ResNet-18 (He et al., 2016) with 11,168,000 parameters.
For our method, we retain a subnetwork with only 42,438 of them (i.e. 0 .
See Appendices D and E for more experimental details and results, respectively.
Following Ovadia et al.
Antor´an et al.
(2020), we perform the following two benchmark experiments: 1) We train all methods on MNIST and evaluate their predictive distributions on increasingly rotated digits ( rotated MNIST ).
We train all methods on CIFAR10 and evaluate on data subject to 16 diﬀerent corruptions with 5 levels of intensity each (Hendrycks and Dietterich, 2019) ( corrupted CIFAR ).
Results are shown in Fig.
While all methods perform well on the original test sets, their accuracy degrades quickly for increasing levels of rotation/corruption.
Our approach matches a MAP-estimated network in terms of predictive error as local linearization makes their predictions the same.
Ensembles and SWAG are the most accurate, and, out of our baselines, perform best in terms of log-likelihood.
Even so, subnetwork inference diﬀerentiates itself by being the least overconﬁdent, outperforming all baselines in terms of log-likelihood at all rotation/corruption levels.
Subnetwork inference makes accurate predictions in-distribution while assigning higher uncertainty than the baselines to out-of-distribution inputs.
These results suggest that subnetwork inference results in better uncertainty calibration and robustness to distribution shift than other popular uncertainty quantiﬁcation methods .
4. Conclusion
We develop a practical and scalable method for expressive yet tractable Bayesian deep learning.
We approximate the posterior over a subnetwork within a NN while keeping all other weights deterministic.
Computational cost is decoupled from network size, allowing us to scale expressive approximations, such as full-covariance Gaussians, to modern DNNs.
Our method can be applied post-hoc to any pre-trained model, making it particularly attractive for practical use.
Our experiments suggest that subnetwork inference 1) retains more predictive uncertainty than crude approximations over the full model, and 2) is competitive with state-of-the-art uncertainty quantiﬁcation methods on real-world scale problems. Acknowledgments
We thank Matthias Bauer and Andrew Y. K. Foong for helpful discussions and feedback.
ED acknowledges funding from the EPSRC and Qualcomm.
JA acknowledges support from Microsoft Research, through its PhD Scholarship Programme, and from the EPSRC.
JUA acknowledges funding from the EPSRC and the Michael E. Fisher Studentship in Machine Learning. This work has been performed using resources provided by the Cambridge Tier-2 system operated by the University of Cambridge Research Computing Service (http://www.hpc.cam.ac.uk) funded by EPSRC Tier-2 capital grant EP/P020259/1. Appendix A. Additional Related Work Bayesian Deep Learning.
There have signiﬁcant eﬀorts to characterise the posterior distribution over NN weights p ( W |D ).
Hamiltonian Monte Carlo (Neal, 1995) remains the golden standard for approximate inference in BNNs to this day.
Although asymptotically unbiased, sampling based approaches are diﬃcult to scale to the large datasets (Betancourt, 2015).
As a result, approaches which ﬁnd the best surrogate posterior among an approximating family (most often Gaussians) have gained popularity.
The ﬁrst of these was the
Laplace approximation, introduced by MacKay (1992), who also proposed approximating the predictive posterior with that of the linearised model (Khan et al., 2019; Immer et al., 2020).
The popularisation of larger NN models has made surrogate distributions that capture correlations between weights computationally intractable.
Thus, most modern methods make use of the mean ﬁeld assumption (Blundell et al., 2015; Hern´andez-Lobato and Adams, 2015; Gal and Ghahramani, 2016; Mishkin et al., 2018; Osawa et al., 2019).
This comes at the cost of limited expressivity (Foong et al., 2019a) and empirical under-performance (Ovadia et al., 2019; Antor´an et al., 2020) of uncertainty estimates.
Our proposed approach recovers predictive posterior expressivity while maintaining tractability by lowering the dimensionality of the weight space considered.
This allows us to scale up approximations that do consider weight correlations (MacKay, 1992; Louizos and Welling, 2016; Maddox et al., 2019; Ritter et al., 2018).
Neural Network Linearization.
In the limit of inﬁnite width, NNs converge to
Gaussian process (GP) behaviour (Neal, 1995; Matthews, 2017; Garriga-Alonso et al., 2018).
Recently, these results have been extended to ﬁnite width BNNs when the surrogate posterior is Gaussian (Khan et al., 2019).
We draw upon these results to formulate a subnetwork selection strategy for BNNs.
Neural linear methods perform inference over only the last layer of a NN, while keeping all other layers ﬁxed (Snoek et al., 2015;
Riquelme et al., 2018;
Ovadia et al., 2019; Ober and Rasmussen, 2019; Pinsler et al., 2019; Kristiadi et al., 2020).
These represent a diﬀerent generalised linear model in which the basis functions are deﬁned by the l −
ﬁrst layers of a NN.
They can also be viewed as a special case of subnetwork inference, in which the subnetwork is simply deﬁned to be the last NN layer.
Inference over Subspaces.
The subﬁeld of NN pruning aims to increase the computational eﬃciency of NNs by identifying the smallest subset of weights which are required to make accurate predictions.
Approaches trade-oﬀ computational cost with compression eﬃciency, ranging from those that require multiple training runs (Frankle and Carbin, 2019) to those that prune before training (Wang et al., 2020).
Our work diﬀers in that it retains all NN weights but aims to ﬁnd a small subset over which to perform probabilistic reasoning.
More closely related work to ours is that of Izmailov et al.
(2019), who propose to perform inference over a low-dimensional subspace of weights; e.g. one constructed from the principal components of the SGD trajectory.
Moreover, several recent approaches use low-rank parameterizations of approximate posteriors in the context of variational inference (Rossi et al., 2019;
Swiatkowski et al., 2020; Dusenberry et al., 2020).
This could also be viewed as doing inference over an implicit subspace of weight space.
In contrast, we propose a technique to ﬁnd subsets of weights which are relevant to predictive uncertainty, i.e., we identify axis aligned subspaces.
Appendix B. Principled Subnetwork Selection for Linear(ized) Models
We next analyze the subnetwork inference procedure described in Section 2 for the case of a generalized linear model (GLM)
(Nelder and Baker, 1972), which models the expected response y n given the basis function expansion of the covariates φ n = φ ( x n ) as
[Equation]
Here, w ∈ R D is the vector of model parameters (which subsumes a scalar bias β 0 for notational convenience) and g − 1 ( · ) denotes a link function such that g − 1 : R (cid:55)→ µ y | φ .
In particular, we consider a Bayesian GLM, by specifying a prior distribution p ( w ) over model parameters and aiming to infer the posterior distribution p ( w | y , Φ ) ∝ p
( y | Φ , w ) p ( w ), where Φ =
[ φ 1 , ...
φ N ] T .
Point Estimation.
Obtain the MAP estimate, w MAP = arg max w log p
( y | Φ , w ) + log p ( w ).
For commonly-used link functions (e.g. the identity in case of a Gaussian likelihood for regression, or the sigmoid/softmax function in case of a categorical likelihood for classiﬁcation) and commonly-used priors (e.g. a Gaussian), the log-posterior ∝ log p ( y | Φ , w ) + log p ( w ) is concave.
This allows for simple gradient-based MAP optimisation.
It also makes a full-covariance Gaussian, estimated via Laplace, a faithful approximation to the true, uni-modal posterior, i.e.
[Equation]
(6) where H is the Hessian deﬁned in Section 2.
Note that for the GLM we consider, the
Hessian H is equivalent to the GGN (cid:101)
H deﬁned in Eq.
(2), meaning that an ordinary
Laplace approximation is equivalent to a linearized Laplace approximation (Martens, 2016).
For the case of an identity link function (i.e. a Gaussian likelihood with noise variance σ 20 ) and a Gaussian prior w ∼ N ( 0 , Λ − 1 0 ), the MAP estimate even has a closedform expression, w MAP =
( Φ T Φ + σ 20 Λ 0 ) − 1 Φ T y .
Here, the Laplace approximation in
Eq.
(6) exactly corresponds to the true posterior , i.e. (cid:101) p ( w | y , Φ )
=
p ( w | y , Φ ).
We will thus refer to the posterior (cid:101) p ( w | y , Φ ) in Eq.
(6) as the full posterior .
Subnetwork Selection.
Select a subset of S model weights via a method of choice, yielding a binary vector m ∈ R D where m d = 1 if the d -th weight is part of the subset, and m d = 0 otherwise.
For convenience, we deﬁne the binary mask matrix
M S =
mm (cid:62) ∈ R D
× D which contains 1s in the rows/columns corresponding to the S subnetwork weights 2 , and 0s otherwise.
Bayesian Inference.
Compute the posterior over the subnetwork via a Laplace approximation:
[Equation]
Firstly, note that the mean of the subnetwork posterior in Eq.
(7) is the MAP estimate w MAP and thus equal to the mean of the full posterior (cid:101) p ( w | y , Φ ) in Eq.
Secondly, note that the covariance matrix of the subnetwork posterior in Eq.
(7) is the element-wise product M S (cid:12) H − 1 , which masks the (co-)variances of all weights not belonging to the subnetwork to zero, eﬀectively making them deterministic.
More precisely, the subnetwork covariance matrix, M S (cid:12)
H − 1 , is a D × D matrix that is equal to the full posterior covariance matrix H − 1 in the rows/columns of the S weights in the subnetwork, and zero in the rows/columns of all other D − S weights.
We consider what we term the posterior gap —the Wasserstein distance 3 (in particular the squared 2-Wasserstein distance) between the posterior distribution over the full network and the posterior distribution over the subnetwork.
Proposition 1 (Posterior Gap) For a subnetwork of size S < D , the Wasserstein gap between the full posterior (cid:101) p ( w | y , Φ ) in Eq.
(6) and the subnetwork posterior p S ( w | y , Φ ) in
Eq.
(7) is:
[Equation]
Proof Note that the posterior distributions (cid:101) p ( w | y , Φ ) and p S ( w | y , Φ ) are both Gaussian.
We thus consider the squared 2-Wasserstein distance between two Gaussian distributions
N ( µ 1 , Σ 1 ) and N ( µ 2 , Σ 2 ), which has the following closed-form expression (Givens et al., 1984) 4 :
[Equation]
Plugging in µ 1 = µ 2 = w MAP , Σ 1 =
H − 1 and Σ 2 = M S (cid:12) H − 1 , we obtain
W [
[Equation]
( w
[Equation]
The optimal subnetwork should then minimize the posterior gap in Eq.
However, for full covariance matrices H − 1 and a large number of weights D , this will generally be infeasible as Eq.
(8) depends on all entries of the D × D
-matrix H − 1 , which is intractable to compute/store.
To derive a practical subnetwork selection strategy, we assume the covariance matrix to be diagonal.
Corollary 2 (Optimality of Maximum Variance Subnetwork Selection under Decorrelation) For a generalized linear model with posterior covariance matrix H − 1 = diag ( σ 21 , . . .
, σ 2 D ) , the optimal subnetwork under the Wasserstein gap is comprised of the S weights with the largest variances
σ 2 d
Proof For H − 1 = diag( σ 21 , . . .
, σ 2 D ), the Wasserstein posterior gap in Eq.
(8) simpliﬁes to
[Equation]
The optimal subnetwork selection strategy amounts to choosing the binary vector m =
[ m dd ] Dd =1 with (cid:80) Dd =1 m d = S
(i.e., we select S out of D parameters) s.t.
the posterior gap in
Eq.
(10) is minimized .
Observing that the contribution of the d -th parameter to the posterior gap is (1 + 1) σ 2 d − 1 × 2 σ 2 d = 0 if it is selected (i.e. if m dd = 1), and (1 + 0) σ 2 d − 0 × 2 σ 2 d = σ 2 d if it is not selected (i.e. if m dd = 0), we see that the optimal subnetwork comprises of the S weights with the largest variances σ 2 d .
Finally, since a GGN-linearized neural network, as in Eq.
(3), corresponds to a GLM with basis expansion φ n = J W MAP ( x n ) =
∂ f ( x n , W MAP ) /∂
W MAP (see Step #3 in Section 2),
Theorem 2 implies that the optimal subnetwork comprises of the weights with the largest variances (assuming decorrelation).
In practice, even just computing the diagonal of the covariance matrix is challenging, so we use a diagonal Laplace approximation which instead computes the inverse of the diagonal of the GGN (see e.g. Ritter et al.
We will also refer to the procedure described in Theorem 2 as Wasserstein pruning in the experiments in Section 3, as it eﬀectively prunes the variances of all weights that do not belong to the subnetwork found via Wasserstein posterior gap minimization.
Finally, note that we only have to make the decorrelation assumption for the purposes of subnetwork selection – when doing posterior inference over the selected subnetwork, we estimate a full covariance matrix for maximal expressiveness, as described in Step #3 in Section 2.
Appendix C. Subnetwork inference in large models vs full inference over small models
We study the following natural question: “Why should one use subnetwork inference in a large model when one can just perform full network inference over a smaller model?”
We explore this by considering 4 fully connected models of increasing size.
These have numbers of hidden layers h d
= { 1 , 2 } and hidden layer widths w d = { 50 , 100 } .
For a dataset with input dimension i d , the number of weights is given by D =(
i d +1)
w d
+( h d − 1) w 2 d .
Ourhidden layer, 100 hidden unit models have a weight count of the order 10 4 .
Full covariance inference in these models borders the limit of computational tractability on commercial hardware.
We ﬁrst obtain a MAP estimate of each model’s weights and our homoscedastic likelihood function’s noise variance.
We then perform full network GGN Laplace inference for each model.
We also use our proposed Wassertein rule to prune every network’s weight variances such that the number of variances that remain matches the size of every smaller network under consideration.
In all cases, we employ the linearized predictive in Eq.
Consequently, networks with the same number of weights make the same mean predictions.
Increasing the number of weight variances considered will thus only increase predictive uncertainty.
We employ 3 UCI (Dua and Graﬀ, 2017) datasets of increasing size (input dimensionality, n. points): wine (11, 1439), kin8nm (8, 7373) and protein (9, 41157).
We consider their
[Figure]
standard train-test splits (Hern´andez-Lobato and Adams, 2015) and their gap variants (Foong et al., 2019b), designed to test for out-of-distribution uncertainty.
For each split, we set aside 15% of the train data as a validation set.
We use these for early stopping when ﬁnding MAP estimates and for selecting the weights’ prior precision.
We keep other hyperparameters ﬁxed across all models and datasets.
Results are in Fig.
We present mean test log-likelihood (LL) values, as these take into account both accuracy and uncertainty.
Larger models tend to perform better when doing MAP inference, with wine-gap and protein-gap being exceptions.
We also ﬁnd larger models improve over their respective MAP LLs more than small ones when performing approximate inference over the same numbers of weights.
We conjecture this is due to an abundance of degenerate directions (weights) in the weight posterior of all models (Maddox et al., 2020).
Full network inference in small models captures information about both useful and non-useful weights.
In larger models, our pruning strategy allows us to dedicate a larger proportion of our resources to modelling informative weight variances and covariances.
In 3 out of 6 datasets, we ﬁnd abrupt increases in LL as we increase the number of weights over which we perform inference, followed by a plateau.
Such plateaus might be explained by all of the most informative weight variances having already been accounted for.
These results suggest that subnetwork inference in a large model is superior to full network inference in a small one.
Appendix D. Details on Image Classiﬁcation Experiments
We use deep ensembles of 5 DNNs, as suggested by (Ovadia et al., 2019), and 16 samples for MC Dropout, diagonal Laplace and SWAG.
We use a Dropout probability of 0 . 1 and a prior precision of λ = 40 , 000 for diagonal Laplace, found via grid search.
For subnetwork inference, we compute the linearized predictive distribution in Eq.
(4) via the closed-form approximation for integrals between Gaussians and multi-class cross-entropy likelihoods described in (Gibbs, 1998).
We use Wasserstein pruning to retain only 0 .
38% of the weights, yielding a subnetwork with only 42,438 weights.
This is the largest subnetwork for which we
[Figure]
CIFAR
[Figure]
MNIST vs Fashion
[Figure]
can tractably compute a full covariance matrix.
Its size is 42 , 438 2 × 4 Bytes ≈ 7 . 2 GB.
We use a prior precision of λ = 500, found via grid search.
Appendix E. Additional Image Classiﬁcation Results
[Table]
Ours
MAP
SWAG
[Table]
, respectively (Nalisnick et al., 2019).
[Figure]
Corrupted
[Figure]
LL
[Figure]
corruption
[Figure]
[Table]
without corruptions or rotations.