Expressive yet Tractable Bayesian Deep Learning via Subnetwork Inference
Authors: Erik Daxberger
Eric Nalisnick
James Allingham
Javier Antoran
José Miguel Hernández-Lobato
Venue:  (2021)

Abstract: The Bayesian paradigm has the potential to solve some of the core issues in modern deep learning, such as poor calibration, data inefficiency, and catastrophic forgetting. However, scaling Bayesian inference to the high-dimensional parameter spaces of deep neural networks requires restrictive approximations. In this paper, 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 would otherwise be intractable for the full model. In particular, we develop a practical and scalable Bayesian deep learning method that first 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 maximally preserve posterior uncertainty. We empirically demonstrate the effectiveness of our approach compared to point-estimated networks and methods that use less expressive posterior approximations over the full network.

Under review as a conference paper at ICLR 2021
EXPRESSIVE YET TRACTABLE BAYESIAN DEEP LEARNING VIA SUBNETWORK INFERENCE
Anonymous authors Paper under double-blind review
ABSTRACT
The Bayesian paradigm has the potential to solve some of the core issues in mod- ern deep learning, such as poor calibration, data inefficiency, and catastrophic forgetting. However, scaling Bayesian inference to the high-dimensional parameter spaces of deep neural networks requires restrictive approximations, In this paper, we propose performing inference over only a small subset of the model parameters while keeping all others point estimates. This enables to use expressive posterior approximations that would otherwise be intractable for the full model. In particular, develop practical and scalable Bayesian deep learning method that first 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 maximally preserve posterior uncertainty. We empirically demon- strate the effectiveness of our approach compared to point-estimated networks and methods that use less expressive posterior approximations over the full network.
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 shift in the train and test distributions (Nguyen et al., 2015) Guo et al. 2017). To reliably inform decision making, DNNs must be able robustly quantify the uncertainty 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 principled way to capture predictive uncertainty via the posterior distribution over model parameters. Unfortunately, due to their non- linearities, exact posterior inference is intractable in DNNs. Despite recent successes in the field of Bayesian deep learning (Blundell et al., 2013] Gal & Ghahramani 2016 Osawa et al., 2019
Maddox al. 2019; Dusenberry al. 2020), existing methods only made scalable 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 (Oadiaa et al.) 2019; Fort et al. 2019 Foong et al. 2019a Ashukha et al. (020a)).
Due to the heavy overparameterization of DNNs, their accuracy well-preserved by small sub- network (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 uncer- tainty quantification. These observations prompt the following question for 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 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 Bayesian deep learning approach that performs Bayesian inference over only small subset the model weights and keeps all other weights deterministic. This allows to use expressive posterior approximations that typically intractable in DNNs.

---

Under review as a conference paper at ICLR 2021
06
d6
Point Estimation
Subnetwork Sclection
(c) Bayesian Inference
(d) Prediction
Figure/ Schematic illustration of our proposed approach. (a) We train neural network using standard techniques to obtain a point estimate of the weights. (b) We identify small subset of the weights. (c) We estimate posterior distribution over the selected subnetwork via Bayesian inference techniques. (d) We make predictions using the full network of mixed Bayesian/detemiiniiie weights.
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 full- covariance Gaussian posterior over subnetwork within point-estimated neural network.
3. We formally characterize the discrepancy between the posterior distributions over a subnet- work 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, various benchmarks, that our method compares favourably against point-estimated networks and other Bayesian deep learning methods, experimentally confirm- ing that expressive subnetwork inference is superior to crude inference full networks.
SUBNETWORK POSTERIOR APPROXIMATION
Bayesian neural networks (BNNs) aim to capture model uncertainty, i.c., uncertainty about the choice of weights W which arises due to multiple plausible explanations y is the dependent variable (e.g. classification label) and X is the feature matrix. A prior distribution pl (W) is specified over the BNN's weights. We wish to infer their full posterior distribution
the training data {y. X). Here,
p( p(Wly. X)
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.c.
/ w px"" W)p(W X X)W.
ply" X"", y. X)
Unfortunately, due to the size of modern deep neural networks, 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 commonly employed (Blundell et al. 2015 Hernández-Lobatt & Adams 2015
Kingma et al. 2015 Khan et al. 2018 Osawa al. 2019), i.e. p(Wly, X) I ((wa) where Wa denotes the d-th weight in the D-dimensional neural network weight vector W E R D (the concatenation and flattening of all layers' weight matrices), Clearly, this is very wishful assumption; In practise, suffers from severe pathologies (Foong 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 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 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 NN's predictive power might be isolated to small subnetwork.
2

---

Under review as a conference paper at ICLR 2021
2. Inference over submodels: Previous work has provided evidence that inference can be effective even when not done on the full parameter space. Izmailov et al. (2019) performed inference over 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 al.) 2018 Ober Rasmussen 2019).
We thus combine these ideas, making the following two-step approximation of the posterior in Eq.
p(W X) ~ p( X) II 8(w, w;) ~ ((Ws) II 6(w, - w,).
The first approximation decomposes the full neural network posterior pl Wly. X) into posterior p( p(Ws)y y. X) over the subnetwork Ws and delta functions 5(w, w over all remaining weights }r. keeping them at fixed values w E IR. This can be viewed as praming the variances of the weights (w,)r zero, which in contrast to ordinary weight pruning methods that set the weights [w,.r themselves to zero. The second approximation is result of posterior inference over the subnetwork still being intractable. In turn, we introduce the approximate distribution q( s). 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 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 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 WAAP. This estimate should respect the Bayesian model given Eq. [T] and therefore we optimize the maximum a-posteriori (MAP) objective:
MMAP = arg maxw [logp(y] X, W) + log (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. (a).
Step #2: Subnetwork Selection. The second step is to identify small subnetwork Ws. Ideally, would like to identify the subnetwork whose posterior is 'closest' the full-network posterior. We formalize this argument in Section 4 Hnn 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. (b) for an illustration of this step.
Step #3: Bayesian Inference. Given the subnetwork point estimate W AP. the Laplace approximation to infer a full-covariance Gaussian posterior distribution over the subnetwork Ws:
W sly. X)
((Ws) s) = N (W W WMAP, --1
(5)
where the posterior covariance matrix H-1 1 E R DDD corresponds the inverse of the average Hessian of the negative log posterior, i.e. H = NE [- poy((1 X W)/awa] AI. Here, the expectation is w.r.t. the data generating distribution and is the precision of zero-mean factorized Gaussian prior p( P(W -W; -1). In practice, approximate the Hessian H with the generalized Gauss-Newton (GGN) matrix H (Schraudolph. 2002). i.e.
i = E n2 HnJn + AI, with J, 0f(zn, W) and H, 02 f2f(xn W) 1=1
82 LL f(In, W))
(6)
Ow
See Scction for thorough discussion of related work.
3

---

Under review as a conference paper at ICLR 2021
where J, ROxD is the Jacobian of the neural network features fz W) RD w.r.t the weights W, and H E ROx0 is the Hessian of the loss Lign. f W)) w.r.t. the features f W The GGN H has clear practical advantages over the Hessian H; see Martens & Sutskever (201] 1) 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 W) at its MAP estimate WMAP,
fin AM (x, W) = f(x, WMAP) + JWMAP (x) - WMAP)
(7)
where J (x) Of (æ. WMAr)/ MAP E ROxD (Immer al., 2020). Note that the model in Eq. (7) is linear in W. as only the term JWMAP ( depends linearly on W, while the other terms constant w.r.t. W and can thus be subsumed into an additive bias term (Khan et al., (209)).
The GGN approximation thus locally turns the underlying probabilistic model from a Bayesian neural network into (generalized) linear model, with basis function expansion J (ar) of covariate
[Immer et al. 2020). Put differently, linearized Laplace in the neural network fx, W) is equivalent
to ordinary Laplace in the linear model A (æz, W) in Eq. (7p. the GGN H corresponding to f(x, W) )) in Eq. (6) is equivalent to the Hessian H corresponding to fM (x, W) in Eq. 7 (Khan cct al. 2019). This is a useful property that will allow us to derive a principled subnetwork selection A This step is illustrated in Fig. (c). We emphasize that this whole procedure (i.e. Steps #1 #1-#3) is perfectly valid mixed inference strategy, performing full Laplace inference over the selected subnetwork and MAP inference over all remaining weights.
strategy in Section
Step #4: Prediction. Given the linearized Laplace approximation over the subnetwork Ws in Eqs. (5) and (6). i.e. ((Ws) = NNWs:: s: W M.P.: 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(yl (ax W)), Immer et al. (2020) recently suggested that, since inference was (implicitly) done in the GGN-linearized model, it is more princi- pled to instead predict using the linearized likelihood Eq. 7 i.e. p(y X, W) plylf f. (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
X" y. X)=
I
W MAP, Ø-1
1[sw, - w() dW. (8)
There are two ways to compute Eq. 8) Firstly, via a Monte Carlo approximation ply" X y, X)
M ΣM= | AM (x*, Wm)) by sampling Wm from ((W M AP., H 1) and II, 8(w, - w
the latter of which is trivial. Secondly, due to linearity of p(y"Iff MAP (X*, W)), there are closed- form 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. (d).
4 PRINCIPLED SUBNETWORK SELECTION FOR LINEAR(IZED) MODELS
We next analyze the subnetwork inference procedure described in Section 3 for the case of generalized linear model (GLM) (Nelder & Baker 1972). which models the expected response you given the basis function expansion of the covariates S2,, = <)(xr.) as
Ely,| - و 9 -w (.).
Here, w E R is the vector of model parameters (which subsumes a scalar bias Bg for notational convenience) and g (-) denotes link function such that g : R + Myld- In particular, we consider Bayesian GLM, by specifying prior distribution p(w) over model parameters and aiming to infer the posterior distribution p( w y, 4) x p(y) Φ w w))(w), where = (01, @N!"
1. Point Estimation. Obtain the MAP estimate, WMAP- = arg max log p(y) 4. w) + log (w). For commonly-used link functions (e.g. the identity in case of a Gaussian likelihood for re- gression, or the sigmoid/softmax function in case of a categorical likelihood for classification) and commonly-used priors (e.g. a Gaussian), the log-posterior x log p w) 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.
p(wly, +) R p( w/ y. +) NW;; WAA H H
(10)
4

---

