Dangers of Bayesian Model Averaging under Covariate Shift
Authors: Pavel Izmailov
Patrick Nicholson
Sanae Lotfi
Andrew Gordon Wilson
Venue: NeurIPS 2021 Poster (2021)
Keywords: Bayesian
Bayesian neural networks
neural networks
covariate shift
out-of-distribution generalization

Abstract: Approximate Bayesian inference for neural networks is considered a robust alternative to standard training, often providing good performance on out-of-distribution data. However, Bayesian neural networks (BNNs) with high-fidelity approximate inference via full-batch Hamiltonian Monte Carlo achieve poor generalization under covariate shift, even underperforming classical estimation. We explain this surprising result, showing how a Bayesian model average can in fact be problematic under covariate shift, particularly in cases where linear dependencies in the input features cause a lack of posterior contraction. We additionally show why the same issue does not affect many approximate inference procedures, or classical maximum a-posteriori (MAP) training. Finally, we propose novel priors that improve the robustness of BNNs to many sources of covariate shift.

Dangers of Bayesian Model Averaging under Covariate Shift
Pavel Izmailov NYU
Patrick Nicholson Covera Health
Sanae Lotfi NYU
Andrew Gordon Wilson NYU
Abstract
Approximate Bayesian inference for neural networks is considered a robust alterna- tive to standard training, often providing good performance out-of-distribution data. However, Bayesian neural networks (BNNs) with high-fidelity approximate inference via full-batch Hamiltonian Monte Carlo achieve poor generalization under covariate shift, even underperforming classical estimation. We explain this surprising result, showing how Bayesian model average in fact be problematic under covariate shift, particularly in cases where linear dependencies in the input features cause lack of posterior contraction. We additionally show why the same issue does not affect many approximate inference procedures, or classical maxi- mum a-posteriori (MAP) training. Finally, we propose novel priors that improve the robustness BNNs to many sources covariate shift.
1 Introduction
The predictive distributions of deep neural networks are often deployed in critical applications such as medical diagnosis [Gulshan et al., 2016, Esteva et al., 2017, Filos et al., 2019], and autonomous driving (Bojarski et al., 2016, Al-Shedivat et al., 2017, Michelmore al., 2020]. These applications typically involve covariate shift, where the target data distribution is different from the distribution used for training [Hendrycks and Dietterich, 2019, Arjovsky, 2021]. Accurately reflecting uncertainty is crucial for robustness to these shifts [Ovadia et al., 2019, Roy et al., 2021]. Since Bayesian methods provide a principled approach to representing model (epistemic) uncertainty, they are commonly benchmarked on out-of-distrrn out-o--distribution (OOD) generalization tasks [e.g., Kendall and Gal, 2017, Ovadia et al., 2019, Chang et al., 2019, Dusenberry et al., 2020, Wilson and Izmailov, 2020].
However, Izmailov et al. [2021] recently showed that Bayesian neural networks (BNNs) with high fidelity inference through Hamiltonian Monte Carlo (HMC) provide shockingly poor OOD generalization performance, despite the popularity and success of approximate Bayesian inference in this setting [Gal and Ghahramani, 2016, Lakshmmnnaraaannn et al., 2017, Ovadia al., 2019, Maddox et al., 2019, Wilson and Izmailov, 2020, Dusenberry et al., 2020, Benton et al., 2021].
In this paper, seek to understand, further demonstrate, and help remedy this concerning behaviour. We show that Bayesian neural networks perform poorly for different types of covariate shift, namely test data corruption, domain shift, and spurious correlations. In Figure (a)) we see that a ResNet-20 BNN approximated with HMC underperforms a maximum a-posteriori (MAP) solution by 25% on the pixelate-corrupeed CIFAR-10 10 test set. This result particularly surprising given that on the in-distribution test data, the BNN outperforms the MAP solution by over 5%.
Intuitively, we find that Bayesian model averaging (BMA) be problematic under covariate shift as follows. Due to linear dependencies in the features (inputs) of the training data distribution, model parameters corresponding to these dependencies do affect the predictions on the training data. As an illustrative special case of this general setting, consider MNIST digits, which always have black corner pixels (dead pixels, with intensity zero). The corresponding first layer weights always multiplied by zero and have no effect on the likelihood, Consequently, these weights are simply
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

