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": "cMs6TnfeaH", "reviewer": "Reviewer_SBkb", "summary": "This paper begins with the observation that BMA underperforms MAP or ensembling under distribution shift and corruptions.\nIt observes that this could well be because of weights in the first layer reverting to the prior along dimensions which are off the data-manifold.\nThese weights can take on high values, even though their mean is low, and the resulting influence on the prediction might be high because of the nonlinearity.\nThe authors empirically support their analysis by proposing a novel data-dependent prior which has low variance along these dimensions.\nThey find that in some settings this prior improves accuracy under corruption, although perhaps not always.", "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": "# Overall review\nThere were parts of this paper that I thought were really excellent.\nThe analysis in section 5 is very strong, and I particularly enjoyed the experiment in Figure 4, which is a clever test of the hypothesis.\n\nI am not currently recommending acceptance, although I am uncertain and borderline.\nI felt that there might be some alternative explanations for the benefit of ensembling/MAP over BMA which were not considered or properly ruled out (see below).\nAt the same time, the main empirical justification for your diagnosis of the problem lies in the effectiveness of the EmpCov prior, but you noted that the prior was not effective in MNIST for CNNs, although I couldn't find any figures showing this result.\nOff the back of the strength of the analysis of section 5, I believe that you have indeed identified *an* issue for BMA under distribution shift, but I don't yet feel confident that this is the main reason for the performance gap that you discuss.\n\n## Originality\nOn the whole, the work makes some novel contributions.\nHowever, large parts of the text and figures build heavily off of Izmailov et al. \"What are Bayesian Neural Network Posteriors Really Like?\". ICML. 2021. It isn't really until page 5 that things really start to take off with novel work.\n\n## Quality\nThe quality was generally high, and I was impressed by the number of ablations performed.\nI make a note below about the terminology \"covariate shift\", which I think may not be correct here.\nThere were also some textual descriptions that referred to figures that I could not find.\n\n## Clarity\nOn the whole, I felt that the text and figures were clear.\n\n## Significance\nI could not quite tell how significant the work seems likely to be.\nThe authors suggest in the introduction that their novel prior could be generally adopted, but don't heavily discuss the risks and limitations of data-dependent priors.\nSetting aside the prior, it seems likely that the work could motivate further use of MAP and ensembles, however since these are already heavily used it is unclear how much this impact will extend beyond what was already known in prior work.\nThe analysis of S5 feels quite satisfying though, and perhaps it will unlock deeper insights in the future.\n\n# Major comments\n\n### Suppose the HMC is working, is your MC integration working?\n\nI'll grant the assumption HMC provides a good set of samples from the true posterior p(w|D) (though reasonable people could disagree about this assumption).\nBut it doesn't necessarily follow that the MC integral (your Equation 2) using HMC samples is a good estimate of the marginal p(y|x).\nAn integration failure seems especially likely given the comparatively small number of samples used for integration (90) compared to the dimensionality of the space.\nWould you be able to validate the accuracy of your integration?\n\nI understand that the non-linearity means that even under perfect integration the expectation will not match the MAP result.\nIt would be nice to demonstrate directly that this noise is *actually* amplified by the non-linearity, rather than being cancelled out higher up the network.\nI wouldn't be surprised if in large and deep networks this effect is somewhat corrected for naturally (for variational inference, at least, it has been found to be somewhat corrected by 'over-pruning' early weights that are set near the prior (see Trippe and Turner \"Overpruning in Variational Bayesian Neural Networks\". 2018).\nThe question is how big the effect you've identified is compared to orthers.\nI understand you are computationally constrained from getting much bigger models with HMC, of course.\n\n### Comparison with Domingos, P. \"Bayesian averaging of classifiers and the overfitting problem.\" ICML 2000.\n\nHe made a similar observation to yours: empirical failure of principled BMA compare to the more naive ensembling which is not weighted by p(w|D).\nBut his work was in the context of trees, rather than NNs with HMC.\nHis explanation of the failure of BMA, is that the BMA overfits to parameter-values that have particularly high probability under the data.\nHMC won't have the same overfitting problem, but I think your analysis might point to something similar going on based on the fact that the training dataset is only a manifold of the overall-dataspace.\nI don't think this is inconsistent with your analysis, and it might be either a special case or competing explanation, but seems relevant to your S5.4.\n(Note some commentary in Minka TP. \"Bayesian Model Averaging Is Not Model Combination\". 2000. which may also shed some light on the issues.)\n\n### Comparison with Izmailov et al. \"What are Bayesian Neural Network Posteriors Really Like?\". ICML. 2021.\n\nIt seems like there's a reasonable amount of overlap between much of this work and results in this paper by Izmailov et al.\nIn particular, it isn't until section 5 that any substantially new results or observations appear to be introduced.\nThe main differences before that point (as far as I can tell) are the addition of MNIST and the use of smaller LeNet-style networks rather than ResNet-20 (presumably for computational reasons).\nFigure 1(b,c) is also completely novel, but I would have actually liked more investigation of this interesting figure!\nThis is not necessarily a huge issue---the analysis and discussion of the results and observations extends beyond the other work---but I wonder if the paper would be stronger if more space was devoted to new results.\n\n### Figure 4 and S5.2 are really excellent\nThis is such a clever experiment and a great result.\nIt should be more heavily emphasised earlier on.\n\n### Use of the term 'covariate shift'\nln 15: strictly speaking, 'covariate shift' is a special case of the target distribution differing from the training distribution, as you note in S2.\nThe 'OOD generalization tasks' you refer to not only change p(x) but also change p(y|x), and are therefore not covariate shift.\nI think the CIFAR-10-C datasets are probably also not covariate shift.\nFor example, if you take some dataset X and add Gaussian noise to X such that $X' = X + \\epsilon$, but keep $\\mathbf{y}$ constant, then it seems unlikely that p(y|x) is unchanged.\nAfter all, there are presumably images in the original (un-noised) data distribution which could only have been from class 1 but which are now ambiguous between class 1/2.\nI'm not certain that this undermines anything you have specifically introduced, but it means that some of the literature you relate to or draw on might need to be interpreted more carefully and that results that are specific to covariate shift might not actually apply in the settings which you are exploring.\n\nSentences like ln 116 \"In this case, there is no semantic distribution shift:\" I'm not sure if this is true given the above.\nOr something like MNIST to SVHN, it isn't clear what p(y|x) in MNIST is for SVHN images---MNIST can be perfectly labelled with a p(y|x) that has zero accuracy on crop/greyed SVHN and vice versa.\nTo be clear, I don't really think that this point is all that important to the contribution of your paper---and it hasn't formed a significant part of my score.\n\n### EmpCov in CNN on MNIST\n>In Appendix I, we provide a detailed analysis of the performance of the CNN architecture on MNIST.\nSurprisingly, we found that using the EmpCov prior by itself does not provide a large improvement in this case.\n\nI was expecting to find a chart of EmpCov on CNNs in MNIST in appendix I and I did not, did I miss it?\nI would perhaps have expected to find it in Figure 7 or 11, but it seemed not to be there.\nWhen you say in line 311 it \"does not provide a large improvement\", does it provide any improvement? Does it hurt?\nIf so, why?\n\nIn line 312 you write that the issue is \"specific to this particular setting\", and in the appendix you explain that the \"brightness\" and \"fog\" corruptions are especially problematic.\nBut I don't immediately see why this would be much worse for MNIST than CIFAR, it seems worth substantiating that a little.\nPresumably with a MNIST CNN equivalent of Figure 7 we could see if there was a similar effect for things other than \"brightness\" and \"fog\".\nIf I understand it right, what you've written is consistent with there being a benefit to EmpCov for MNIST CNNs for all the corruptions except those two.\nIs that the case?\n\nGiven that the proposal of a new prior is both a major contribution of your paper and one of the main empirical evaluations that the issue you have identified is a major issue for the difference between BMA and ensembling, I would definitely devote more space to examining these issues and the SumFilter at the expense of some of the repetition of results from prior work.\n\n# Minor remarks\nln 39: \"similar to a draw from the Gaussian prior.\"\nIt's not clear what sense of 'similar' you are referring to, and you should probably substantiate it.\nDo you just mean it is noisy?\nHow much like a draw from a Gaussian is it?\nIn what way?\nWhat do the means of the model average look like for these weights? (i.e., it would be interesting to contrast this plot with (d) BNN weights averaged over N samples and see if it looks like (c))\n\nln 298: I'd like a little more discussion of the use of this as a prior.\nIn particular, it is likely to lead to weight posteriors which are highly overconfident in some parts of the space.\nIf this were used for any kind of sequential inference with multiple datasets, say, it would be quite inappropriate.\nThat's all fine, it's just worth noting.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "This is fine on the whole. More acknowledgement in the introduction of the fact that the prior is data-dependent and the limitations this entails, as well as the fact that they seem to work slightly inconsistently, might help slightly.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "Erh6NrCuxSi", "reviewer": "Reviewer_8WYY", "summary": "This paper shows that Bayesian neural networks are vulnerable to distribution shifts. The authors provide theoretical results showing that this vulnerability is caused by the linear dependencies in the training samples, which are supported by their experiments as well. Based on this theoretical insight, the authors propose a simple initialization method to resolve this undesirability. \n", "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": "Evaluating and improving the robustness of Bayesian neural networks to distribution shifts is an important research area getting increasing attention in the community. In this regard, this paper provides valuable insights that explain why Bayesian neural networks are vulnerable to covariate shifts and how to overcome this vulnerability. \n\nIn addition, the paper is well-written, intuitive, and easy to follow. I think that the motivating examples in Fig. 1 and Chapter 4 are informative to both experts and readers who are not familiar with the Bayesian neural networks’ vulnerability to the distribution shift. Also, the structure of Chapter 5, starting from the simplest case in which all training samples have a constant zero at a particular feature with feed-forward neural nets and then generalizing the case, helps readers get insights into the phenomenon. Also, the theoretical result in Proposition 3 covers sufficiently large cases with mild assumptions. Finally, the proposed initialization method is simple and effective. \n\nI believe this is a good paper, and I have only two minor comments to authors: \n\n1. Regarding Fig. 1 (b), the authors present BNN weights sampled HMC and claim that they are unstructured. However, the features obtained by marginalization might have the structured form, e.g., zeros, or small values, on the boundaries. Considering the posterior weights on the boundary location would follow zero-mean Gaussian (or Laplace depending on the prior), I think this case is highly likely. If so, I’m wondering if the authors can still argue that the BNN weights are unstructured. \n\n2. Looking at the theoretical results in Chapter 5, there are no restrictions on the way to perform posterior inference of BNNs. However, I’m wondering why the authors did not provide results for BNNs with stochastic VI, even though they are more frequently used in practice. \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors properly addressed the limitation in the paper. Also, the authors provided the potential negative social impact caused by the intensive computational load of HMC.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "15"}, "extra_text": {}}, {"review_id": "EgHh3PtTi6y", "reviewer": "Reviewer_SrTJ", "summary": "The paper conducts empirical and theoretical analysis on why Bayesian neural networks (BNNs) with high-fidelity inference schemes (HMC) are vulnerable to out-of-distribution data. The authors empirically noticed that BNNs perform worse than the point estimation MAP on the corrupted test dataset. They propose that it is the linear dependency in the input features that causes this lack of robustness. This input linear dependency will not provide likelihoods for the weight distributions in the orthogonal direction, thus won't modify the weights' prior distributions. Those unmodified weight distributions are vulnerable to input perturbations. Based on these observations, the authors propose a new prior covariance to providing more robustness.  ", "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": "The paper is well-written, and the delivery is clear. The theoretical perspective is interesting. But the empirical demonstration is not that novel because a similar observation exists as the author mentioned, although the paper provides more evidence. Furthermore, the paper would be more motivated if the following two points can be addressed:\n- In the first place, the analyses are all interesting, but why would people want to apply their model to out-of-distribution data? Aren’t people expected to detect the ODDs first and exclude them from the testing dataset? \n- In Fig. 4, it seems that the prior already has a quite small scale. So sampling from those priors will generate weights almost zero (if my understanding is correct), which is somewhat the opposite of the paper's argument. Will this small error during testing propagate through the network? (Please correct me if my interpretation is wrong.)", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The mentioned limitations are good. Practitioners are also interested to see the results or mentions of another approximate inference scheme -- variational inference -- which is more computationally efficient and widely used. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3.5"}, "extra_text": {}}, {"review_id": "BYRHxeaEB0i", "reviewer": "Reviewer_rUts", "summary": "The paper explores when predictions relying on the posterior predictive distribution of a Bayesian neural network can be poor in the presence of covariate shift and provides a plausible explanation. The main argument being that weakly informative likelihoods cause a lack of posterior contraction, causing the posterior to revert to the prior for at least a subset of the model parameters. The resulting posterior-predictive reverts to the prior predictive and can be detrimental in covariate shift if the shifted data need to exploit this subset of parameters that lack posterior contraction.", "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": "I enjoyed reading this paper. It is well written, clearly demonstrates the empirical phenomenon of Bayesian neural networks struggling under certain (but not all) types of covariate shift, and proposes a plausible explanation for the observed phenomenon. Moreover, the fact that the authors used full-batch HMC leads me to believe that the observed issues are likely not a result of approximate inference procedures producing poor approximations and lends credence to the proposed explanation. I only have a few minor quibbles with the paper.\n\n- The paper does a good job of demonstrating that Bayesian model averaging (BMA) is sensitive to perturbations orthogonal to the data subspace when the training data is low rank. What is less clear is whether BMA is as robust as its point-estimated MAP counterpart when the data is full rank or when the perturbations cause shifts in the subspace rather than out of it. I believe Figure 4b was an attempt at showing this. Still, it would be nice to separate it from any idiosyncrasies introduce by the highly curated MNIST dataset and is perhaps best illustrated on simulated data.\n- It is interesting that in Figure 4a, for the MLP and CNN (N(0, 0.01) for the high variance (larger PC) components, the posterior is more diffuse than the prior, and there appears to be no posterior contraction. Do the authors have an explanation for this? Is this a case of not seeing enough data (relative to the large number of parameters in the model) or a case of prior misspecification (apriori variances are too small?)? \n- Bayesian neural networks are sometimes sold as being effective in the small data regime as a guard against overfitting. However, one would expect the posterior to similarly display a lack of contraction in those cases (there isn’t enough data to condition on). I am curious if the authors believe that the posterior-predictive in such cases would revert to the prior-predictive and result in the BNN ignoring the data? \n- Finally, the title of the paper likely needs a qualifier considering that these issues stem from severely over-parameterized models. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The limitations are adequately discussed. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4.5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper explores when predictions relying on the posterior predictive distribution of a Bayesian neural network can be poor in the presence of covariate shift and provides a plausible explanation. Weakly informative likelihoods cause a lack of posterior contraction, causing the posterior to revert to the prior for a subset of the model parameters. The resulting posterior-predictive reverts to the prior predictive and can be detrimental in covariate shift.

The paper convincingly demonstrates the empirical phenomenon of Bayesian neural networks struggling under certain types of covariate shifts and proposes a plausible explanation for the observed phenomenon. It is solid in terms of its empirical analysis. Furthermore, it presents a new prior that does not suffer from the downsides of the conventional priors commonly used.  

Initial doubts about a lack of samples to evaluate the posterior predictive distribution were successfully addressed in the discussion period. Likewise, initial concerns about a lack of novelty with respect to [Izmailov et al., ICML 2021] were also clarified. Another concern was that the proposed phenomenon might not be the only explanation for BNNs failing in OOD situations, but the reviewers didn't consider this a significant weakness of the paper. 
Decision: Accept (Poster)