Under review as a conference paper at ICLR 2021
where H is the Hessian defined in Section Note that for the GLM we consider, the Hessian H is equivalent to the GGN H defined in Eq. meaning that ordinary Laplace approximation is equivalent to a linearized Laplace approximation [Martens] 2016). For the case of an identity link function (i.e. Gaussian likelihood with noise variance og) and Gaussian prior the MAP estimate even has closed-form expression, APP (Φ" + oi Ao) "Tyy Here, the Laplace approximation in Eq. (T)) exuctly corresponds to the true posterior, i.e. p(wly, +) p( (wy,,, b) We will thus refer to the posterior p( ((wly, ②) in Eq. (TO) as the full posterior.
(0 A
2. Subnetwork Selection. Select a subset of S model weights via a method of choice, yielding binary vector
R. where m. = 1 if the d-th weight is part of the subset, and ma 0 otherwise. For convenience, we define the binary mask matrix Mg = mm E RP> which contains 1s in the rows/columns corresponding the S subnetwork weight and Os otherwise.
3. Bayesian Inference. Compute the posterior over the subnetwork via a Laplace approximation: ps ps(wlp
(wly y. +) N(w; WMA Ms ④ H H-)
(11)
1) is the MAP estimate WMAP and thus equal to the mean of the full posterior p( ([wly, +)) in Eq. 1Opp. Secondly, note that the covariance matrix of the subnetwork posterior in Eq. (11 is the element-wise product Mg H-1, which masks the (co-) (co-)variances of all weights not belonging to the subnetwork to zero, effectively making them deterministic. More precisely, the subnetwork covariance matrix, Mg @ 1 is a D x D matrix that is equal to the full posterior covariance matrix H- in the rows/columns of the S weights in the subnetwork, and zero in the rows/columns of all other D weights.
Firstly, note that the mean of the subnetwork posterior in Eq. (11)
We consider what we term the posterior tup-the Wasserstein distance (in particular the squared 2- Wasserstein distance) between the posterior distribution over the full network and the posterior distribution over the subnetwork. The proofs for all results below will be presented in Appendix[A] Proposition 1 (Posterior Gap). For subnetwork of size S < D, the Wasserstein gap between the full posterior p(w( y, ②) in Eq. (10) and the subnetwork posterior ps (w y, 4) in Eq. (II) is:
W [[W]]] y. + ps y. 2) E (1 maa) o3 trace (2 (Ms 2 H 12 (12)
The optimal subnetwork should then minimize the posterior gap in Eq. 12) However, for full covariance matrices H-1 and a large number of weights D, this will generally be infeasible as Eq. (12 depends on all entries of the D x D-matrix H H which is intractable to compute/store. To derive a practical subnetwork selection strategy, we assume the covariance matrix be diagonal. Corollary 1.1 (Optimality of Maximum Variance Subnetwork Selection under Decorrelation). For generalized linear model with posterior covariance matrix H = diag (of....... ....... T.), the optimal subnetwork under the Wasserstein gap is comprised of the S weights with the largest variances o²-
Finally, since a GGN-linearized neural network, as in Eq. 7 corresponds to GLM with basis ex- pansion b J (2x,) (x.. AMAP) /2W M A (see Step #3 in Section B). Corollary [.] implies that under decorrelation, the optimal subnetwork comprises of the weights with the largest variances. In practice, even just computing the diagonal of the covariance matrix is challenging, use a diagonal Laplace approximation which instead computes the inverse of the diagonal of the GGN (see e.g. |Ritter al. (2018). Finally, note that 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 3 In our experiments in Scction we empirically show that making the decorrelation assumption for subnetwork selection but then using full-covariance Gaussian for inference performs significantly better than directly making the decorrelation assumption for inference (e.g. mean-field variational inference, diagonal Laplace).
EMPIRICAL ANALYSIS
We empirically assess the effectiveness of subnetwork inference compared to point-estimated NNs and methods that do less expressive inference over the full network. We consider three tasks: 1) small-scale toy regression, 2) medium-scale tabular regression, and 3) large-scale image classification.
p For consistency, we will keep referring to the S selected linear model weights as a "subnetwork".. 3 We use the Wasserstein distance instead of the more common Kullhack-Leei divergence because the Wasserstein is well-defined for degenerate distributions and is an actual distance metric (i.e. symmetrie).
5

---

Under review as a conference paper at ICLR 2021
Full Cov (2600)
Wass 50% (1300)
Wass 3% (78)
Wass 1% (26)
MAP (0)
2
0-
2-
Diag (2600)
Rand 50% (1300)
Rand 3% (78)
Rand 1% (26)
Final layer (50)
2
-2
-2
-2
-2
Figure 2: Predictive distributions (mean + std) for 1D regression. The numbers in brackets denote the number of parameters over which inference was done (out of 2600 total). Wasserstein-base subnetwork inference maintains richer predictive uncertainties at smaller parameter counts.
5.1 How DOES SUBNETWORK INFERENCE RETAIN POSTERIOR PREDICTIVE UNCERTAINTY?
We first assess how the predictive distribution of full-covariance Gaussian posterior selected subnetwork qualitatively compares to that obtained from 1) a full-covariance Gaussian over the full network (Full Cov), 2) factorised Gaussian posterior covariance Gaussian over only the (Final layer) of the network (Kristiadi et al.) 2020), and 4) point estimate (MAP). For subnetwork inference, we consider both Wasserstein (Wass) (as described in Section 4 and uniform random subnetwork selection (Rand) to obtain subnetworks that comprise of only 50%, 3% and 1% of the model parameters. Note that while for this toy example, we could in principle use the full covariance matrix for the purpose of subnetwork selection, we still just use its diagonal described Section for consistency. Our NN consists ReLU hidden layers with 50 hidden units each. We employ a homoscedastic Gaussian likelihood function where the noise variance optimised with maximum likelihood. We use GGN Laplace inference over network weights (not biases) in combination with the linearized predictive distribution in Eq. (8). Thus, all approaches considered share their predictive mean, allowing us to better compare their uncertainty estimates. All approaches share single prior precision of X = 3.
the full network (Diag), 3) full-
We use a synthetic ID regression task with two separated clusters of inputs (Antorán et al. 2020). allowing probe for 'nnbbewwenn uncertainty (Foong et al.) (2019b). Results shown Fig. Subnetwork inference preserves more of the uncertainty of full network inference than diagonal Gaussian final layer inference while doing inference correlations, subnetwork inference retains uncertainty in between clusters of data. This is true for both random and Wasserstein subnetwork selection. However, the latter preserves more uncertainty with smaller subnetworks. Finally, the strong superiority to diagonal Laplace shows that making a diagonal assumption for subnetwork selection but then using full-covariance Gaussian for inference (as we do) performs much signficantly better than making a diagonal assumption for the inferered posterior directly. These results suggest that expressive inference over a carefully selected subnetwork retains more predictive uncertainty than crude approximations over the full network.
fewer weights. By capturing weight
5.2 SUBNETWORK INFERENCE IN LARGE MODELS VS FULL INFERENCE OVER SMALL MODELS
Secondly, we study the following natural question: "Why should one use subnetwork inference in large model when one can just perform full network inference over smaller model?" We explore this by considering 4 fully connected models of increasing size. These have numbers of hidden layers h 2} and hidden layer widths w (50, 100). For a dataset with input dimension ia. the number of weights is given by D=(f++1>) w ( (ia--))xt 1 Our hidden layer, 100 hidden unit models have a weight count the order 10 Full covariance inference in these models borders the limit of computational tractability commercial hardware. We first obtain MAP estimate of each

---

Under review as a conference paper at ICLR 2021
model's weights and our homoscedastic likelihood function's '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. 77) 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.
kindem
kinlnm-gap
proters
protein-gop
-200
-5000
-2.20
-1109
3.85
2200
1.85
2200
2305
1.45
-0000
130 200
3130 1220
200
200 1200
500
11000
502 1000
postestor dim.
posterior dim
posizrior dim
posterter dim
posterioe dim
postertor cltn
50 Apl
10:00 Apd
5i k:2
vy:1000, Ar2
Figure 3: Mean test log-likelihood values obtained on UCI datasets across all splits. Different markers indicate models with different numbers of weights. The horizontal axis indicates the number of weights over which full covariance inference is performed. 0 corresponds to MAP parameter estimation, and the rightmost setting for each marker corresponds to full network inference.
We employ 3 tabular datasets of increasing size (input dimensionality. n. points): wine (11, 1439). kin8nm (8, 7373) and protein (9, 41157). We consider their standard train-test splits (Hernández Lobato & Adams 2013) and their gap variants (Foong et al. (019b), designed to test for out-of- distribution uncertainty. Details are provided in Appendix C.4 For each split, we set aside 15% of the train data as validation set. We use these for early stopping when finding MAP estimates and for selecting the weights' prior precision. We keep other hyperparameters fixed across models and datasets. Results are in Fig.
We present mean test log-likelihook (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 find 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 al. 2020). Full network inference in small models captures information about both useful and non-useful weights. In larger models, our subnetwork selection strategy allows us dedicate a larger proportion of our resources to modelling informative weight variances and covariances. In out of 6 datasets, we find abrupt increases in LL as we increase the number of weights over which we perform inference, followed by plateau. Such plateaus might be explained by all of the most informative weight variances having already been accounted for. These results suggest that, given the same amount compute, larger models benefit more from subnetwork inference than small ones.
5.3 SCALING TO IMAGE CLASSIFICATION WITH DISTRIBUTION SHIFT
We now assess the robustness of large convolutional neural networks with subnetwork inference to distribution shift on image classification tasks compared to the following baselines: point-estimated networks (MAP), Bayesian deep learning methods that do less expressive inference over the full network: MC Dropout (Gal & Ghahramani] 2016)). diagonal Laplace, VOGN (Osawa et al.) 2019) (all of which assume factorisation of the weight posterior), and SWAG (Maddox et al. 2019) (which assumes diagonal plus low-rank posterior). We also benchmark deep ensembles (Lakshminarayannal
et al. 2017). The latter is considered tathe-of-the-art for uncertainty quantification in deep learning
(Ovadia et al. 2019 Ashukha et al., 2020a). We use ensembles of DNNs, as suggested by 2019 and 16 samples for MC Dropout, diagonal Laplace and SWAG. We use Dropout probability of 0.1 and a prior precision of À = 40, 000 for diagonal Laplace, found via grid search. We apply all approaches to ResNet-18 18 (He et al. (016)). which is composed of an input convolutional block, 8 residual blocks and a linear layer, for total of 11.16,,000 parameters. For subnetwork inference, we compute the linearized predictive distribution in Eq. 8 via the closed-form
(Ovadia et al.
7

---

Under review as a conference paper at ICLR 2021
Rotated MNIST
Corrupted CIFAR10
0.25
emrr
0.30
Oues
Ensemble
Dagelap
0.25
Dogont
SWAG
VOGN
0.00
0.0
-5.0
-7.
0
30
60
90
120
150
180
5
rotation
corruption
Figure 4: Results on the rotated MNIST (left) and the corrupted CIFAR (right) benchmarks of Ovadia et al. (2019). showing the mean + std of the error (top) and log-likelihood (bottom) across three different seeds. Subnetwork inference retains better uncertainty calibration and robustness to distribution shift than point estimated networks and other Bayesian deep learning approaches.
approximation for integrals between Gaussians and multi-class cross-entropy likelihoods described in (Gibbs, 1998). We use Wasserstein subnetwork selection to retain only 0.38% of the weights, yielding a subnetwork with only 42,438 weights. This is the largest subnetwork for which we can tractably compute full covariance matrix. Its size is 42, 4382 precision of X = 500, found via grid search. Finally, to assess to importance of principled subnetwork selection, we also consider the baseline where we select the subnetwork uniformly at random (called Ours (Rand)). We perform the following two experiments, with results in Fig. See Appendix for additional results.
Bytes 7.2 GB. We use prior
Following Ovadia et al. (2019); Antorán et al. 2020 we train all methods on MNIST and evaluate their predictive distributions on increasingly rotated digits. While all methods perform well on the original MNIST test set, their accuracy degrades quickly for rotations larger than 30 degrees. terms of LL, ensembles perform best out of our baselines. Subnetwork inference obtains significantly larger LL values than almost all baselines, including ensembles. The only exception is VOGN, which achieves slightly better performance. It was also observed in Ovadia et al. (2019) that mean-field -ield variational inference (which VOGN also is instance of) is very strong on MNIST, but its performance deteriorates on larger datasets. Subnetwork inference makes accurate predictions in-distribution while assigning higher uncertainty than the baselines to -iitt-i--ittiittioon points. Corrupted CIFAR: Again following Ovadia et al. (2019): Antorán et al. (2020). we train on CIFARIO and evaluate on data subject 16 different corruptions with 5 levels intensity each (Hendrycks & Dietterich, 2019). 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. Even subnetwork inference differentiates itself by being the least overconfident, outperforming all baselines in terms log- log-likelihhod at all corruption levels. Here, VOGN performs rather badly; while this might appear in stark contrast to its strong performance benchmark, the behaviour that mean-field VI performs well on MNIST but not on larger datasets was also observed in Ovadia et al. (2019).
Rotated MNIST:
the MNIST
On both benchmarks, we furthermore find that randomly selecting the subnetwork performs sub- stantially worse than using more principled subnetwork selection strategy. This highlights the importance of the way subnetworks are selected. These results suggest that subnetwork inference results in better uncertainty calibration and robustness to distribution shift than other popular uncertainty quantification approaches.
6 RELATED WORK
Bayesian Deep Learning. There have significant efforts to characterise the posterior distribution
over NN weights p pD)) Hamiltonian Monte Carlo (Neal 1995) remains the golden standard
8