09
5%
Accuracy
07
25%
MAP
BNN
No skot: Corruption Intensity
-025 0.00
0.55
0.05
0.00
0.05
(a) ResNet-20, CIFAR-IO-C
(b) BNN weights
(c) MAP weights
Figure 1: Bayesian neural networks under covariate shift. (a): Performance of ResNet-20 the pixelate corruption in CIFAR-IO-C. For the highest degree corruption, a Bayesian model average underperforms MAP solution by 25% (44% against 69%) aceuracy. See Izmailov al. (b): Visualization of the weights in the first layer of Bayesian fully-connected network MNIST sampled via HMC. (c): The corresponding MAP weights. We visualize the weights connecting the input pixels to a neuron in the hidden layer as 28 x 28 image, where each weight is shown in the location of the input pixel it interacts with.
[2021] for details.
sampled from the prior. at test time the corner pixels are not black, e-g., due to corruption, these pixel values will be multiplied by random weights sampled from the prior, and propagated to the next layer, significantly degrading performance. On the other hand, classical MAP training drives the unrestricted parameters towards zero due to regularization from the prior that penalizes the parameter norm, and will not be similarly affected by noise at test time. Here we see major difference in robustness between optimizing posterior for MAP training in comparison to posterior weighted model average.
As motivating example, in Figure (b,, c) we visualize the weights in the first layer of a fully- connected network for a sample from the BNN posterior and the MAP solution on the MNIST dataset. The MAP solution weights are highly structured, while the BNN sample appears extremely noisy, similar a draw from the Gaussian prior. In particular the weights corresponding to dead pixels (i.e. pixel positions that are black for all the MNIST images) near the boundary of the input image set near zero (shown white) by the MAP solution, but sampled randomly by the BNN, If at test time the data is corrupted, e.g. by Gaussian noise, and the pixels near the boundary of the image are activated, the MAP solution will ignore these pixels, while the predictions of the BNN will be significantly affected.
Dead pixels are a special case of our more general findings: we show that the dramatic lack of robustness for Bayesian neural networks is fundamentally caused by any linear dependencies in the data, combined with models that are non-linear in their parameters. Indeed, we consider a wide range of covariate shifts, including domain shifts. These robustness issues have the potential to impact virtually every real-world application of Bayesian neural networks, since train and test rarely come from exactly the same distribution.
Based on our understanding, we introduce novel prior that assigns a low variance the weights in the first layer corresponding to directions orthogonal to the data manifold, leading to improved generalization under covariate shift. We additionally study the effect of non-zere mean corruptions and accordingly propose a second prior that constrains the sum of the weights, resulting in further improvements in OOD generalization for Bayesian neural networks.
Our code is available here.
2 Background
Bayesian neural networks. A Bayesian neural network model is specified by the prior distribution p(w) over the weights w of the model, and the likelihood function p(y|æ, w), where z represents the input features and y represents the target value. Following Bayes' rule, the posterior distribution over the parameters w after observing the dataset D {(x,, y:)li 1,..., n) is given by
P(D[w) p(w)
I pmu (ri, w) p(w)
P(w Jow p( (D|w') w') p(w" ddw'
J.n 1I" p( x. w() pl u/ 'dw!

---

where we assume that the likelihood is independent over the data points. The posterior in Equation 1 then used make predictions for a new input x according to the Bayesian model average:
p(y]=) = p(yla, w) p(w)
Unfortunately, computing the BMA in Equation is intractable for Bayesian neural networks. Hence, a number of approximate inference methods have been developed. In this work, we use full-batch HMC [Neal et al., 2011, I, Izmailov al., 2021], J. as it provides high-accuracy posterior approximation. For more detailed discussion of Bayesian deep learning see Wilson and Izmailov [2020].
Maximum a-posteriori (MAP) estimation. In contrast with Bayesian model averaging. MAP estimator uses the single setting of weights (hypothesis) that maximizes the posterior density WMAP argmax (egmaxp(w[D) p(w(D) argmax(log p p(D/w) + p(w)), where the log prior can be viewed
as a regularizer. For example, if we use a Gaussian prior on w, then l(g)(w) will penalize the & norm of the parameters, driving parameters that do not improve the log likelihood log p(Dlw) zero. MAP is the standard approach to training neural networks and our baseline for classical training throughout the paper. We perform MAP estimation with SGD unless otherwise specified.
Covariate shift. In this paper, we focus on the covariate shift setting. We assume the training
dataset Darin consists of i.i.d. samples from the distribution Purain ( v) = Praii (x:) p(y]zr). However, the test data may come from a different distribution Pres (v. 1) = Pteaa (a) pr (z). For concreteness, assume the conditional distribution p(y]=:) remains unchanged, but the marginal distribution of the input features (r) differs from Prrain (z); we note that our results do not depend on this particular definition of covariate shift. Arjovsky [2021] provides detailed discussion of covariate shift.
3 Related work
Methods to improve robustness to shift between train and test often explicitly make use of the test distribution in some fashion. For example, it is common to apply semi-supervised methods to the labelled training data augmented by the unlabelled test inputs [c.g.. Daume III and Marcu, 2006, Athiwaratkun al., 2018], or to learn shared feature transformation for both train and test [c.g.. Daumé III, 2009]. In a Bayesian setting, Storkey and Sugiyama [2007] and Storkey [2009] propose such approaches for linear regression and Gaussian processes under covariate shift, assuming the data comes from multiple sources. Moreover, Shimodaira [2000] propose to re-weight the train data points according to their density in the test data distribution. Sugiyama et al. [2006, 2007] adapt this importance-weighhing approach to the cross cross- validation setting.
We focus on the setting of robustness to covariate shift without any access to the test distribution (e.g., Daume III and Marcu, 2006). Bayesian methods are frequently applied this setting, often motivated by the ability for a Bayesian model average to provide a principled representation of epistemic uncertainty: there are typically many consistent explanations for out of distribution points, leading to high uncertainty for these points. Indeed, approximate inference approaches for Bayesian neural networks are showing good and increasingly better performance under covariate shift [e.g.. Gal and Ghahramani, 2016, Lakshminarayanan al., 2017, Ovadia et al., 2019, Maddox et al., 2019, Wilson and Izmailov, 2020, Dusenberry et al., 2020, Benton et al., 2021].
Many works attempt to understand robustness to covariate shift. For example, for classical training Neyshabur et al. [2020] show that models relying on features that encode semantic structure in the data are more robust to covariate shift. Nagarajan et al. [2020] also provide insights into how classically trained max-margin classifiers can fail under covariate shift due to a reliance on spurious correlations between class labels and input features. BNN robustness to adversarial attacks is related area of study, but generally involves much smaller perturbations to the test covariates than covariate shift. Carbone et al. [2020] prove that BNNs are robust to gradient-based adversarial attacks in the large data, overparameterized limit, while Wicker et al. [2021] present a framework for training BNNs with guaranteed robustness adversarial examples.
In this work, we propose novel priors that improve robustness of BNNs under covariate shift. In Fortuin et al. [2021]. the authors explore wide range of priors and, in particular, show that the distribution of the weights of SGD-trained networks is heavy-tailed such as Laplace and Student-l, and does not appear Gaussian. Other heavy-tailed sparsity inducing priors have also been proposed
3

---

in the literature [Carvalho et al., 2009, Molchanov et al., 2017, Kessler et al., 2019, Cui et al., 2020, Izmailov al., 2021, Fortuin, 2021]. Inspired by this work, we evaluate Laplace and Student-f priors for BNNs, but find that they do not address the poor performance of BNNs under covariate shift.
Domingos [2000], Minka [2000], Masegosa [2019] and Morningstar al. [2020] explore failure modes of Bayesian model averaging when the Bayesian model does not contain reasonable solution in its hypothesis space, causing issues when the posterior contracts. This situation is orthogonal to the setting in our paper, where we know the Bayesian model does contain a reasonable solution in its hypothesis space, since the MAP estimate is robust to covariate shift. In our setting, robustness issues are caused by lack of posterior contraction.
In general, understanding and addressing covariate shift is large area of study. For a comprehensive overview, see Arjovsky [2021]. To our knowledge, no prior work has attempted to understand, further demonstrate, or remedy the poor robustness of Bayesian neural networks with high fidelity approximate inference recently discovered in Izmailov et al. [2021].
4 Bayesian neural networks are not robust to covariate shift
In this section, we evaluate Bayesian neural networks under different types of covariate shift. Specifi- cally, we focus on two types of covariate shift: test data corruption and domain shift. In Appendix D, we additionally evaluate BNNs in the presence of spurious correlations in the data.
Methods. We evaluate BNNs against two deterministic baselines: a MAP solution approximated with stochastic gradient descent (SGD) with momentum [Robbins and Monro, 1951, Polyak. 1964] and a deep ensemble of 10 independently trained MAP solutions [Lakshminarayanan et al., 2017]. For BNNs, we provide the results using Gaussian prior and a more heavy-tailed Laplace prior following Fortuin et al. [2021]. Izmailov et al. [2021] conjectured that cold posteriors [Wenzel et al., 2020] can improve the robustness of BNNs under covariate shift; to test this hypothesis. we provide results for BNNs with a Gaussian prior and cold posteriors at temperature 10-2, For all BNN models, we run single chain of HMC for 100 iterations discarding the first 10 iterations as burn-in, following Izmailov et al. [2021]. We provide additional experimental details Appendix A.
Datasets and data augmentation. We run all methods on the MNIST [LeCun et al., 2010] and CIFAR-10 10 [Krizhevsky et al., 2014] datasets. Following Izmailov al. [2021] we do not use data augmentation with any of the methods, as it is not trivially compatible with the Bayesian neural network framework [e.g., Izmailov et al., 2021, Wenzel et al., 2020].
Neural network architectures. On both the CIFAR-10 and MNIST datasets we use a small convolutional network (CNN) inspired by LeNet-5 [LeCun et al., 1998], with 2 convolutional layers followed by 3 fully-connected layers. On MNIST we additionally consider a fully-connected neural network (MLP) with 2 hidden layers of 256 neurons each. We note that high-fidelity posterior sampling with HMC is extremely computationally intensive. Even on the small architectures that we consider, the experiments take multiple hours on NVIDIA Tesla V-100 GPUs or 8-core TPU-V3 devices [Jouppi et al.. 2020]. See Izmailov et al. [2021] for details on the computational requirements of full-batch HMC for BNNs.
4.1 Test data corruption
We start by considering the scenario where the test data is corrupted by some type of noise. In this case, there is no semantic distribution shift: the test data is collected in the same way as the train data, but then corrupted by generic transformation.
In Figure 2, we report the performance of the methods trained on MNIST and evaluated on the MNIST-corrupted (MNIST-C) test sets [Mu and Gilmer, 2019] under various corruptions'. We report the results for a fully-come conneecd network and a convolutional network.
With the CNN architecture, deep ensembles consistently outperform BNNs. Moreover, even single MAP solution significantly outperforms BNNs on many of the corruptions. The results are especially striking on brightness and fog corruptions, where the BNN with Laplace prior shows accuracy at the
In addition to the corruptions from MNIST-C, we consider Gaussian noise with standard deviation 3, as this corruption considered by Izmailov al. [2021].

---

MLP MLP Acouaty
Domais
Nose
Blu
Afline
Other Corruct ons
Aug
CNN Accusty
2 2 a a
HOUST
Stow
Gavs Nosee
Shodt Nooee
Inpedd Rosse
Coos Bisr
Modban Bla
Scat
Trenttee
Brightouse
Figt
Singer
Spedder
Tigent
Gendy Edges
Avonge
Dantet
Deep Ens
Laplace
BNN Temp-10" 2
Figure 2: Robustness on MNIST. Accuracy for deep ensembles, MAP and Bayesian neural networks trained on MNIST under covariate shift. Top: Fully-connected network; bottom: Convolutional neural network. While on the original MNIST test set BNNs provide competitive performance, they underperform deep ensembles most of the corruptions. With the CNN architecture, all BNN variants lose to MAP when evaluated on SVHN by almost 20%.
level of random guessing, while the deep ensemble retains accuracy above 90%. Gaussian noise and impulse noise cоmuptiins also present significant challenges for BNNs. The BNNs show the most competitive performance in-distribution and on the corruptions representing affine transformations: shear, scale, rotate and translate. The results are generally analogous for the MLP: across the board the BNNs underperform deep ensembles, and often even a single MAP solution.
While the cold posteriors provide an improvement some of the noise corruptions, they also hurt performance on Brightness and Stripe, and do not improve the performance significantly on average across all corruptions compared to standard BNN with Gaussian prior. We provide additional results on cold posterior performance in Appendix G.
Next, we consider the CIFAR- SEPAR-10-corrun 10-corrppedd dataset (CIFAR-10-C) [Hendrycks and Dietterich, 2019]. CIFAR-10-C consists of 18 transformations that are available at different levels intensity (1-5). We report the results using corruption intensity 4 (results for other intensities are in Appendix C) for each of the transformations in Figure Similarly to MNIST-C. BNNs outperform deep ensembles and MAP on in-distribution data, but underperform each over multiple corruptions. On CIFAR-10-C. BNNs are especially vulnerable to different types of noise (Gaussian noise, shot noise, impulse noise, speckle noise). For each of the noise corruptions, BNNs underperform even the classical MAP solution. The cold posteriors improve the performance on the noise corruptions, but only provide a marginal improvement across the board.
We note that Izmailov al. [2021] evaluated a ResNet-20 model on the same set of CIFAR-10-C corruptions. While they use much larger architecture, the qualitative results for both architectures are similar: BNNs are the most vulnerable noise and blur corruptions. We thus expect that our paper's s analysis is not specific to smaller architectures and will equally apply to deeper models.
4.2 Domain shift
Next, consider a different type of covariate shift where the data and train data come from different, but semantically related distributions.
First, we apply our CNN and MLP MNIST models to the SVHN test set [Netzer et al., 2011]. The WOIST-to-SVHH domain shift task is a common benchmark for unsupervised domain adaptation: both datasets contain images of digits, although MNIST contains hand-written digits while SVHN represents house numbers. In order to apply our MNIST models to SVHN we crop the SVHN images and convert them to grayscale. We report the results in Figure 2. While for MLPs all methods perform similarly, with the CNN architecture BNNs underperform deep ensembles and MAP by nearly 20%.
Next, we apply our CIFAR- 10 CNN model to the 10L--10 dataset [Coates et al., 201 1]. Both datasets contain natural images with 9 shared classes between the two datasets². We report the accuracy of
CIIARR-I 10 has class frog and STL-10 10 has monkey. The other nine classes coincide.
5

---

Image Processing
Computtions
Avg
0.6
Accuney
0.6
0.2
CH CPAR-10
STL STL-d0
Statt Nolee
Impolie Noste
Speckle Molie
Bi
Oedoraa Blue
Gros Bhe
Rie
Brightsess
Contrac
Elaste
JPEG
Proelat
Fient
Syotte
Avergge
Goun
Doep Ers
MAP
BNN Ganasim Prior
BNN Laplace Price
ENN Temp-10
Figure Robustness on CIFAR-10. Accuracy for deep ensembles, MAP and Bayesian neural networks using CNN architecture trained on CIFAR- 10 under covariate shift. For the corruptions from CIFAR-10-C 10-C, we report results for corruption intensity While the BNNs with both Laplace and Gaussian priors outperform deep ensembles on the in-distribution accuracy, they underperform even a single MAP solution on most corruptions.
the CIFAR-10 models on these 9 shared classes in STI-10 in Figure 3. While BNNs outperform the MAP solution, they still significantly underperform deep ensembles.
The results presented in this section highlight the generality and practical importance of the lack of robustness BNNs: despite showing strong performance in-distribution, BNNs underperform even single MAP solution (classical training) over an extensive variety of covariate shifts.
Understanding Bayesian neural networks under covariate shift
Now that we have established that Bayesian neural networks are highly susceptible to many types of covariate shift, we seek to understand why this is the case. In this section, we identify the linear dependencies the input features as one of the key issues undermining the robustness BNNs. We emphasize that linear dependencies in particular are not simply chosen for the simplicity of analysis, and their key role follows from the structure the fully-connected and convolutional layers.
5.1 Motivating example: dead pixels and fully-connected layers
To provide an intuition for the results presented in this section, we start with simple but practically relevant motivating example. Suppose we use fully-connected Bayesian neural network on a dataset D = ( y:) W with m input features, where I; E IR.. We use the upper indices to denote the features _ ...,, x and the lower indices to denote the datapoints r;, Then, for each neuron j the first hidden layer of the network, the activation can be written 2j < <bEn xdw,y + bJ). where w/y is the weight of the first layer of the network corresponding to the input feature and hidden neuron j. and b) the corresponding bias. We show the following Lemma:
Lemma Using the notation introduced above, suppose that the input feature a' is equal to zero for all the examples F3 in the training dataset D. Suppose the prior distribution over the parameters p(W) factorizes p(W)= p(w?z) p(W \ why)) for some neuron j in the first layer, where W why represents all the parameters W of the network except w/y- Then, the posterior distribution p(W] ()) will also factorize and the marginal posterior over the parameter wiy will coincide with the prior:
P(WID) = p(W \ w/y D) -((w))).
(3)
Consequently, the MAP solution will set the weight why the value with maximum prior density.
Intuitively, Lemma I says that if the prior for some parameter in the network is independent of the other parameters, and the value of the parameter does not affect the predictions of the model on any of the training data, then the posterior of this parameter will coincide with its prior. In particular, if one of the input features is always zero, then the corresponding weights will always be multiplied by zero, and will not affect the predictions of the network. To prove Lemma we simply note that the posterior proportional to the product of prior and likelihood, and both terms factorize with respect to w/y- We present a formal proof in Appendix H.

---

So, for any sample from the posterior, the weight wi will be a random draw from the prior distribution. Now suppose test time we evaluate the model on data where the feature x² is no longer zero. Then for these new inputs, the model will be effectively multiplying the input feature x' by a random weight way leading to instability in predictions. In Appendix H, formally state and prove the following proposition:
Proposition 1 (informal) Suppose the assumptions of Lemma 1 hold. Assume also that the prior distribution p( (/is)) has maximum density at O and that the network uses ReLU activations. Then for any test input ž. the expected prediction under Bayesian model averaging (Equation 2) will depend on the value of the feature H', while the MAP solution will ignore this feature.
For example, in the MNIST dataset there is a large number of dead pixels: pixels near the boundaries of the image that have intensity zero for all the inputs. In practice, we often use independent zero- mean priors (c.g. Gaussian) for each parameter of the network. So, according to Lemma 1, the posterior over all the weights in the first layer of the network corresponding to the dead pixels will coincide with the prior. If at test time the data is corrupted by e.g., Gaussian noise, the dead pixels will receive non-zero intensities, leading to significant degradation in the performance of the Bayesian model average compared to the MAP solution.
While the situation where input features are constant and equal to zero may be rare and easily addressed, the results presented in this section can be generalized to any linear dependence in the data. We will now present our results in the most general form.
5.2 General linear dependencies and fully-connected layers
We now present our general results for fully-connected Bayesian neural networks when the features are linearly dependent. Intuitively, if there exists a direction in the input space such that all of the training data points have a constant projection on this direction (i.e. the data lies in a hyper-plane), then posterior coincides with the prior in this direction. Hence, the BMA predictions are highly susceptible to perturbations that move the test inputs in a direction orthogonal the hyper-plane srper-plane. The MAP solution on the other hand is completely robust to such perturbations. In Appendix H we prove the following proposition.
Proposition 2 Suppose that the prior over the weights why and biases bi} in the first layer is an i..d Gaussian distribution N a2), independent of the other parameters in the model. Suppose all the inputs x1 ... Zn in the training dataset D lie in an affine subspace of the input space: Σi= 2/ co for all i = =,...., . n and some constants e such that Σ ct = 1. Then,
(1) For any neuron in the first hidden layer, the posterior distribution of random variable wj. Eis cawlyy cob) (the projection of parameter vector (wit.--.,,uu why, bz) on direction
(C1, .i...... Cm. ..,, Co)) will coincide with the prior (0,
(2) The MAP solution will set wy to zero.
(3) (Informal) Assuming the network uses ReLU activations, at test time, the BMA prediction will be susceptible to the inputs Tt that lie outside of the subspace, i.e. the predictive mean will depend on >ー F'c - Cg. The MAP prediction will not depend on this difference.
Empirical support. To test Proposition 2, examine the performance of fully-connected BNN on MNIST. The MNIST training dataset is not full rank, meaning that has linearly dependent features. For fully-connected BNN, Proposition 2 predicts that the posterior distribution of the first layer weights projected onto directions corresponding to these linearly dependent features will coincide with the prior. In Figure 4(a) we test this hypothesis by projecting first layer weights onto the principal components of the data. As expected, the distribution of the projections on low- low-variance PCA components (directions that are constant or nearly constant in the data) almost exactly coincides with the prior. The MAP solution, on the other hand, sets the weights along these PCA components close to zero, confirming conclusion (2) of the proposition. Finally, in Figure 4(b) we visualize the performance of the BMA and MAP solution as we apply noise along high-variance and low-variance directions the data. As predicted by conclusion (3) of Proposition 2, the MAP solution is very robust noise along the low-variane directions, while BMA is not.
7

---

MIl Prier A10 100
CNR, Prior A7o 100
ZNN Prier A00.:A00..0 A10.30
Top 50 PC
Last 50 PC
Veights Projectoon
1.0
3.0
Accuncy
0.5
16
PCA componint
PCA component
PCA compenent
Noise
Scale
BNN
MAP
(a) Weight projections
(b) Robustness along PCs
Figure 4: Bayesian inference samples weights along low-variance principal components from the prior, while MAP sets these weights zero. (a): The distribution (mean 2 std) of projections of the weights of the first layer the directions corresponding to the PCA components of the data for BNN samples and MAP solution using MLP and CNN architectures with different prior scales. In each case, MAP sets the weights along low-variance components to zero, while BNN samples them from the prior. (b): Accuracy of BNN and MAP solutions on the MNIST test set with Gaussian noise applied along the 50 highest and 50 lowest variance PCA components of the train data (left and right respectively). MAP is very robust to noise along law-nariance PCA directions, while BMA is not; the two methods are similarly robust along the highest-variants PCA components.
5.3 Linear dependencies and convolutional layers
Finally, we extend Proposition 2 to convolutional layers. Unlike fully-connected layers, convolu- tional layers include weight sharing such that no individual weight corresponds to a specific pixel in the input images. For example, dead pixels will not necessarily present an issue for convolutional layers, unlike what is described in subsection 5.1. However, convolutional layers are still susceptible to special type of linear dependence. Intuitively, we can think of the outputs of the first convolutional layer on all the input images as the outputs of cully-connected layer applied to all the K x K patches of the input image. Therefore the reasoning in Proposition 2 applies the convolutional layers, with the difference that the linear dependencies in the K x K patches cause the instability rather than dependencies in the full feature space. Appendix H we prove the following proposition.
Proposition 3 Suppose that the prior over the parameters of the convolutional filters and biases in the first layer is an i.i.d. Gaussian distribution N NO, a²), independent of the other parameters in the model. Suppose that the convolutional filters in the first layer are of size K x K x C, where C is the number of input channels. Then, consider the set DD of size of all the patches of size K x K x C extracted from the training images in D after applying the same padding as in the first convolutional layer. Suppose all the patches 21 2N.. in the dataset D lie in an affine subspace of a.h.
the space R. K RKXKXt xKxC XC:
"c,,.b "t for all i = 1,...,N N and some constants :EC -E=1 -K E 24
e, such that EC Ed EK Proposition 2 (see the Appendix for the details).
+ 1. Then, we can prove results analogous to (1)-(3 in
Empirical support. In Figure 4(a), we visualize the projections of the weights in the first layer of the CNN architecture on the PCA components of the x K patches extracted from MNIST. Analogously to fully-connected networks, the projections of the MAP weights close to zero for low-variance components, while the projections of the BNN samples follow the prior.
5.4 What corruptions will hurt performance?
Based on Propositions 2. 3, expect that the corruptions that are the most likely to break linear dependence structure in the data will hurt the performance the most. In Appendix we argue that noise corruptions are likely to break lincar dependence, while the affine corruptions are more likely to preserve it, agreeing with our observations in section 4.
5.5 Why do some approximate Bayesian inference methods work well under covariate shift?
Unlike BNNs with HMC inference, some approximate inference methods such as SWAG [Maddox et al., 2019], MC dropout [Gal and Ghahramani, 2016), deep ensembles [laaksminarryyaann et al., 2017] and mean field VI [Blundell et al., 2015] provide strong performance under covariate shift
8

---

[Ovadia et al., 2019, Izmailov et al., 2021]. For deep ensembles, we can easily understand why: deep ensemble represents an average of approximate MAP solutions, and we have seen in conclusion (3) of Proposition 2 that MAP is robust to covariate shift in the scenario introduced in subsection 5.2. Similarly, other methods are closely connected to MAP via characterizing the posterior using MAP optimization iterates (SWAG), or modifying the training procedure for the MAP solution (MC Dropout), We provide further details, including theoretical and empirical analysis of variational inference under covariate shift, in Appendix J.
Towards more robust Bayesian model averaging under covariate shift
In this section, we propose a simple new prior inspired by our theoretical analysis. In section 5 we showed that linear dependencies in the input features cause the posterior to coincide with the prior along the corresponding directions in the parameter space. In order to address this issue, we explicitly design the prior for the first layer of the network so that the variance is low along these directions.
6.1 Data empirical covariance prior
Let us consider the empirical covariance matrix of the inputs T₁. Assuming the input features are all preprocessed to be pero-mean Σ;= I; = 0, we have Σ = - Ei=1 xạæe For fully-connected networks, we propose use the EmpCov prior p(w") = NN aL + cf) on the weights w of the first layer of the network, where is small positive constant ensuring that the covariance matrix is positive definite. The parameter
determines the scale of the prior.
Suppose there a linear dependence in the input features of the data: x p = e for some direction p and constant Then p will be an eigenvector of the empirical covariance matrix with the correspond- ing cigenvalue equal to 0: Σp = n-1 Σi=1 x4x; n- Ei= 24 = O. Hence the prior over w will have variance of c along the direction p.
More generally, the EmpCov prior is aligned with the principal components of the data, which are the eigenvectors of the matrix Σ. The prior variance along each principal component p; is equal to ao² + e where of is its corresponding explained variance. In Appendix K, we discuss more general family of priors aligned with the principal components of the data.
Generalization to convolutions. We can generalize the EmpCov prior to convolutions by replacing the empirical covariance of the data with the empirical covariance of the patches that interact with the convolutional filter, denoted by Á in Proposition 3.
Is EmpCov valid prior? EmpCov constructs valid prior by evaluating the empirical covariance matrix of the inputs. This prior does not depend on the train data labels 34. unlike the approach known as Empirical Bayes [see e.g. Bishop, 2006, section 3.5]. which commonly used to specify hyperparameters in Gaussian process and neural network priors [Rasmussen and Williams, 2006, MacKay, 1995].
6.2 Experiments
In Figure 5 we report the performance of the BNNs using the EmpCov prior. In each case, we apply the EmpCov prior to the first layer, and a Gaussian prior to all other layers. For more details, please see Appendix A. On both MLP on MNIST and CNN on CIFAR-10, the EmpCov prior significantly improves performance across the board. In both cases, the BNN with EmpCov prior shows competitive performance with deep ensembles, especially in the domain shift experiments. EmpCov is also particularly useful on the noise corruptions.
In Appendix L, we provide a detailed analysis of the performance of the CNN architecture on MNIST. Surprisingly, we found that using the EmpCov prior by itself does not provide a large improvement in this case. In Appendix L,, identify issue specific to this particular setting, and propose another targeted prior that substantially improves performance.
7 Discussion
We consider the generality of our results, additional perspectives, and future directions.

---

MNIST. MLP
CIFAR-10. CNN
1.0
Accuracy
0.6
MNSS
SWeN
Couss Noine
Impollo No
Motess Blter
Stata
Transtare
Biginiees
Foo
Stripe
Conny Edgns
Averag
CIPAR 10
STL-10
Couss Noine
Impollo Noie
Gonusa Biler
Defocas Bilor
Brightoous
Contrnt
Pioslate
Foa
Fost
Averge
Deep
BNN Gauesian Prior
BNN Emp
Figure 5: EmpCov prior improves robustness. Test accuracy under covariate shift for deep en- sembles, MAP optimization with SGD, and BNN with Gaussian and EmpCov priors. Left: MIP architecture trained on MNIST. Right: CNN architecture trained on CIFAR-10. 10. The EmpCov prior provides consistent improvement over the standard Gaussian prior. The improvement is particularly noticeable on the noise corruptions and domain shift experiments (SVHN, STL-10).
Is the poor generalization of BNNs under covariate shift surprising? The results presented in this paper of crucial practical importance, relevant to the applicability of Bayesian neural networks in virtually every real-world setting since the train distribution often not exactly the same as the test distribution. Ultimately, whether or not result is surprising is subjective, but there many reasons to find the dramatic performance degradation of Bayesian neural networks under shift surprising, which we outline in Appendix B.
Why focus on linear dependencies?
This focus is dictated by the structure of both fully funne-con and convolutional neural networks, which apply activation functions to linear combinations of features. Our results can be directly extended to other models, where different types of dependencies between input features would lead to the same lack of robustness to covariate shift. For example, in Appendix M we derive analogous results to subsection 5.2 for multiplicative neural network architectures [Trask al., 2018], where different form of dependence between features causes poor robustness.
Intermediate layers. While our analysis in section 5 is focused on the linear dependencies in the input features, similar conclusions can be made about intermediate layers of the network. For example, weights connected to dead neurons, which output zero, do not affect predictions of the model. We provide more details in Appendix N.
Linear Bayesian models. In models that are linear in parameters, such as Bayesian linear regression and Gaussian process regression, we typically saved from the perils of BMA under covariate shift discussed in Section 5, because the MAP and the predictive mean under BMA coincide. We provide further details in Appendix O.
BNNs in low- regime. While we focus on covariate shift, our results also suggest that BNNs may generalize poorly when the training dataset is extremely small. Indeed, in low-data regime may observe linear combinations of the features that are constant in the training data, but not on test. In Appendix E we show empirically that BNNs underperform MAP in low-data regime.
An optimization perspective on the covariate shift problem. Section 5.1, we argued that SGD is more robust to covariate shift than HMC because the regularizer pushes the weights that correspond to dead pixels towards zero. This effect is mainly obtained through explicit regularization, without which these weights will remain at their initial values. We study the effect of initialization and explicit regularization on the performance of SGD under covariate shift in Appendix P. We also show in Appendix P.2 that other stochastic optimizers such Adam [Kingma and Ba, 2014] and Adadelta [Zeiler, 2012] behave similarly to SGD under covariate shift. Finally, we discuss the effect of test data corruptions on the loss landscape in Appendix P.3 and argue that the relative sharpness of low density posterior samples makes these solutions, which are included in a BMA but not MAP, more vulnerable covariate shift.
Limitations. Due to the intense computational requirements of HMC, our experiments are limited to smaller models and datasets. Our analysis is focused on issues arising models that are non-line non-linear
10

---

in their parameters. Moreover, while our proposed priors help improve robustness, they do not entirely resolve the issue. For example, the Contrast dataset of Figure (right panel), the BMA is still underperforming MAP.
Conclusion. Our work has demonstrated, both empirically and theoretically, how linear depen- dencies in the training data cause Bayesian neural networks to generalize poorly under covariate shift - explaining the important and unexpected findings Izmailov al. [2021]. The scope of this research is exceptionally broad, relevant to the safe deployment of Bayesian methods in virtually any real-world setting. While the two priors we introduce achieve some improvement in BNN performance under covariate shift, we are only beginning to explore possible remedies. Our work intended as step towards understanding the true properties of Bayesian neural networks, and improving the robustness of Bayesian model averaging under covariate shift.
Acknowledgements
We thank Martin Arjovsky, Behnam Neyshabur, Vaishnavh Nagarajan, Marc Finzi, Polina Kirichenko, Greg Benton and Nate Gruver for helpful discussions. This research is supported with Cloud TPUs from Google's s TPU Research Cloud (TRC), and by Amazon Research Award, NSF I-DISRE 193471, NIH RO1DA048764-01A2 NSF I18-1910266, and NSF 1922658 NRT-HDR: FUTURE Foundations, Translation, and Responsibility for Data Science.
References
Maruan Al-Shedivat, Andrew Gordon Wilson, Yunus Saatchi, Zhiting Hu, and Eric P Xing. Learning scalable deep kernels with recurrent structure. The Journal of Machine Learning Research, 18(1): 2850-2886, 2017.
Martin Arjovsky. Out of distribution generalization in machine learning. arXiv preprint arXiv:2103.02667, 2021.
Ben Athiwaratkun, Mare Finzi, Pavel Izmailov, and Andrew Gordon Wilson, There many consis- tent explanations of unlabeled data: Why you should average. arXiv preprint arXiv: 180..000,,. 2018.
Gregory W Benton, Wesley J Maddox, Sanae Lotfi, and Andrew Gordon Wilson. Loss surface simplexes for mode connecting volumes and fast ensembling. arXiv preprint 102. 2021.
Christopher M. Bishop. Pattern Recognition and Machine Learning. Springer, 2006.
Charles Blundell, Julien Corebiie, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural networks. arXiv preprint arXivv1505.05424 2015.
Mariusz Bojarski, Davide Del Testa, Daniel Dworakowski, Bernhard Firner, Beat Flepp, Prasoon Goyal, Lawrence D Jackel, Mathew Monfort, Urs Muller, Jiakai Zhang, et al. End to end learning for self-driving cars. arXiv preprint arXiv: 160..07316 2016.
Ginevra Carbone, Matthew Wicker, Luca Laurenti, Andrea Patane, Luca Bortolussi, and Guido Sanguinetti. Robustness of bayesian neural networks to gradient-based attacks. arXiv preprint 20XY:22002.04359, 2020.
Carlos M Carvalho, Nicholas G Polson, and James Scott. Handling sparsity via the horseshoe. In Artificial Intelligence and Statistics, pages 73-80. PMLR, 2009.
Billiams-King, Oscar Chang, Yuling Yao, David Williams King, and Hod Lipson. Ensemble model patching: A parameter-efffcient variational bayesian neural network. arXiv preprint wxXY::9905.09433. 2019.
Adam Coates, Andrew Ng, and Honglak Lee. An analysis of single-layer networks in unsupervised feature learning. Proceedings of the fourteenth international conference on artificial intelligence and statistics, pages 215-223. JMLR Workshop and Conference Proceedings, 2011.

---

Tianyu Cui, Aki Havulinna, Pekka Marttinen and Samuel Kaski. Informative gaussian scale mixture priors for bayesian neural networks. arXiv e-prints, pages arXiv-2002. 2020.
Hal Daumé III. Frustratingly easy domain adaptation. arXiv preprint urXiv:0907.1815, 2009.
Hal Daume III and Daniel Marcu. Domain adaptation for statistical classifiers. Journal of artificial Intelligence research, 26:101-126, 2006.
Erik Daxherger, Eric Nalisnick, James Urquhart Allingham, Javier Antorán, and José Miguel Hernández-Lobato. Expressive yet tractable bayesian deep learning via subnetwork inference. arXiv preprint arXiv:2010. 14689, 2020.
Pedro Domingos. Bayesian averaging of classifiers and the overfitting problem. In ICML, volume 747, pages 223-230. Citeseer, 2000.
Michael Dusenberry, Ghassen Jerfel, Yeming Wen, Yian Ma, Jasper Snoek, Katherine Heller, Balaji Lakshminarayanan, and Dustin Tran. Efficient and scalable bayesian neural nets with rank-1 factors. In International conference on machine learning. pages 2782-2792, 2020.
Andre Esteva, Brett Kuprel, Roberto Novoa, Justin Ko, Susan M Swetter, Helen M Blau, and Sebastian Thrun. Dermatologist-level classification of skin cancer with deep neural networks. nature. (2277639):115-111 2017.
Angelos Filos, Sebastian Farquhar, Aidan N Gomez, Tim GJ Rudner, Zachary Kenton, Lewis Smith, Milad Alizadeh, Arnoud de Kroon, and Yarin Gal. A systematic comparison of bayesian deep learning robustness in diabetic retinopathy tasks. arXiv preprint arXiv: 2019.
Vincent Fortuin. Priors in bayesian deep learning: A review. arXiv preprint arXiv:2105.06868, 2021.
Vincent Fortuin, Adrià Garriga-Alonso, Florian Wenzel, Gunnar Rätsch, Richard Turner, Mark van der Wilk, and Laurence Aitchison. Bayesian neural network priors revisited. arXiv preprint arXiv:2102.06571, 2021.
Yarin Gal and Zoubin Ghahramani. Dropout Bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pages 1050-1059, 2016.
Varun Gulshan, Lily Peng, Mare Coram, Martin C Stumpe, Derek Wu, Arunachalam Narayanaswamy, Subhashini Venugopalan, Kasumi Widner, Tom Madams. Jorge Cuadros, al. Development and validation of a deep learning algorithm for detection of diabetic retinopathy in retinal fundus photographs. Jama, 316(22):2402222) 2410, 2016.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pages 1026-1034, 2015.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. arXiv preprint 1rXhv:1903.12261 2019.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 911:1144, 1997.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In Uncertainty in Artificial Intelligence (UAI), 2018.
Pavel Izmailov, Sharad Vikram, Matthew D Hoffman,, and Andrew Gordon Wilson. What are bayesian neural network posteriors really like? arXiv preprint arXiv:2104 104.1441, 2021.
Jörn-Henrik Jacobsen, Jens Behrmann, Richard Zemel, and Matthias Bethge. Excessive invariance causes adversarial vulnerability. arXiv preprint arXiv: 1.00401, 2018.
Norman P Jouppi, Doe Hyun Yoon, George Kurian, Sheng Li, Nishant Patil, James Laudon, Cliff Young, and David Patterson. domain-specific supercomputer for training deep neural networks. Communications of the ACM, 63(7):67-78, 2020.
12

---

Alex Kendall and Yarin Gal. What uncertainties do we need in Bayesian deep learning for computer vision? Advances in neural information processing systems, pages 5574-5584, 2017.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: 160.0066,0 2016.
Samuel Kessler, Vu Nguyen, Stefan Zohren, and Stephen Roberts. Hierarchical indian buffet neural networks for bayesian continual learning. arXiv preprini urXi:: 2019.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6990, 2014.
Diederik P Kingma and Max Welling, Auto-encoding variational bayes. arXiv:1312.6114, 2013.
arXiv preprint
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. The CIFAR-10 dataset. 2014. http::/ww..c... toronto. toronto.edu/krree edu/kriz/cifar .tml..
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive
uncertainty estimation using deep ensembles. In Advances in Neural Information Processing Systems, pages 6402-6413, 2017.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11))2278-2222, 1998.
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann. wwwp:/wwwwww.....eee/eeee//e wttp/////////iiitt..tmmm 2, 2010.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. arXiv preprint 17 2017.
Dong C Liu and Jorge Nocedal. On the limited memory bigs method for large scale optimization. Mathematical programming, 45(1-3):503-558, 1989.
David JC MacKay. Probable networks and plausible predictions-a review of practical bayesian methods for supervised neural networks. Network: computation in neural systems, 6(3):469-505, 1995.
Wesley Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson.
A simple baseline for baycsian uncertainty in deep learning. Advances in Neural Information Processing Systems, 32:13153-13164 2019.
Andres R. Masegosa. Learning under model misspecification: Applications to variational and ensemble methods, 2019.
Rhiannon Michelmore, Matthew Wicker, Luca Laurenti, Luca Cardelli, Yarin Gal, and Marta Kwiatkowska. Uncertainty quantification with statistical guarantees in end-to-end autonomous driving control. In 2020 IEEE International Conference on Robotics and Automation (ICRA), IEEE, 2020.
pages
John Miller, Rohan Taori, Aditi Raghunathan, Shiori Sagawa, Pang Wei Koh, Vaishaal Shankar, Percy Liang Yair Carmon, and Ludwig Schmidt. Accuracy on the line: On the strong correlation between nut-of-diitribution and in-distribution generalization. International Conference on Machine Learning, pages 7721-7735. PMLR, 2021.
Thomas Minka. Bayesian model averaging model combination. 2000.
Dmitry Molchanov, Arsenii Ashukha, and Dmitry Vetrov. Variational dropout sparsifies deep neural networks. In International Conference on Machine Learning, pages 2498-2 2598-2507. PMLR, 2017.
Warren R Morningstar, Alexander A Alemi, and Joshua V Dillon. Pac m-bayes: Narrowing the empirical risk gap in the misspecified bayesian regime. arXiv preprint arXiv:2010.09629, 2020.
13