---

Under review as a conference paper at ICLR 2021
for approximate inference in BNNs to this day. Although asymptotically unbiased, sampling based approaches are difficult to scale to the large datasets (Betancourt) 2015). As a result, approaches which find the best surrogate posterior among an approximating family (most often Gaussians) have gained popularity. The first of these was the Laplace approximation, introduced by MMaKaa (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 field assumption ( (Blundell et al. 2013 Hernández-Lobate & Adams 2013 Gal Ghahramani] 2016; Mishkin et al.. 2018 Osawa et al. 2019). This comes at the cost of limited expressivity (Foong et al. (019a) and empirical under- performance (Ovadia et al. 2019 Antorán et al.) 2020) of uncertainty estimates. Our proposed approach recovers predictive postenor 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 & Welling 2016; Maddox et al. 2019 Ritter et al. 2018).
Neural Network Lincarization. In the limit of infinite width, NNs converge to Gaussian process (GP) behaviour (Neal [995] Matthews 2017 Garriga-Alonse et al. 2018). Recently, these results have been extended to finite width BNNs when the surrogate posterior is Gaussian (Khan et al.) 2019). We draw upon these results 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 fixed (Snoek et al. 2015 Riquelme et al. 2018 Ovadia al. 2019 Ober Rasmussen] 2219;) Pinsler itt al. 2019; Kristiadi et al. 2020). These represent a different generalised linear model in which the basis functions are defined by the /-1 1 first layers of a NN. They can also be viewed as a special case of subnetwork inference, in which the subnetwork is simply defined to be the last NN layer.
Inference over Subspaces. The subfield of NN pruning aims to increase the computational efficiency of NNs by identifying the smallest subset of weights which are required to make accurate predictions. Approaches trade-off computational cost with compression efficiency, ranging from those that require multiple training runs (Frankle& & Carbin] 2019) to those that prune hefore training (Wang et al. 2020). Our work differs in that it retains all NN weights but aims to find 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 al., 2020; Dusenberry et al. 2020 This could also be viewed doing inference over an implicit subspace of weight space. In contrast, we propose a technique to find subsets of weights which are relevant to predictive uncertainty, i.c., we we identify axis aligned subspaces. Finally, there have been recent works studying neural network sparsity/ / pruning from a Bayesian perspective
(Ghosh Doshi- Velez 2017 Polson & Ročková 2018, Cui et al. 2020. Louizos et al. 2017 Molchanov et al. 2017] Gomez et al. 2019 et al. 2018). While these seem conceptually related at first glance, their goal is fundamentally different to ours: While those methods aim to perform model selection/ / sparsification by either explicitly or implicitly pruning unnecessary weights, our goal is to make inference more tractable. More precisely, while those sparse Bayesian deep learning methods prune individual weights, we instead just prune the variances of certain weights, which, importantly retains the full predictive power of the full network to retain high predictive accuracy.
7 CONCLUSION
In this paper, we develop a practical and scalable method for expressive yet tractable probabilistic inference in deep neural networks. We approximate the posterior over a subset of the weights while keeping all other weights deterministic. Computational cost decoupled from network size, allowing us to scale expressive approximations, such as full-covariance Gaussian distributions, to real-world sized NNs. Our approach can be applied post-hoc to any pre-trained model, making it particularly attractive for practical use. Our empirical analysis suggests that subnetwork inference 1) is more expressive and retains more uncertainty than crude approximations over the full network, 2) allows us employ larger NNs, which fit broader range of functions, without sacrificing the quality of our uncertainty estimates, and 3) is competitive with like deep ensembles (Laakshminaryynnan al. 2017), on real-world scale problems.
uncertainty quantification methods,
9

---

Under review as a conference paper at ICLR 2021
REFERENCES
Dario Amodei, Chris Olah, Jacob Steinhardt, Paul Christiano, John Schulman, and Dan Mané. Concrete problems in ai safety. arXiv preprint arXiv:1606. 06600, 2016.
Javier Antorán, James Urquhart Allingham, and José Miguel Hermández-Lobato. Depth uncertainty in neural networks, 2020.
Arsenii Ashukha, Alexander Lyzhov, Dmitry Molchanov, and Dmitry Vetrov. Pitfalls of in-domain uncertainty estimation and ensembling in deep learning. In ICLR, 2020a.
Arsenii Ashukha, Alexander Lyzhov, Dmitry Molchanov, and Dmitry P. Vetrov. Pitfalls of in-domain uncertainty estimation and ensembling in deep learning. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020b.
Michael Betancourt. The fundamental incompatibility of scalable hamiltonian monte carlo and naive data subsampling. volume 37 of Proceedings of Machine Learning Research, pp. 533-540. Lille, France, 07-09 Jul 2015. PMLR. URL httt:://proceedd////////// ml www.sress/v3
betancourt15.c htm1
Christopher M Bishop. Puttern recognition and machine learning. springer. 2006.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight Uncertainty in Neural Networks. In Proceedings of The 32nd International Conference on Machine Learning (ICML), pp. 1613-1622, 2015.
Yu Cheng, Duo Wang, Pan Zhou, and Tao Zhang. survey of model compression and acceleration for deep neural networks. arXiv preprint arXiv:1710.09288 2017.
Tianyu Cui, Aki Havulinna, Pekka Marttinen, and Samuel Kaski. Informative gaussian scale mixture priors for bayesian neural networks. arXiv preprint arXiv:2002. 10243, 2020.
Dheeru Dua and Casey Graff. UCI machine learning repository, 2017. URL httt::t/ttcchiii
ics uci. www.udu/m.....c...momm edduml
Michael W Dusenberry, Ghassen Jerfel, Yeming Wen, Yi-an Ma, Jasper Snoek, Katherine Heller, Balaji Lakshminarayanan, and Dustin Tran. Efficient and scalable bayesian neural nets with rank-1 factors. arXiv preprint 2rXiv:2005.07186, 2020.
Angelos Filos, Sebastian Farquhar, Aidan N. Gomez, Tim G. J. Rudner, Zachary Kenton, Lewis Smith, Milad Alizadeh, Arnoud de Kroon, and Yarin Gal. Benchmarking bayesian deep learning with diabetic retinopathy diagnosis. patps://githnn 2019a.
com/ http://w.....o////////////////.. http:///AAA//oodddoeeeeccarrrrsr benchmarks
Angelos Filos, Sebastian Farquhar, Aidan Gomez, Tim GJ Rudner, Zachary Kenton, Lewis Smith, Milad Alizadeh, Arnoud de Kroon, and Yarin Gal. Benchmarking bayesian deep learning with diabetic retinopathy diagnosis. Preprint, 2019b.
Andrew YK Foong, David Burt, Yingzhen Li, and Richard E Turner. On the expressiveness of approximate inference in bayesian neural networks. arXiv, pp. arXiv-1909, 2019a.
Andrew YK Foong, Yingzhen Li, José Miguel Hemández-LLbbtto and Richard E Turner. In-between uncertainty in bayesian neural networks. ICML Workshop on Uncertainty and Robustness in Deep Learning, 2019b.
Stanislav Fort, Huiyi Hu, and Balaji Lakshminarayanan. Deep ensembles: A loss landscape perspec- tive. arXiv preprint arxiv:19/2.02257.
2019. 2019.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference Learning Representations, 2019.
Lex Fridman, Li Ding, Benedikt Jenik, and Bryan Reimer. Arguing machines: Human supervision of black box ai systems that make life-critical decisions. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops, pp. 0-0, 2019.
10

---

Under review as a conference paper at ICLR 2021
Yarin Gal. Uncertainty in deep learning. University of Cambridge, 1:3, 2016.
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pp. 1059,1059, 2016.
Adrià Garriga-Alonso, Carl Edward Rasmussen, and Laurence Aitchison. Deep convolutional networks shallow gaussian processes. In International Conference on Learning Representations, 2018.
Zoubin Ghahramani. Probabilistic machine learning and artificial intelligence. Nature, 521(7553): 452-459, 2015.
Soumya Ghosh and Finale Doshi-Velez, Model selection in bayesian neural networks via horseshoe priors. arXiv preprint 1rXiv 10388, 2017.
Mark N Gibbs. Bayesian Gaussian processes for regression and classification. PhD thesis, Citeseer, 1998.
Clark R Givens, Rae Michael Shortt, et al. A class of wasserstein metrics for probability distributions. The Michigan Mathematical Journal, 31(2):2311244 1984.
Aidan Gomez, Ivan Zhang, Kevin Swersky, Yarin Gal, and Geoffrey E. Hinton. Learning sparse networks using targeted dropout. arXiv preprint arXiv: 1905. 13678, 2019.
Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume 1. MIT Press, 2016.
Priya Goyal, Piotr Dollár, Ross B. Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training ImageNet in 1 hour. CoRR,
2017.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In Proceedings of the 34th International Conference on Machine Learning- pp. 1321-1330. JMLR. org, 2017.
Dearning-Volume Volume 70,
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification. In 2015 IEEE International Conference on Computer Vision, ICCV 2015, Santiago, Chile, December 7-13, 2015, pp. 1026-1034. IEEE
Computer Society, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 770-778, 2016.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness common corruptions and perturbations. arXiv preprint arXiv: 1et.. 2019.
José Miguel Hernández-Lobal and Ryan Adams. Probabilistic backpropagation for scalable learning of bayesian neural networks. International Conference on Machine Learning, pp. 1861-1869, 2015.
Alexander Immer, Macicj Korzepa, and Matthias Bauer. Improving predictions of bayesian neural networks via local lincarization. ICML Workshop on Uncertainty and Robustness in Deep Learning. 2020.
Pavel Izmailov, Wesley Maddox, Polina Kirichenko, Timur Garipov, Dmitry Vetrov, and An- drew Gordon Wilson. Subspace inference for bayesian deep learning. In 35th Conference on Uncertainty in Artificial Intelligence, UAI 2019, 2019.
Mohammad Emtiyaz Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivas- tava. Fast and scalable bayesian deep learning by pertut--etttraation in adam. arXiv preprint arXiv: 1806 04854, 2018.
11

---

Under review as a conference paper at ICLR 2021
Mohammad Emtiyaz E Khan, Alexander Immer, Ehsan Abedi, and Maciej Korzepa. Approximate in- ference turns deep networks into gaussian processes. In Advances in neural information processing systems. pp. 3094-3104, 2019.
Durk Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameterization trick. In Advances neural information processing systems, pp. 2575-2583, 2015.
Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. Being bayesian, even just a bit, fixes overconfidence in relu networks. arXiv preprint arXiv:2002.10116 2020.
Alex Krizhevsky and Geoffrey Hinton. Learning Multiple Layers of Features from Tiny Images. Technical report, University of Toronto, 2009.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances Neural Information Processing Systems, pp. 6402-6413, 2017.
Neil David Lawrence. Variational inference in probabilistic models. PhD thesis, University of Cambridge, 2001.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. tradient-basee learning applied to document recognition. Proceedings of the IEEE, 86(11):228-2324, 1998.
Juho Lee, Saehoon Kim, Jaehong Yoon, Hae Beom Lee, Runho Yang, and Sung Ju Hwang. Adaptive network sparsification via dependent variational beta-bernoulli dropout. 2018.
Ekaterina Lohacheva, Nadezhda Chirkova, Maxim Kodryan, and Dmitry P Vetrov. On power laws in deep ensembles. Advances in Neural Information Processing Systems, 33, 2020.
Christos Louizos and Max Welling, Structured and efficient variational deep learning with matrix gaussian posteriors. International Conference on Machine Learning, pp. 7708-7716, 2016.
Christos Louizos, Max Welling, and Diederik P Kingma. Learning sparse neural networks through 1_0 regularization. arXiv preprint 17200:17 2017.
David JC MacKay. A practical bayesian framework for backpropagation networks. Neural computa- tion, 4(3):448447) 1992.
Wesley Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainty in deep learning. In Advances in Neural Information Processing Systems, pp. 13132-13143, 2019.
Wesley Maddox, Gregory Benton, and Andrew Gordon Wilson. Rethinking parameter counting in deep models: Effective dimensionality revisited. arXiv preprint arXiv:2003.021399 2020.
James Martens. 2016.
Second-order optimization for neural networks. University of Toronto (Canada),
James Martens and Ilya Sutskever. Learning recurrent neural networks with hessian-free optimization. In Proceedings of the 28th international conference on machine learning (ICML-11), pp. 1033- 1040. Citeseer, 2011.
Alexander Graeme de Garis Matthews. Scalable Gaussian process inference using variational methods. PhD thesis, University of Cambridge, 2017.
Aaron Mishkin, Frederik Kunstner, Didrik Nielsen, Mark Schmidt, and Mohammad Emtiyaz Khan. Slang: Fast structured covariance approximations for bayesian deep learning with natural gradient. In Advances Neural Information Processing Systems, pp. 2255,225, 2018.
Dmitry Molchanov, Arsenii Ashukha, and Dmitry Vetrov. Variational dropout sparsifies deep neural networks. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 2498-2507. JMLR. org, 2017.
12

---

Under review as a conference paper at ICLR 2021
Eric Nalisnick, Akihiro Matsukawa, Yee Whye Teh, Dilan Gorur, and Balaji Lakshminarayanan.
Do Deep Generative Models Know What They Don't 't Know? In International Conference on Learning Representations (ICLR), 2019.
Radford M. Neal. Bayesian Learning for Neural Networks. PhD thesis, CAN, 1995. AAINN02676.
John Ashworth Nelder and R Jacob Baker. Generalized Linear Models. Wiley Online Library, 1972.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading Digits in Natural Images with Unsupervised Feature Learning. In NeurIPS Workshop on Deep Learning and Unsupervised Feature Learning, 2011.
Anh Nguyen, Jason Yosinski, and Jeff Clune. Deep neural networks are easily fooled: High confidence predictions for unrecognizable images. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 427-436, 2015.
Sebastian W Ober and Carl Edward Rasmussen. Benchmarking the neural linear model for regression. arXiv preprint arXiv: 992.08416.. 2019.
Kazuki Osawa, Siddharth Swaroop, Anirudh Jain, Runa Eschenhagen, Richard E Turner, Rio Yokota, and Mohammad Emtiyaz Khan. Practical deep learning with Bayesian principles. arXiv preprint 1xXiv:1906.000 02506, 2019.
Yaniv Ovadia, Emily Fertig. Balaji Lakshminarayanan, Sebastian Nowozin, D Sculley, Joshua Dillon, Jie Ren, Zachary Nado, and Jasper Snock. Can you trust your model's s uncertainty? evaluating predictive uncertainty under dataset shift. In Advances in Neural Information Processing Systems, pp. 13969- 13980, 2019.
Robert Pinsler, Jonathan Gordon, Eric Nalisnick, and José Miguel Hernández-Lobato. Bayesian batch active learning as sparse subset approximation. Advances in Neural Information Processing Systems, pp. 6359-6370, 6370, 2019.
Nicholas G Polson and Veronika Ročková. Posterior concentration for sparse deep learning. In Advances in Neural Information Processing Systems, pp. 930-999--900 930-941,220 2018.
Carlos Riquelme, George Tucker, and Jasper Snoek. Deep bayesian bandits showdown: An empirical comparison of bayesian deep networks for thompson sampling. In International Conference on Learning Representations, 2018.
Hippolyt Ritter, Aleksandar Botev, and David Barber. scalable laplace approximation for neural networks. In International Conference on Learning Representations, 2018.
Simone Rossi, Sebastien Marmin, and Maurizio Filippone. Walsh-hadamanl variational inference for bayesian deep learning. arXiv preprint arXiv: 1rXiv::905.1 //48, 2019.
Nicol N Schraudolph. Fast curvature metrix-vector products for second-order gradient descent.
Neural computation, 14(7):172--1738, 2002.
Jasper Snoek, Oren Rippel. Kevin Swersky, Ryan Kiros, Nadathur Satish, Narayanan Sundaram, Mostofa Patwary, Mr Prabhat, and Ryan Adams. Scalable bayesian optimization using deep neural networks. In International conference on machine learning, pp. 2171-2180, 2015.
Jakub Swiatkowski, Kevin Roth, Bastiaan S Veeling, Linh Tran, Joshua V Dillon, Stephan Mandt, Jasper Snock, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. The k-tied normal distribution: A compact parameterization of gaussian mean field posteriors in bayesian neural networks. arXiv preprint arXiv:2002.02655: 2020.
Chaoqi Wang, Guodong Zhang, and Roger Grosse. Picking winning tickets before training by preserving gradient flow. In International Conference on Learning Representations, 2020. URL ew net / Forrumiid=Sggsstteerer
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-MNIST: novel image dataset for bench- marking machine learning algorithms. 2017.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In Richard C. Wilson, Edwin R. Hancock, and William A. P. Smith (eds.), Proceedings of the British Machine Vision Conference 2016, BMVC 2016, York, UK, September 19-22, 2016. BMVA Press, 2016.
13

---