---

Norman Mu and Justin Gilmer. Mnist-c: A robustness benchmark for computer vision. arXiv preprint arXiv:1906.023333 2019.
Vaishnavh Nagarajan, Anders Andreassen, and Behnam Neyshabur. Understanding the failure modes of out-of-distribution generalization. arXiv preprint arXiv:2010. 2020.
Radford M Neal et al. Meme using hamiltonian dynamics. Handbook of markov chain monte carlo, 2(11):2, 2011.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.
Behnam Neyshabur, Hanie Sedghi, and Chiyuan Zhang. What is being transferred in transfer learning? arXiv preprint arXiv:2008. / /687, 2020.
Jorge Nocedal. Updating quasi nuwsi-newton matrices with limited storage. Mathematics of computation, 35(151):773-78,, 1980.
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D Sculley, Sebastian Nowozin, Joshua V Dillon, Balaji Lakshminarayanan and Jasper Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. arXiv preprint arXiv:/906.000
2019.
Boris T Polyak. Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics, 4(5):1-17, 1964.
C. E. Rasmussen and C. K. I. Williams. Gaussian processes for Machine Learning. The MIT Press, 2006.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pages 400-407, 1951.
Abhijit Guha Roy, Jie Ren, Shekoofeh Azizi, Aaron Loh, Vivek Natarajan, Basil Mustafa, Nick Pawlowski, Jan Freyberg, Yuan Liu, Zach Beaver, et al. Does your dermatology classifier know what it doesn't know? detecting the long-tail of unseen conditions. arXiv preprint arXiv:2104.03829 2021.
Hidetoshi Shimodaira. Improving predictive inference under covariate shift by weighting the log- likelihood function. Journal of statistical planning and inference, 90(2):227-244, 2000.
Amos Storkey. When training and test sets different: characterizing learning transfer. Dataset
shift in machine learning, 30:3-28, 2009.
Amos Storkey and Masashi Sugiyama. Mixture regression for covariate shift. Advances in neural information processing systems, 19:1337 1337, 2007.
Masashi Sugiyama, Benjamin Blankertz, Matthias Krauledat, Guido Dornhege, and Klaus-Robert Müller. weepora cross-validation for covariate shift. In Joint Pattern Recognition Symposium, pages 354-363. Springer, 2006.
Masashi Sugiyama. Matthias Krauledat, and Klaus-Robert Müller, Covariate shift adaptation by importance weighted cross validation. Journal of Machine Learning Research, 8(5), 2007.
Andrew Trask, Felix Hill, Scott Reed, Jack Rae, Chris Dyer, and Phil Blunsom. Neural arithmetic logic units. arXiv preprint 1rxiv:1808. 00508, 2018
Florian Wenzel, Kevin Roth, Bastiaan S Veeling, Jakub Šwiatkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How good is the Bayes posterior in deep neural networks really? arXiv preprint 0tl::2002.02403: 2020.
Matthew Wicker, Luca Laurenti, Andrea Patane, Zhuotong Chen, Zheng Zhang, and Marta Kwiatkowska. Bayesian inference with certifiable adversarial robustness. International Confer- ence on Artificial Intelligence and Statistics, pages 2431-2439. PMLR, 2021.
Andrew Gordon Wilson and Pavel Izmailov. Bayesian deep learning and probabilistic perspective of generalization. arXiv preprint urXiv:2002.08791 2020.
Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprint arXiv: 1xXiv:1212.570.) 2012.
14