Under review as a conference paper at ICLR 2021
A PROOFS FOR THE THEORETICAL RESULTS
We now provide the proofs for the results Section
A.1 PROOF OF PROPOSITION
I
Proof. Note that the posterior distributions p( w y, +) and ps( (, +) are both Gaussian. We thus consider the squared 2- 2-Wasserstein distance between two Gaussian distributions. M(n. Σ1) and N2, E2), which has the following closed-form expression (Givens et al. 1984
W (W Σ1) W (M2 Σz) = |1 Hell trace Σ
Σg - 2 2E,E
(13)
Plugging in μn = H2 = WMAP, E₁ H and E₂ Ms 3 H we obtain
W [ w/w/ X) ps(yy)
W [W(WMAP, H 1 Ms
|WaAP W WMr trace 17-1-(M5 2 @ (-1
trace ((1 + Ms) o =1 trace (2 (Ms o 11 11
maa) ag trace (2(H-1 (Ms 2 (-
A.2 PROOF OF COROLLARY
Proof. For H diag( o ..... .f)......,)............... the Wasserstein posterior gap in Eq. (12) simplifies to
D
piwly. b) 11 6)
Σ ((1 maa) o3 2mada
(14)
d=1
The optimal subnetwork selection strategy amounts to choosing the binary vector m madd d=1 with
ED my S (i.c., we select S out of D parameters) Observing that the contribution of the d-th parameter to the posterior gap is (1 1)a 1 2o
the posterior gap in Eq. 14 is minimized.
it is selected (i.e. if mas = 1), and (1 + +)oz - 0 x 2o = =o if is not selected (i.e. if mad = 0),
we see that the optimal subnetwork comprises of the S weights with the largest variances oz
B ADDITIONAL IMAGE CLASSIFICATION RESULTS
Table 1: AUC-ROC scores for out-of-distribution detection, using CIFARIO SVHN and MNIST vs FashionMNIST as in- (source) and out-of -iit-----itriuttion (target) datasets, respectively (Nalisnick
al. 2019).
SOURCE
TARGET
OURS
OURS (RAnd)
DROPOUT
DIAG-LAP
ENSEMBLE
MAP
SWAG
CIFARIO
SVII
0.8510.03
0.8610.02
0.8510.01
0.862.002
0.01
0.8610.02
0.80
MNIST Fashion
0.9216.05
0.7510.02
0.8210.12
0.75.0..1
0.90.1..9
0.7210.03
0.97
Table 2: MNIST rotation.
OURS
OURS (RANd)
DROPOUT
DIAG-LAP
ENSEMBLE
VOGN
LL
0.07
0.014..000
+0.01
-0.0120.0
-0.0140.00
-0.01
0.01+0.00
Liman
or1or
0.01
0.00+...00
0.00
0.01=0.01
0.00+a.00
0.00
0.00#6.00
0.01 sman
ECE
0.05=20.01
0.00+a..0
0.00
0.002..00
0.00#0.00
0.0020.00
0.00+..00
0. 10
brier
0.02
0.01#0.00
0.01
0.02+0.01
0.0140.00
0.01
0.01#0.00
0.04
"This also holds for case of degenerate Gaussian with singular covariance matrix Givens et al.
1984
14

---

Under review as a conference paper at ICLR 2021
MNIST Rotation
CIFARIO Corruption
005
-2
-20
-25
30
N
N
Figure 5: Rotated MNIST (left) and Corrupted CIFAR1O (right) results for deep ensembles (Lak shminarayanan al. 2017) with large numbers of ensemble members (i.e. up to 55). Horizontal axis denotes number ensemble members, and vertical axis denotes performance in terms log- likelihood. Straight horizontal lines correspond to the performance of our method, Colors denote different levels of rotation (left) and corruption (right). It can clearly be observed that the performance of deep ensembles saturates after around 15 ensemble members, meaning that adding more members yields strongly diminishing returns. This is in agreement with recent works (Antorán al.) 2020 Ashukha et al., 2020al [Lobacheva et al. 2020). Our method significantly outperforms even very large deep ensembles, especially for high degrees of rotation/corruttin...
reference.
CIFARIO SVHN
MNIST Fashion
1.0
1.0
0.8
accuracy
0.8
0.6
0.4
0.6
0.2
0.0
0
25
50
75
100
25
50
75
100
% rejected
% rejected
Ours
Dlag-Lap
Dropout
(Rand)
Ensemble
MAP
SWAG
Figure 6: Rejection-classiffattien plots. We simulate realistic OOD rejection scenario (Filos et al. 2019b) by jointly evaluating our models on an in-distribution and an OOD test set. We allow our methods reject increasing proportions of the data based on predictive entropy before classifying the rest. All predictions on OOD samples are treated as incorrect. Following (Nalisnick et al. 2019), we use CIFARIO vs SVHN and MNIST FashionMNIST as in- and out-of diitt--iittiitttion datasets, respectively. Note that the SVHN test set is randomly sub-sampled down to size of 10,000.
15

---

Under review as a conference paper at ICLR 2021
Rotated MNTST
Corrupted CIFAR10
erurr
Ours
Ersemble
Diag-Lap
Dropout
0.1
Ours (Rand)
VOGN
110
15
L1
20
-6
-30
-8
0.35
0.6
0.25
0.4
0.20
0.15
0.2
0.30
0.00
1.2
0.6
brikr score
1.0
0.5
0.2
0.0
0.1
30
60
120
150
180
rotation (*)
corruption
Figure 7: Full MNIST rotation and CIFARIO corruption results, for ResNet-18, reporting predictive error, log-likeel likelihood (L..), expected calibration error (ECE) and brier score, respectively (from top to
bottom).
16

---

Under review as a conference paper at ICLR 2021
0.8
0.6
LI
0.0
30
60
90
120
150
180
30
50
90
120
150
180
1.5
brier score
1.0
ECE
0.4
0.0
0.0
30
60
90
120
150
180
0
60
120
150
180
rotation (*)
rotation (?)
Figure 8: MNIST rotation results for ResNet-50, reporting predictive error, log- likeliidl (LL), expected calibration error (ECE) and brier score. We choose a subnetwork containing only 0.167% (39,190/ 23, 23,466,560) of the parameters of the full network. We see that subnetwork inference still results in an improvement in the calibration of predictive uncertainty. As expected, however, for ResNet-50 the improvement over MAP smaller than for ResNet- 18 where were able choose a subnetwork containing 0.38% of the parameters.
Table 3: MNIST - 15t rotation.
OurS
OuRS (RAnd)
DROPOUT
DIAG-LAv
ENSEMBLE
MAP
SWAG
VOGN
I.I.
-0.1420.0
-0.0520.00
0.05
-0.1 -0.1120.08
-0.04.0.00
-0.0520000
0..42.0.00
-0.19anan
0.02
0.02+..00
0.0320.02
0.02
0.01+0.00
0.02
ECE
0.08
0.001:0000
0.0029.00
0.0120.01
0.0010.00
0.0020000
0.001.0.00
0.12
brier score
0.05
0.03+a.00
0.02
0.05+0.00
0.02+0.00
0.02
0.02:0.00
0.07
Table 4: MNIST 30 rotation.
OURS
OURS (RANd)
DROPOUT
DIAG-LAP
ENSEMBLE
VOGN
LL
0.4220.04
0.36+a.01
0.322s.00
-0.4120.08
-0,28+0.02
-0.39
0.30=a..0
0.51
error
0.11
0.10+a.00
0.09
0.12#0.01
0.08+a.01
0.10a0.00
0.08
0.14
ECE
0.10
0.01.a.01
0.03
0.0620.01
0.02#0.00
0.05
0.01+0.00
0.13
brier
0.1920
0.1622.00
0.14g0.0
0.182002
0.12g0mm
0.162000
0.1240.00
0.23
Table 5: MNIST 45 rotation. ENSEMBLE
OURS
OURS (RANd)
DROPOUT
DIAG-LAP
MAP
SWAG
VOGN
midrule LL
1.099...0
1.00eo.cs
1.44so.11
1.68ss.20
1.36a..07
1.75a..00
1.3550..c
1.15mtmm
error
0. 0.3
0.35
0.33
0.35
0.31
0.35
0.29
0.40
0.03
0.2220001
0.1920.0
0.2240.m
0.1720.01
0.23
0.1820
brier
0.49:4..2
0.55
0.52s0.02
0.35:0.04
0.482n.02
0.56s0.02
0.46:0.0:
0.53tmam
Table 6: MNIST 600 rotation.
OURS
OURS (RAND)
DROPOUT
DIAG-LAP
ENSEMBLE
MAP
SWAG
VOGN
LL
-2, 2.10ss0.0
3,85+a.18
3.5
-4.1120.08
3.600a..10
4.29
-2 2.952a.0s
1.92aman
0.63
0.63+..01
0.6240.01
0.62:0.s
0.61:0.1
0.63
0.53
0.64
ECE
0.25
0.1648.02
0.43:0.02
0.4700.00
0.12a0.01
0.4820.02
0.36+0.02
0.172mam
brier score
0.85.20.z
1.04ga.03
1.00
1.05g0.10
0.984..02
1.07
0.86g..03
0.80aman
17

---

Under review as a conference paper at ICLR 2021
Table 7: MNIST 75 rotation.
OuRS
OURS (RANd)
DROPOUT
DIAG- LAP
ENSEMBLE
MAP
SWAG
VOGN
LL error
=.02
-5.9320.28
55.99
-6.92+0.22
-5.74
-6.63
4.46+0.18
-2.44tmm
0.80=0.02
0.79
0.79
0.81+0.00
0.78
0.79
0. 0.72+0.02
0.77
ECE
2.01
0.622..03
a.
0.6520.01
0.58
0.64
0.51
0.262man
brier
1.0
1.34.0.04
1.30ao.02
1.39+0.11
1.29
1.37
1.17+0.04
0.95
Table 8: MNIST - -90" rotation.
OURS
OURS (RAND)
DROPOUT
DIAG-LAP
ENSEMBLE
MAP
SWAG
VOGN
LL
-3 -3.35
16.66%..15
-6 -6.18
-7 -..32
-6.39
-7.18
- 5.63
- -2.11 imam
error
0.8120.02
0.81+..01
0.81a0.01
0.85+0.01
0.84
0.81s0.01
0.82+a.02
2ma
ECE
0.4320.08
0.64
0.02so.01
0.66+0.00
0.62#a.01
0.6620.04
0.60+:0.01
0.29
brier score
1.1300.0
1.40+a.05
1.37
1.11.0.0
1.3640.01
1-43a0.05
1.3/1..02
1.02
Table 9: MNIST 105" rotation.
OURS
OURS (RANd)
DROPOUT
DIAG-LAP
ENSEMBLE
MAP
SWAG
VOGN
LL.
-0.50
-7.061..45
-6.7020.52
-7.6920.88
-7.01+0.17
-7.87
-6.2810.19
33.10mmm
errer
0.85=2002
0.84s..02
0.81s0.01
0.85+0.01
0.84:0.01
0.81s0.02
0.81+0.00
0.81
ECE
0.47
0.67
0.63
0.67+0.05
0.64
0.68
0.61#0.01
0.34
brier
1.17=0.00
1.44+0.07
1.38
1.41:004
1.40+0.01
1.46a0.07
1.31+0.02
1.07
Table 10: MNIST 120 rotation.
OURS
OURS (RANd)
DROPOUT
DIAG-LAP
ENSEMBLE
MAP
SWAG
VOGN
L.L
-3 -..8320.o
6.73
6.6220.59
-7 -7.9220.5
-6.73
-.533
- 6.44920.3
- 3.07
error
0.8020.0z
0.79
0.78a0.m
0.8110001
0.7810.0
0.79
0.762.0.2
0.762mon
ECE
0.40
0.62
0.58
0.6510004
0.59
0.63
0.5810.05
0.30
brier score
1.100..0
1.35+a.07
1.29
1.3920.00
1.30
1.36
1.27+0.04
1.01aman
Table 11: MNIST 135 rotation.
OuRS
OURS (RAnd)
DROPOUT
DIAG-LAv
ENSEMBLE
MAP
SWAG
VOGN
3.224.000
6.4311..6
-6.662022
-7.552088
-6.5720.10
-7.24
-6.4020.37
-2.80
error ECE
20.02
0.71
0.70ao.0
0.7120.01
0.70g9.01
0.71 a.00
0.70a...2
0.672mam
0.32
0.55
0.52
0.5620.02
0.52
0.56
0.5320.02
0.25:tman
brier score
0.0920.07
1.2120.05
1.17
1.222003
1.1720m
1.233000
1.1820.0
0.94
Table 12: MNIST 150 rotation.
-100
OURS
OURS (RAND)
DROPOUT
DIAG-LAP
ENSEMBLE
MAP
SWAG
VOGN
LL error
03.25
0.5
6.02s0.33
-7 -7.0120.00
-6.884..11
7.41 a0.1
6.39
2.69
0.63320.0z
0.63
0.63
0.6529.01
0.62a0.01
0.63
0.632..01
0.60
ECE
0.29
0.5020.01
0.48
0.5210.1
0.4810.1
0.51
0.491.001
0.23
brier score
0.0222.oz
1.10su.02
1.07a0.0
1.1320.0
1.06
1.11
1.08
0.85
Table 13: MNIST 165e rotation.
OURS
OURS (RAnd)
DROPOUT
DIAG-LAP
ENSEMBLE
MAP
SWAG
VOGN
LL error
-3. 3.42
-7.0140.15
-7.08
-7.80s0.12
-7.51+0.1
-7.91
6.6322.22
2.67
0.58
0.58
0.58
0.5840.00
0.57
0.58
0.50
0.56
ECE
0.32=0002
0.19+a.01
0.48
0.4920.01
0.48+0.00
0.51
0.48+0.00
0.25
brier
0.9022.0z
1.05a0.01
1.04ao.0
1.052001
1.03
1.07
1.03+0.1
0.82
Table 14: MNIST 180 rotation.
OURS
OURS (RAND)
DROPOUT
DIAG-LAP
ENSEMBLE
MAP
SWAG
VOGN
LL error
-3.32
6.63+
6.87
-7. 7.100a0.47
-7.1640.18
-7.33
6.61+9.22
- -2.71
0.5622001
0.561.0.11
0.5610.00
0.552001
0.555..00
0.5620.0
0.57:4..0
0.55
ECE
0..00
0.4644.01
0.4540.00
0.4640.00
0.4640.01
0.48
0.4749.01
0.25
brier score
0.8622.02
1.00++0.01
0.99
0.9920.01
0.99+...0
1.01
1.01a0.01
0.82aman
18

---

Under review as a conference paper at ICLR 2021
Table 15: CIFARIO no corruption.
OURS
OURS (RAnd)
DROPOUT
DIAG-LAP
ENSEMBLE
MAP
SWAG
VOGN
LL
=0.07
0.443t0.01
-0.37
-0.50s0.02
-0.21+0.01
-0.46
0.48
imma
crrer
0.09g0.00
0.08
0.08
0.09+0.00
0.06+0.00
0.08
0. 0.11
0.21 2man
ECE
0.01
0.06
0.04
0.06+0.00
0.01
0.06
0.07
i.03:
brier score
0.13220.0
0.14+...0
0.13+0.00
0.15+0.00
0.09
0.1
0.17+0.00
0.30
Table 16: CIFAR10 level corruption.
OurS
OURS (RANd)
DROPOUT
DIAG- LAP
ENSEMBLE
MAP
SWAG
VOGN
2051
-0.91.0.01
0.80
1.0020..
0.50
-0.96
0.80
0.90
error
0.17=0001
0.16+:000
0.1620.00
0.1740.00
0.13#0.00
0.1620.00
0.17440.0
0.32#mmm
ECE
0.0320000
0.1 0.11
0.10a0.00
0.1320.00
0.04
0.12ao.
0.11
0.03gman
brier score
0.24g0.00
0.27+0.00
0.25
0.29
0.19
0.27
0.29
0.44
Table 17: CIFARIO level 2 corruption.
OURS
OURS (RAnd)
DROPOUT
DIAG-LAP
ENSEMBLE
MAP
SWAG
VOGN
LL
-07330o.0
-1.29
-1 1.20a0.02
-1.0a0.12
-0.80+a.01
1.40.o.00
-1.21=0.00
-1 1.31
error
0.23
0.22+0.01
0.2230.00
0.231001
0.19:
0.22
0.22:0.00
0.40:nan
ECE
0.0620.00
0.16#:.01
0.11s0.00
0.1740.01
0.07
0.1620.00
0,15+0.00
0.10
brier score
0.33
0.37
0.35
0.4020.02
0.28
0.37
0.37.0.00
0.56
Table 18: CIFARIO- 3 corruption.
OURS
OURS (RAnd)
DROPOUT
DIAG-LAP
ENSEMBLE
MAP
SWAG
VOGN
LL
1.0620.02
-2.064..12
1.85
2.13310.1
-1.2840.03
-2.18
1.63#0.00
1.83
eITor
0.32
0.31+:0.11
0.31
0.31a001
0.28+0.00
0.31 a0.ot
0.28
0.51 sman
ECE
0.11
0.24+0.01
0.21
0.21s001
0.12#0.00
0.21s0.01
0.20+0.00
0.19mmm
brier score
0.46.20.1
0.542.0.2
0.50ao.02
0.5420.03
0.42g9.00
0.54
0.47a0.01
0.72
Table 19: CIFARIO level 4 corruption.
OURS
OURS (Rand)
DROPOUT
DIAG-LA
ENSEMBLE
MAP
SWAG
VOGN
1.2252000
-2.4310.18
22.2810
-.5422018
-1.5612..5
-0.57
-1.9520.00
11.9
error
0.36
0.354.0.01
0.35
0.35g..01
0.3240.01
0.35
0.324..00
0.54
ECE
0.1320.01
0.27
0.24
0.2720.01
0.14
0.27
0.231..00
0.22.tmon
brier
0.51
0.60+a.05
0.57
0.6120.02
0.47#0.01
0.60sa.00
0.53+0.00
0.762mam
Table 20: CIFARIO - level 5 corruption.
OURS
OURS (RAnd)
DROPOUT
DIAG-LAP
ENSEMBLE
MAP
SWAG
VOGN
LL error
-1 17700000
2.82+0..11
2.71
-3. 3.200a.13
-1.8840.0
-3.03
-2.31=0.00
2.00amam
0.41
0.40+4.01
0.40ao.01
0.412001
0.37+
0.40s0.00
0.36
0.54
ECE
0.16220.z
0.31+9.01
0.28#0.01
0.3340.02
0.1749.01
0.31
0.2740.01
0.192mmm
brier score
0.58
0.69
0.65
0.72+0.0a
0.55+0.01
0.69
0.61+0.01
0.75mman
19

---

Under review as a conference paper at ICLR 2021
C EXPERIMENTAL SETUP
C.I TOY EXPERIMENTS
We train single, 2 hidden layer network, with 50 hidden ReLU units per layer using MAP inference until convergence. Specifically, we use SGD with a learning rate of x 10-3, momentum of 0.9 and weight decay of 1 x 10-2 We use batch size of 512. The objective we optimise is the Gaussian log-likelihood of our data, where the mean outputted by the network and the the variance is a hyperparameter learnt jointly with NN parameters by SGD. This variance parameters is shared among all datapoints. Once the network trained, we perform post-hoc inference on it using different approaches. Since all of these involve the linearized approximation, the mean prediction is the same for all methods. Only their uncertainty estimates vary.
Note that while for this toy example, could in principle use the full covariance matrix for the purpose of subnetwork selection, we still just use its diagonal (as described in Section ( for consistency. We use GGN Laplace inference over network weights (not biases) in combination with the linearized predictive distribution in Eq. (8). Thus, all approaches considered share their predictive mean, allowing us to better compare their uncertainty estimates.
All approaches share a single prior precision of X = 3. We chose to select the prior precision such that the full covariance approach (optimistic baseline) presents reasonable results. We use the same value for all other methods. We first tried a precision of and found the full covariance approach to produce excessively large errorbars (covering the whole plot). value of 3 produces more reasonable results.
Final layer inference is performed by computing the full Laplace covariance matrix and discarding all entries except those corresponding to the final layer of the NN. Results for random nub-network selection are obtained with a single sample from a scaled uniform distribution over weight choice.
C.2 UCI EXPERIMENTS
In this experiment, our fully connected NNs have numbers of hidden layers (1, 2} and hidden layer widths (50 100}. For a dataset with input dimension id, the number of weights is given by (1 (ha-1)ww 1. Our 2 hidden layer. 100 hidden unit models have weight count of the order 10 The non- con-lincarity used is ReLU.
We first obtain MAP estimate of each model's weights. Specifically, we use SGD with a learning rate x 10-3, momentum of 0.9 and weight decay of x 10-4. We use a batch size of 512. The objective optimise is the Gaussian log- likelihood of our data, where the mean is outputted by the network and the the variance is a hyperparameter learnt jointly with NN parameters by SGD.
For each dataset split, we set aside 15% of the train data as validation set. We use these for early stopping training. Training runs for a maximum of 2000 epochs but early stops with a patience of 500 if validation performance does not increase. For the larger Protein dataset, these values are 500 and 125. The weight settings which provide best validation performance are kept.
We then perform full network GGN Laplace inference for each model. We also use our proposed Wassertein rule together with the diagonal Hessian assumption to prune every network's 's weight variances such that the number of variances that remain matches the size of every smaller network under consideration. The prior precision used for these steps is chosen such that the resulting predictor's s logliklihood performance on the validation set maximised. Specifically, we employ grid search over the values: X [0.0001, 0.001, 0.1, 0.5, 1, 2, 5, 10, 100, 1000). In all cases, we employ the linearized predictive in Eq. 77. Consequently, networks with the same number weights make the same mean predictions. Increasing the number of weight variances considered will thus only increase predictive uncertainty.
C.3 IMAGE EXPERIMENTS
The results shown in Section 5.3 and Appendix B are obtained by training ResNet-18 (and ResNet-50) models using SGD with momentum. For each experiment repetition, we train different models: The first is for: 'MAP', 'Qurs', 'Ours (Rand)', 'SWAG', 'Diag-Laplace" and as the first element of
20

---

Under review as a conference paper at ICLR 2021
'Ensemble'. We train 4 additional 'Ensemble' elements, network with 'Dropout', and, finally 1 network for "VOGN'. The methods 'Ours', "Ours (Rand)', 'SWAG', and 'Diag- Laplace are applied post training.
For all methods except 'VOGN' we use the following training procedure. The (initial) learning rate, momentum, and weight decay are 0.1, 0.9, and x 10-4, respectively. For 'MAP" use 4 Nvidia P100 GPUs with a total batch size of 2048. For the calculation of the Jacobian in the subnetwork selection phase we use single P100 GPU with batch size of 4, For the calculation of the hessian we use a single P100 GPU with a batch size of 2. We train on 1 Nvidia P100 GPU with a batch size of 256 for all other methods. Each dataset is trained for a different number of epochs, shown in Table 21 We decay the learning rate by factor of 10 at scheduled epochs, also shown in Table []] Otherwise, all methods and datasets share hyperparameters. These hyperparameter settings are the defaults provided by PyTorch for training on ImageNet. We found them to perform well across the board. We report results obtained at the final training epoch. We do not use separate validation set to determine the best epoch as we found ResNet-18 and ResNet-50 to not overfit with the chosen schedules.
Table 21: Per-dataset training configuration for image experiments.
DATASET
No. EPOCTS
LR SCHEDULE
MNIST
90
40, 70
CIFARIO
300
150, 225
For 'Dropout', we add dropout to the standard ResNet-50 model (He et al. 2016) in between the and 3 convolutions in the bottleneck blocks. This approach follows Zagoruyka ko & Komodakis (2016) and AAshukha et al. (2020b) who add dropout in in-between the two convolutions of WideResNet-50's
basic block. Following Antorán et al. (2020), we choose dropout probability of 0.1, as they found to perform better than the value 0.3 suggested by |Ashukha et al. (2020b). We use 16 MC samples for predictions. 'Ensemble" uses 5 elements for prediction. Ensemble elements differ from each other their initialisation, which sampled from the He initialisation distribution (He et al.) 2015). We do not use adversarial training as, inline with AAshukha et al. (2020b), we do not find it to improve results. For 'VOGN' we use the same procedure and as used by Osawa et al. (209)) in their CIFAR1O experiments, with the exception that we use a learning rate of 1 x 10 as we we found a value of x 10 not to result in convergence. We train on a single Nvidia P100 GPU with batch size of 256. See the authors' GitHub for more details: r.thubp.ccmmeeaarrrpppoororrrr/r with batet/hhbbyylllbb/////o// eistributed/cae http:///ee..ee.e..oo//ooo//////a voogn vogn bs256 8gpu.json
We modify the standard ResNet-50 and ResNet-18 architectures such that the first 7 x 7 convolution is replaced with a 3 x 3 convolution. Additionally, we remove the first max-pooling layer. Following Goyal et al. (2017). we zero-initialise the last batch normalisation layer in residual blocks so that they act as identity functions at the start of training.
At test time, we tune the prior precision used for 'Ours', DaplaeL' and 'SWAG' approximation on validation set for each approach individually, as in Ritter et al. (2018); Kristiadi et al. (202)). We use a grid search from 1 x 10-4 to x 10a in logarithmic steps, and then a second, finer-grained grid search between the two best performing values (again with logarithmic steps).
C.4 DATASETS
The 1d toy dataset used in Section 5.1 was taken from Antorán et al. (2020). We obtained it from the authors' github repo: https: /github.....g.. com/ htt::/ammbrirrre Lolg/DDUN Table 22 summarises the datasets used in Section 5.2 Wine and Protein are available from the UCI dataset repository [uua & Graff (201)). Kin8nm is available from wwtps:tppp// //t:://wwww www. openm1 hatpm///r///////////// Foong et al. (2019b). For the standard splits (Hernández-Lobaaee training and 10% for validation. For the gap splits (Foong et al., 2019b) a split is obtained per input dimension by ordering points by their values across that dimension and removing the middle 33% of
& Adams 2015)) 90% of the data is used for
21

---

Under review as a conference paper at ICLR 2021
Table 22: Datasets from tabular regression used in Section 5.2
Dataset
N Train
N Val (15% train)
N Test
Splits
Output Dim
Output Type
Input Dim
Input Type
Wine
1223
216
160
20
Continous
11
Continous
Wine Gap
906
161
532
I
Continous
11
Continous
Kin8nm
6267
1106
819
20
I
Continous
8
Continous
Kin8nm Gap
4642
820
2730
I
Continous
8
Continous
Protein
34983
6174
4573
I
Continous
9
Continous
Protein Gap 25913
4573
15244
9
1
Continous
9
Continous
the points. These are used for validation. The datasets used for our image experiments are outlined in Table 23
Table 23: Summary of image datasets. The test and train sizes shown in brackets, e.g. (test & train).
NAME
SIZE
INPUT DIM.
No. CLASSES
No. SPLITS
MNIST LeCun et al.
1998
70,000 (60.000 10,000)
784 (28 x 28)
Fashion-MINIS
2017
70,000 (60,000 & 10,000)
784 (28 28)
CIFARIO Krizhevsky & Hinion 2009
60.000 (50,000 & 10,000)
3072 (32 32 x 3)
10
2
SVHN
Neizer et al. 2011
99,289 (73,257 & 26,032)
3072 (32 x 32 x 3)
10
22


Reviews:
[{"review_id": "xrYSjH7Nol5", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "#### Summary\n\nThe 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.\n\n#### Strengths\n\n- Scalable approximate inference for Bayesian models is an important research area.\n- Expressive, diverse approximate posteriors is also an important area of research, especially given the limitations of mean-field VI and recent literature.\n- The proposed method demonstrates better results for robustness to dataset shifts, as well as in-between uncertainty that mean-field VI misses.\n- In general, the paper is well-written, clearly-motivated, includes both theoretical and empirical results, and adequately compares to, or discusses, relevant literature in the space.\n\n#### Weaknesses\n\n- The authors push on the idea of *scalable* approximate inference, yet the largest experiment shown is on CIFAR-10. Given this focus on scalability, and the experiments in recent literature in this space, I think experiments on ImageNet would greatly strengthen the paper (though I sympathize with the idea that this can a high bar from a resources standpoint).\n- As I noted down below, the experiments currently lack results for the standard variational BNN with mean-field Gaussians. More generally, I think it would be great to include the remaining models from Ovadia et al. (2019). More recent results from ICML could also useful to include (as referenced in the related works sections).\n\n#### Recommendation\n\nOverall, I believe this is a good paper, but the current lack of experiments on a dataset larger than CIFAR-10, while also focusing on scalability, make it somewhat difficult to fully recommend acceptance. Therefore, I am currently recommending marginal acceptance for this paper.\n\n#### Additional comments\n\n- p. 5-7: Including tables of results for each experiment (containing NLL, ECE, accuracy, etc.) in the main text would be helpful to more easily assess\n- p. 7: For the MNIST experiments, in Ovadia et al. (2019) they found that variational BNNs (SVI) outperformed all other methods (including deep ensembles) on all shifted and OOD experiments. How does your proposed method compare? I think this would be an interesting experiment to include, especially since the consensus in Ovadia et al. (2019) (and other related literature) is that full variational BNNs are quite promising but generally methodologically difficult to scale to large problems, with relative performance degrading even on CIFAR-10.\n\n##### Minor\n- p. 6: In the phrase \"for 'in-between' uncertainty\", the first quotation mark on 'in-between' needs to be the forward mark rather than the backward mark (i.e., $`in-between'$).\n- p. 7: s/out of sitribution/out of distribution/\n- p. 8: s/expensive approaches 2) allows/expensive approaches, 2) allows/\n- p. 8: s/estimates 3) is/estimates, and 3) is/\n- In the references:\n  - Various words in many of the references need capitalization, such as \"ai\" in Amodei et al. (2016), \"bayesian\" in many of the papers, and \"Advances in neural information processing systems\" in several of the papers.\n  - Dusenberry et al. (2020) was published in ICML 2020\n  - Osawa et al. (2019) was published in NeurIPS 2019\n  - Swiatkowski et al. (2020) was published in ICML 2020\n- p. 13, supplement, Fig. 5: error bar regions should be upper and lowered bounded by [0, 1] for accuracy.\n- p. 13, Table 2: Splitting this into two tables, one for MNIST and one for CIFAR-10, would be easier to read.", "extra_scores": {}, "extra_text": {}}, {"review_id": "sjop_3593Hz", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThe paper is generally well written and easy to follow. The idea that BNNs have too many parameters for reliably inferring posterior distributions over all of them is a reasonable one. Splitting the posterior approximation into a deterministic and stochastic component to deal with this issue is interesting.  The experiments do indicate improvements over fully factorized posterior approximations. \n\nConcerns:\n* The process of selecting the sub-network over which the full posterior distribution is inferred is crucial and herein lies my main concern with the approach presented in the paper. Minimizing the Wasserstein distance between the subnetwork posterior p(W_s | data) and the true posterior over all weights p(W | data) is sensible. However, since the true posterior is intractable, the authors instead appear to minimize the distance to an Laplace approximated posterior q(W | data). While this is fine for smaller models when the approximation can use a full covariance,  why does it make sense for larger models where the authors use diagonal approximations to the generalized Gauss Newton matrix (Section 5.3)? If the goal is to do better than such diagonal approximations then why treat the subnetwork with the lowest Wasserstein distance to such crude approximations as optimal? How much worse is the  random selection strategy on the tasks listed in section 5.3? \n* In the toy experiments, when the full covariance posteriors approximations  are available, using the deterministic + stochastic sub-network approximation does a little bit worse than the full covariance approximation but better than the diagonal approximation. If instead the diagonal posterior approximation is used to select the subnetwork in these models, does the resulting approximation still improve upon the diagonal approximation (as seems to be happening in the experiments in 5.3)? \n* The experiment in section 5.3 has other curious issues. Different methods appear to be using different priors (Gaussians with different precisions). How do we then know that the benefits reported stem from the proposed approximation rather than the differences in the model, especially since diagonal Laplace is using a Gaussian prior with a precision that is 80 times smaller?  What priors were used for deep ensembles and SWAG? \n* The grid search by which the prior precisions were selected needs more details. The presented numbers are surprisingly small suggesting that a-priori  with high probability we expect all weights to be zero, and hence the prior predictive functions to be all zero as well. This does not seem like a sensible prior. \n* At the very least there needs to be a discussion about sparse Bayesian deep learning techniques (and preferably an empirical comparison)[1, 2, 3], that use sparsity inducing priors to prune away weights and nodes from a larger network instead of the approach presented here. \n* (Minor) The conclusion of 5.3 that subnetwork posteriors are better calibrated are not supported by Figure 4. I would suggest moving the ECE and Brier score plots from the appendix to Figure 4. \n\nOverall, although I have several concerns they primarily stem from experimental issues in 5.3. Assuming that the authors are able to sufficiently address these in the rebuttal, I am leaning towards a borderline accept. \n\n[1] https://www.jmlr.org/papers/v20/19-236.html\n\n[2] https://papers.nips.cc/paper/7372-posterior-concentration-for-sparse-deep-learning.pdf\n\n[3] https://arxiv.org/pdf/2002.10243.pdf", "extra_scores": {}, "extra_text": {}}, {"review_id": "gErsWdf0h8c", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper presents a new way for approximating posteriors in Bayesian DNN. The network is split into two subnets. One uses only point estimates while another one uses full (non-diagonal) Gaussian approximation. The structure of that subnet is found by taking largest second derivatives of Hessian of linearized DNN (the authors call it generalized Gauss-Newton (GGN) matrix). The authors show that under very specific conditions such choice correposnds to minimization of Wasserstein-2 distance between their approximation and the true posterior. In the experimental part they provide a set of explorative experiments showing that it may be better to use their approximation for inference in large newtork than both using standard (simple) approximations in large network and full Bayesian inference in small network. This is very nice methodologically and I welcome such demonstration but this can only be considered as a (good) proof of concept. The flagship experiment however looks very inconvincing (see below).\n\nPros.\n1. Interesting idea of finding a better approximation for the posterior on subset of parameters.\n2. Methodologically nice PoC.\n3. Thorough comparison with alternative similar techniques such as SWAG.\n\nCons.\n1. The authors claim that they theoretically characterize the descrepancy and derive optimal strategy (see contribution 3) but they (0) consider linearized approximation of DNN (they admit this); (1) do this ONLY for regression problem although their flagship experiment is on classification problems; (2) the method they derive is based on un-natural assumption that covariance matrix is diagonal (it the matrix is diagonal there is no way to approximate with a full submatrix anyway). I would not call it an optimal strategy - it rather looks like a reasonalbe heuristic.\n2. My major concern is section 5.3. The authors claim that their method estimates uncertainty better than all baselines including deep ensembles (DE). I am afraid in its current form the comparison is not fair. They use only DE of size 5 while their method requires approx. (42K)*(42K) = 1756B of parameters that is enough to keep in memory about 160 initial networks. So it would be fair to compare aganist DE of size 160 since we know that larger DE estimate uncertainties better. It is not that surprizing that the suggested method outperforms other baselines since all of them require much less memory. So I would recomend the authors to compare (1) their current model aganist DE that requires similar amount of memory; (2) their reduced model (that requires approximately the same amount of memory as baselines) aganist other baselines to check wether the proposed algorithm may still estimate uncertainty better given the same memory budget.", "extra_scores": {}, "extra_text": {}}, {"review_id": "f-8Qd2qUUCi", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors present a new method for Bayesian deep learning motivated by the difficulty of posterior inference in the \"overparameterized\" regime of deep neural network models. The proposed method provides a principled strategy for selecting a subset of the neural network's parameters (forming a so-called \"subnetwork\") for which a full-covariance approximate posterior can be computed. The authors use the well-studied Laplace approximation with the generalized Gauss-Newton Hessian approximation for the covariance. An empirical analysis is presented which attempts to assess the efficacy of the proposed method in prediction accuracy and uncertainty quantification.\n\nThe presented approach is novel and appears to be a promising contribution to the study of Bayesian neural networks. As a fellow Bayesian, I applaud the authors' efforts. Unfortunately, the paper has a number of significant weaknesses which I detail below. The authors' experimental results appear to me to not sufficiently support some of their claims. There are also a number of formatting issues. As such, I cannot recommend accepting this work in its current state. I would be happy to revisit my rating after revision and discussion.\n\nPros:\n- The basic idea behind the method is very interesting and constructive for the field\n- Theoretical justification is solid\n- Overall well written with only minor clarity issues\n\nCons:\n- Experimental results lack thoroughness\n- Results do not seem to adequately support authors' (somewhat bold) claims\n- Lacking detail in certain areas of the method description\n\n\nI will organize my comments for the authors by section.\n\n#### Section 1\n1. \"In turn, we can apply more expensive, but more faithful, posterior approximations to just that subnetwork to achieve better uncertainty quantification than if we apply cheaper, but more crude, approximations to the full network.\"\nThis sentence is too long and too difficult to read. Please try rephrasing to make it more clear.\n\n#### Section 2\n\n2. It would be nice to see the subnetwork posterior fully defined in terms of the fixed weights, e.g:\n$$p(W_s | y,X,W^*) \\propto p(y|X,W_s,W^*)p(W_s)p(W^*) = p(W_s|y,X)p(W^*)$$\nI don't meant to be pendantic here. It wasn't obvious to me (at first) that the delta functions were, in fact, a degenerate prior over the fixed weights, $W^*$; this would make it more clear.\n\n3. It might be worth discussing here or later what implications the degenerate prior has for the subnetwork posterior. This seems to, at least, preclude any ideas about possibly applying gradient-optimized MAP directly to the subnetwork posterior.\nFurthermore, what are the benefits of this approach as opposed to assigning a non-degenerate Gaussian prior with fixed variance to the remaining weights?\n\n#### Section 4\n\n4. \"We now analyze the following procedure...\"\nAre you analyzing the procedure for selecting a subnetwork? Or the entire procedure described previously? This is not made clear.\n\n5. In step 1, you specify the analytical solution for w_MAP, but in section 3 you describe MAP as being performed using stochastic gradient optimization. Which one are you using? Related to previous point, is this the same step as outlined in section 3? Or a separate step?\n\n6. In general, this section needs work. It is not clear to me where you actually detail how the subnetwork is selected. You discuss approximating the optimal subnetwork w.r.t to the Wasserstein distance in equation 11, but this equation requires M_S to be already available. M_S is generated by \"a (one-shot) procedure of choice\"; I expected this choice to be clearly explained in this section, but this is not the case.\n\n#### Section 5.1\n\n7. It's a bit confusing that you say \"50%, 97%, and 99% of model parameters\" when really you mean that this percentage of weights were pruned. It would be visually more intuitive as well in Figure 2 if you named these 50%, 3%, 1%, since the posterior size is getting *smaller* with each one.\n\n8. Just to clarify, homeoscedastic here is w.r.t to the sequence of data points? i.e. you have one variance rather than one per data point?\n\n9. Please provide some explanation for the prior precision being set to 3.\n\n10. Bolding most of the last sentence in section 5.1 is unnecessary and looks weird. I would suggest bolding individual words or phrases in the sentence, or just not bolding at all.\n\n#### Section 5.2\n\n11. 1e4 ->$10^4$\n\n12. It would be helpful to provide a very brief summary of the three datasets as well as what the \"gap variants\" are (it's fairly straightforward, just creating \"gaps\" in the training data). I cannot find the \"kin8nm\" dataset on UCI, so at bare minimum this needs to be clarified.\n\n13. Figure 3 is quite difficult to interpret at first glance. You should at least use a discretized perceptually uniform colormap to make more visually apparent the pattern in increasing network size. I would also consider using a line plot rather than scatter plot here, but this is a matter of taste.\n\n14. Why are you using the log likelihood as opposed to the full posterior probability? Since we're in a Bayesian setting here, it seems worth considering the priors.\n\n15. While log likelihood can, in principle, serve as an indirect proxy for uncertainty calibration, there are important caveats to consider which make it unconvincing as a standalone measure (see \"Pitfalls of In-Domain Uncertainty Estimation...\", Ashukha et al. 2020). Furthermore, a quick scan of recent literature (cited in this work) confirms that most authors tend to use multiple methods of assessment (e.g. Brier score, accuracy calibration curves, etc), not just log likelihood. The raw likelihood scores also don't provide any interpretability in *how much better* one model is than another, i.e. what concrete impact does the difference have on uncertainty quantification? Thus, **it's unconvincing that here (and elsewhere) you rely entirely on log ilkelihood** to demonstrate the alleged superiority of your method.\n\n16. It's noteworthy that 2/3 of the \"gap\" datasets produce inconsistent results, particularly considering that they are designed specifically to test out-of-sample uncertainty. This needs to be discussed.\n\n17. typo: modelling -> modeling\n\n18. In light of points 15 and 16, I do not think you have sufficient empirical basis to make the claim that \"...it is better to perform subnetwork inference in a large model than full network inference in a small one\". The results are simply too weak to support this. I suggest either running a more comprehensive experiment or dialing back the confidence of this claim.\n\n#### Section 5.3\n\n19. You mention \"grid search\" multiple times but do not provide any indication of the grid you searched over.\n\n20. typo: sitribution -> distribution\n\n21. Please specify the error metrics being used in figure 4 (and discussed in the text).\n\n22. Once again, marginally higher likelihood scores are not, in my view, sufficient basis to claim that your method assigns \"high uncertainty to out of distribution points\" or that it is better calibrated. This is especially the case considering that there seems to be little to no improvement in the error. The case for improved \"robustness\" also appears to be equally unfounded on these grounds.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper propose an approach to efficient Bayesian deep learning by applying Laplace approximations to sub-structures within a larger network architecture. In terms of strengths, scalable approximate Bayesian inference methods for deep learning models are an important and timely topic. The paper includes an extensive set of experiments with promising results. 

In terms of issues, the reviewers originally raised many concerns and the authors provided a large update to the paper. However, following that update and the discussion, several concerns remain. First, the reviewers noted that the originally submitted draft made claims about the optimality of the sub-network selection procedure that were incorrect due to the use of a diagonal approximation. The authors subsequently retracted these claims and re-focused on the idea that the subset selection approach is theoretically well-motivated heuristic that performs well empirically. Following the discussion, the reviewers continued to express concerns about the heuristic nature of this procedure. 

A second point has to do with scalability. The reviewers noted that the authors had only evaluated their approach on small data sets, leaving open the question of how scalable the method is. The authors responded by adding experiments on the same data sets using larger models, which does not squarely address the issue raised. Third, an additional point was raised regarding the lack of control of resource use in the experiments. The authors note that their approach can use more resources when available while many other methods can not. However, some methods including deep ensembles can also expand to use more resources, as can posterior ensembles produced using MCMC methods like SGLD and SGHMC. The authors need to consider quantifying space-performance and time-performance trade-offs in the same units for different approaches to satisfactorily address this issue. While the authors added one set of experiments looking at deep ensembles in isolation, their conclusions that performance saturates for these models at low ensembles sizes seems to be hasty in some cases (e.g., deep ensembles show continued improvement for large corruptions in Figure 5(right) despite the claim by the authors that the models saturate after 15 epochs). 

In summary, this appears to be a promising approach. While the authors made significant efforts to correct issues and address questions with the original draft, the majority view of the reviewers following discussion is that this paper requires additional work to more carefully expand on the revised results and to address the heuristic status of the sub-network selection approach.
Decision: Reject