Reviews:
[{"review_id": "bv7uOTnDR8", "reviewer": "Reviewer_uotf", "summary": "This paper studies an important and interesting problem that the BNN with HMC inference performs worse than point estimation under covariate shift. The paper conducts theoretical analysis and contributes this to the linear dependencies in the input feature. A data-dependent prior is proposed for improving the robustness under covariate shift.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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": "Originality:\nThis paper is of high originality, as it provides a new perspective to explore the reason of poor robustness of HMC. The EmpCov prior is reasonable with previous analysis.\n\nQuality, clarity:\nThe paper is well-written and easy to read. I like that every claim is back up by empirical results. The graphs are clear and helpful in understanding the paper.\n\nSignificance:\nAlthough this paper provides useful analysis and methods, I would question that if the analysis and method only work on small datasets like MNIST where linear dependence in features are prevalent. How about more complex datasets with rich data?\nThe data dependent approach might not be scalable to the size of data, as the covariance should be calculated.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "OYAhqJdH_Pp", "reviewer": "Reviewer_YvmA", "summary": "**Post-rebuttal update.** Thank you for your response. I did not have any major concerns with this work, and have also read the other reviews and the author responses and did not find anything concerning. Thus I am keeping the rating unchanged.\n\nThis work investigates the performance of Bayesian model averaging (BMA) prediction for NN models under covariate shift. It\n* provides further empirical evidence that, for toy models on image classification tasks, BMA may perform poorly on covariate shifted data even though it has competitive performance on clean data.\n* provides a (partial) explanation by observing the difference of the posterior and MAP estimate of the first-layer parameters: the MAP estimate shrinks prediction on unobserved \"subspaces\" (meaning orthogonal linear subspace for Gaussian prior, and zero-variance feature dimensions for general factorized prior) to zero, while the posterior must behave like prior on these subspaces, leading to more noisy feature extractions and deteriorated performance.\n* proposes to use as the first-layer prior the $N(0, \\Sigma+\\epsilon^2 I)$ distribution, where $\\Sigma$ is the empirical data covariance.  Experiments show that the prior leads to improved robustness.", "questions": "", "limitations": "", "rating": 7, "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": "This work provides new insight on an important problem.  The idea, as summarized above, is obvious (in hindsight) and quite convincing, and is well supported by experiments.  The explanation is incomplete as it does not provide much intuition for the subsequent layers (besides the difference on \"dead neurons\"), and the experiments are limited in scale, but both issues are understandable and acknowledged in discussion.\n\nI do not have any major concern with this work. There are a few issues / questions however:\n\n* There should be more discussion on the past observations that the distribution of NN weights at lower layers are more heavy-tailed (e.g., Fortuin et al. [13]).  Even though the observations are made in different contexts, and experiments here show the Laplace prior does not alleviate the problem, the high-level intuition is similar, and it is possible that a more extreme sparsity-promoting prior is needed to address this issue (e.g. the spike-and-slab prior).\n\n* The most interesting / practically relevant results are Prop 2-3, but they only applies to Gaussian priors. Non-Gaussianity clearly breaks the proof, but does it lead to qualitatively different behavior for feature extraction at the first layer (or is its performance deterioration due to different mechanisms)? It will be helpful to re-do Figure 4 using Laplace or student-t priors for this purpose.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The most important limitations are the limited scale of the experiments, and the lack of complete explanation for the subsequent layers / nonlinear dependencies. They are acknowledged in the paper, and for the latter, some toy-ish scenarios are analyzed in the appendix.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "CKItYbRODzX", "reviewer": "Reviewer_2XLh", "summary": "This paper investigates the behavior of Bayesian model averaging under covariate shift. This is mainly an experimental paper, in which the authors explain the poor generalization ability of Bayesian inference by possible linear dependencies in the input features that cause a lack of posterior contraction, and they show why alternative approximate inference procedures are not affected. They also propose a way to fix the issue via novel priors.", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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": "I really enjoyed the paper, which is clear, well written, and brings novel ideas to the community. It is in line with the recent work of Izmailov et al. (2021) that showed that exact Bayesian inference in deep learning (via Hamiltonian Monte Carlo) generalizes poorly although approximate Bayesian inference is widely used in this setting. Supported with extended numerical experiments, this work validates trough extensive experiments on different types of covariate shift the assumption that the issue of averaging is due to dependencies in the features of the train data distribution that make some model parameters do not affect the predictions on the train data (in which case the posterior coincides with the prior). Hence, the authors use this finding to design new priors that assign low variance to relevant weights. I really appreciated the investigation and the results of the paper.\n\nI have a question regarding the technique used in the paper. The authors say that they use full-batch HMC, and alternative approximate inference techniques such as SWAG and MC Dropout are studied in Appendix G. Both fix the issue, but the influence of the prior in variational inference plays a particular role, in opposite to SWAG for instance which approximates the posterior distribution as a multivariate Gaussian with the SWA solution as its mean. Would the poor generalization issue of the exact posterior be fixed as well for Gaussian Gaussian approximations / in which case, what would be the influence of the prior ?\n\nAlso, I would have enjoyed a more general discussion relating the current study with the slightly different notion of misspecification. Could the authors elaborate on this ?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10"}, "extra_text": {}}, {"review_id": "8I2lv6fdNO", "reviewer": "Reviewer_anCt", "summary": "The paper explores the performance of Bayesian neural networks (BNNs) under covariate shift, comparing primarily to MAP and deep ensembles. The authors explain why, when using standard priors, the test accuracy of HMC BNNs degrades quicker than MAP and deep ensembles when the models are tested on increasing levels of corruption. Using the insights from their explanation, the authors provide a new prior, 'EmpCov', that attempts to address this.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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": "Overall, I am quite conflicted on this paper. On the positive side, it is well-written with clear explanations of the propositions and results, and the proposed 'EmpCov' prior does appear to work quite well while being a relatively simple fix. From what I can tell, the results and derivations are correct. The figures are also well-made and the authors provide extensive appendices covering a large array of related topics.\n\nHowever, I am not convinced by the framing of the central issue of the paper. The authors present the degradation of test *accuracy* that BNNs exhibit under covariate shift specifically as an issue that should be surprising (e.g. lines 26-29). They claim that a BNN should be more robust to covariate shift than a deterministic network, in particular with respect to test accuracy, and test accuracy is the main focus of the work. However, I do not see any reason that this should be the case. Moreover, after searching the literature (in particular the works the paper cites), I could not find any work that specifically claimed that the test accuracy of BNNs *should* be higher than those of deterministic NNs. \n\nFor me, it makes sense that the accuracy would be worse, as the BNN will increasingly match the prior for test points that are increasingly far from train points, and the prior predicts a uniform distribution over classes. In my view the main advantage of using a BNN, which is the claim that is often repeated in the literature, is that it should have better-calibrated uncertainty than a deterministic NN, so that the end user is aware when it has encountered an example it has not seen enough data to reliably classify. Indeed, in Fig. 6 in the appendix (page 17), we see that the standard Gaussian prior BNN far outperforms MAP networks on test log likelihoods (test LLs) and expected calibration error (ECE), and has approximately equal performance to the proposed EmpCov prior BNNs. This indicates that standard BNNs (at least with HMC) are much better calibrated than MAP networks, which on average overconfidently predict the correct class more.\n\nNevertheless, the paper does clearly explain one of the potential mechanisms by which standard BNNs provide predictions closer to uniform for far-away test points, and the proposed prior is able to obtain similar LLs and ECE to standard BNNs with improved accuracy. Therefore, the work clearly has value.\n\nOverall, the paper does a good job investigating the worsened accuracy that BNNs obtain under covariate shift, exploring a mechanism for the worsened accuracy and providing a new prior that addresses this mechanism. However, I believe the framing is not quite correct, in that it is not clear why one would necessarily expect the accuracy of standard BNNs to be better than those of MAP networks on out-of-distribution data. Moreover, standard BNNs are clearly better calibrated when it comes to uncertainty. If the standard BNNs had worse test LLs and ECE then the framing would be far more convincing, but this is clearly not the case. Therefore, I hope the authors would address this with further clarification and by focusing the discussion more on uncertainty calibration (which is the main desired benefit of using a BNN in covariate shift settings). Indeed, it seems as though EmpCov prior can effectively bring together the higher accuracy of the MAP networks and the improved calibration of more standard BNN priors, and I think the paper would be much more effective if it focused more on this rather than claiming that it is surprising that BNNs have reduced accuracy over MAP in OOD regions.\n\nAs I stated at the beginning of the review, I am quite borderline on this paper, and so I am looking forward to hearing the authors' response as well as what the other reviewers think.\n\nAdditional questions/comments:\n- I am not entirely convinced by the arguments concerning linear models given in lines 326 - 330. The predictive mean and MAP solutions only coincide for regression, as the predictive variance will affect the of the output of e.g. a softmax for classification. Along these lines, I would be curious how the authors feel the NNGP [1, 2] would perform in comparison to the finite-width BNNs. I would expect for the MLPs at least that the behavior would be quite similar, because the MLPs used in the paper are 256 units wide and relatively shallow, which should be wide enough to obtain NNGP behavior.\n- In many cases neural networks are trained without explicit regularization in the form of an L2 penalty. Would the authors expect these MLE networks to perform worse in terms of test accuracy on corrupted data? Following the logic of the paper it would seem they should.\n- Izmailov et al. (2021) show that beyond MAP-based methods, MFVI and SGLD also seem to outperform HMC when it comes to covariate shift. Additionally, the test LLs of MAP seem to be much more even with HMC than what you have observed. Do you have any theories why this may be?\n\nMinor/typos:\n- line 120: footnotes should appear after punctuation\n- line 227: 'Propositions' shouldn't be plural\n- line 930: should be 'then' rather than 'than'\n\n[1] https://arxiv.org/abs/1804.11271\n[2] https://arxiv.org/abs/1711.00165", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors briefly describe some of the main limitations of their work. I am satisfied by their exploration of these limitations, and agree that this work is a sufficient first exploration of these issues. They state that the main negative societal impact of their work was the high energy usage that went into the HMC experiments. As the work is focused on fundamental properties of BNNs, I believe this is adequate and do not believer there are any other direct negative societal impacts.", "review_text": "", "extra_scores": {"ethical_concerns": "I do not have any ethical concerns with this work.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers have provided thoughtful and constructive comments. They have responded to the authors' feedback and the most active reviewers who are championing this manuscript continue to lean towards acceptance. I personally find that the authors are presenting an interesting set of contributions in the under-represented and -valued field of model evaluation. I hope the authors will take the reviewers' comments to heart and encourage them to incorporate their thoughts in preparing the camera-ready version of their manuscript.
Decision: Accept (Poster)