Robustness to corruption in pre-trained Bayesian neural networks
Authors: Xi Wang
Laurence Aitchison
Venue: ICLR 2023 poster (2023)

Abstract: We develop ShiftMatch, a new training-data-dependent likelihood for robustness to corruption in Bayesian neural networks (BNNs). ShiftMatch is inspired by the training-data-dependent “EmpCov” priors from Izmailov et al. (2021a), and efficiently matches test-time spatial correlations to those at training time. Critically, ShiftMatch is designed to leave the neural network’s training time likelihood unchanged, allowing it to use publicly available samples from pre-trained BNNs. Using pre-trained HMC samples, ShiftMatch gives strong performance improvements on CIFAR-10-C, outperforms EmpCov priors (though ShiftMatch uses extra information from a minibatch of corrupted test points), and is perhaps the first Bayesian method capable of convincingly outperforming plain deep ensembles.

Published as a conference paper at ICLR 2023
ROBUSTNESS TO CORRUPTION IN PRE-TRAINED BAYESIAN NEURAL NETWORKS
Xi Wang
Laurence Aitehison Department of Computer Science University of Bristol 1 laurence. a
College of Information and Computer Science University of Massachusetts Amherst xwang3@cs umass. edu
ac. uk
ABSTRACT
We develop ShiftMatch', new training-data-depenedd data-dependent likelihood for robustness to corruption in Bayesian neural networks (BNNs). ShiftMatch is inspired by the trainin-ddtt-dependeer "EmpCov" priors from Izmailov et al. (2021a), and efficiently matches test-time spatial correlations to those at training time. Criti- cally, ShiftMatch designed to leave the neural network's 's training time likelihood unchanged, allowing it to use publicly available samples from pre-traines BNNs. Using pre-trained HMC samples, ShiftMatch gives strong performance improve- ments on CIFAR-IO-C. 10-C, outperforms EmpCov priors (though ShiftMatch uses extra information from a minibatch of corrupted test points), and is perhaps the first Bayesian method capable of convincingly outperforming plain deep ensembles.
INTRODUCTION
Neural networks are increasingly being deployed in real-world, safety- settings such as self- driving cars (Bojarski et al., 2016) and medical imaging (Esteva et al., 2017). Accurate uncertainty estimation in these settings critical, and common approach is to use Bayesian neural networks (BNNs) to reason explicitly about uncertainty in the weights (MacKay, 1992; Neal, 2012; Graves, 2011; Blundell et al., 2015; Gal & Ghahramani, 2016; Maddox et al., 2019; Aitchison, 2020a;b; Ober & Aitchison, 2021; Unlu & Aitchison, 2021; Khan & Rue, 2021; Daxberger et al., 2021). BNNs are indeed highly effective at improving uncertainty estimation in the in-distribution setting, where the train and test distributions are equal (Zhang et al., 2019; Izmailov et al., 2021b). Critically, also need to continue to perform effectively (or at least degrade gracefully) when presented with corrupted inputs. Superficially, BNNs seem like a good choice for this setting: we would hope they would give more uncertainty in regions far from the training data, and thus degrade gracefully as inputs become gradually more corrupted, and thus diverge from the training data.
However, recent work has highlighted that BNNs including with gold-standard Hamiltonian Monte Carlo (HMC) inference can fail generalise corrupted images, potentially performing worse than ensembles (Lakshminarayanal) et al., 2017: Ovadia et al., 2019; Izmailov et al., 2021a;b). Izmailov al. (2021a) gave key intuition as to why this failure might occur. In particular, consider directions in input space with zero variance under the training data. As the weights in this direction have little or no effect on the output, any weight regularisation reduces the weights these directions to zero. The zero weights imply that these directions continue to have effeett the outputs, even if corruption subsequently increases variance in these input directions. However, BNNs do not work like this. BNNs sample weights in these cero-input-variin variance directions from the prior. That is fine in the training data domain, as there is no variance in the input in those directions, so the non-zero weights do not affect the outputs. However, if corruption subsequently increases input variance in those directions, then those new high-variance inputs will interact with the non-zero weights to corrupt the output.
Izmailov et al. (2021a) suggested an approach for fixing this issue, by modifying the prior over weights the input layer to reduce the variance in the prior over weights directions where the inputs have little variance. While their approach did outperform BNNs with standard Gaussian priors, performed comparably to deep ensembles (Izmailov al., 2021a, their Figs. 4,11,12). This failure is surprising: part of the promise of BNNs is that they should perform well for corrupted inputs by giving more uncertainty away from the training data.
Code available https: /ttp::///iiii tttpb/tttttttgoe/gg com/ didul u/Shi FtMatch

---

Published as a conference paper at ICLR 2023
0.0
-0.2
0.85
Log-Lii Likelikkihodd
Accuracy
0.80
0.75
0.70
-0.8
0.65
110
0.60
-1.2
0.55
0.50
-1.4
4
Corruption Intensity
Corruption Intensity
HMC
SGD
Deep Ens
SGLD
MFVI
ShittMatch (CIFAR-10.1)
ShiftMMtthh
Figure 1: Performance of HMC samples ,SGD, Deep Ensemble, SGLD, mean mield VI extracted from (Izmailov et al., 2021b) and ShiftMatch (applied on HMC samples from Izmailov et al., 2021b) on CIFAR-IO-C averaged over different corruption types with intensity level (1 to 5). ShiftMatch (CIFAR-10.1 Recht et al., 2019 11) uses the 2000 new images from CIFAR-10.1 to compute the training statistics, demonstrating that we get only a small performance penalty the full training data is not available, but can get access data from roughly the same distribution. Intensity 0 stands for clean CIFAR-10 test set without corruption. ShiftMatch significantly improves the robustness against corruption compared with plain HMC and even outperforms Deep Ensemble, but uses additional information from a minibatch of corrupted test datapoints, which is not used by the other methods.
However, the performance of any Bayesian method on in-distribution and corrupted inputs in depends heavily on the choice of model (prior and likelihood combined). We begin by noting that the as
EmpCov priors from Izmailov et al. (2021a) can be equivalently viewed likelihoods. We then develop new raining-data-dependeee ShiftMatch, which has two advantages over EmpCov priors (Izmailov et al., 2021 First, EmpCov priors apply only to the input layer, so might not be effective for more complex corruptions which are best understood and corrected at later layers. In contrast, our likelihoods modify the activity at every layer in the network, so have the potential to fix complex, nonlinear corruptions. Second, EmpCov modifies the prior over weights. preventing the of publically available samples from BNNs with standard Gaussian priors (Izmailov et al., 2021b), which is especially important as some gold-standard Bayesian sampling methods are extremely expensive (e.g- by Hamiltonian Monte Carlo, HMC in Izmailov et al. (2021b) took one hour to get sample on a ResNet trained on CIFAR-16 10 using "a cluster of 512 TPUs" Izmailov al., 2021b). In contrast, ShiftMatch keeps the prior and training-time likelihood unchanged, allowing us to directly re-use gold-standard HMC samples from Izmailov et al. (2021b). Indeed, ShiftMatch is highly efficient does not require further retraining or fine-tuning at test time, allowing us to e.g. use very large batch size.
likelihood.
We found that ShiftMatch improved considerably over all previous baselines for BNN robustness to corruption, including BNNs with standard Gaussian and EmpCov priors, and non non-Bayesian methods like plain deep ensembles (Fig. 1,3). Further, ShiftMatch can combined with non-Bayesian methods: we found significant improvement for all methods tested: stochastic gradient descent (SGD), ensembles, and Bayes (HMC) (Fig. 4), and ShiftMatch can be scaled to ImageNet (Fig. 5), where it offers improved performance over test-time batchnorm (Nado et al., 2020).
2 BACKGROUND
In BNNs, we use Bayes theorem to compute posterior distribution over the weights w given the training input, X and labels, Yusin,
P (Ytrain train X train w) W)P(W)
P (w) Xarinn Virain)
P (y train X
2

---

Published as a conference paper at ICLR 2023
At test-time, we use Bayesian model averaging (BMA) to compute predictive distribution test labels, Ylest: given test inputs, Xust, by integrating over the posterior,
/ P (Yrsst X w) P (w Xmain Yerain) dw.
P (Yrett Xrests Xunin Strain)
(2)
In practice, directly using Bayes theorem is intractable in all but very low dimensional settings. As such, alternative approaches, ranging from (HMC) (Neal et al., 2011; Izmailov et al., 2021b), variational inference (Graves, 2011; Blundell al., 2015), stochastic gradient Langevin dynamics (SGLD) (Welling & Teh, 2011; Zhang et al., 2019) or expectation propagation (EP) (Hernández- Lobato Adams, 2015) are usually used. HMC is usually regarded as the "gold-standard", as it will produce exact samples given enough time to sample, while the other methods here will not.
3 RELATED WORK
The closest prior work is Izmailov et al. (2021a), which introduces the "EmpCov" prior that uses the training inputs, X X E IR Nem x M, to specify the prior over weights,
PempCoov (w mrain) = N/ (w: 0, 1. ( (Niann
(~21 X X X
(3)
Here, Nuaan is total number of training examples, M is the number of features in the input layer, and w E R is weight vector mapping from the inputs to single feature in the first hidden layer. The covariance of the Gaussian prior over weights, MNris covariance of the training inputs. Thus, the prior suppresses weights along directions with low input variance. Interestingly, the idea of using the inputs to specify prior over weights arises earlier, in the Zellner G-prior (Zellner, 1986). Remarkably, Izmailov et al. (2021a) and Zellner (1986) propose priors that are in some sense "opposite": while the EmpCov prior suppresses weights in directions with low input variance, the Zellner-G prior magnifies weights in those directions, as to ensure that all directions in input space have the same impact on the outputs.
1 X X X is proportional to the empirical
Pzellner-G (w|X (w(Xtain) -V[ w: O, 3I
Xtrain Xtrain
Of course, the best prior depends on the setting: we suspect that the Zellner-G prior is most appropriate when informative features have very different scales, whereas the EmpCov prior is most appropriate for e.g. images, where the low-variance input directions are likely to be uninformative.
One approach to using unlabelled data to improve robustness to corruption to keep batchnorm on at test-time, so the batch statistics are computed on corrupted test data (Li et al., 2016; Carlucci et al., 2017; Nado et al., 2020; Schneider al., 2020; Benz et al., 2021). There considerable confusion about whether batchnorm can be applied to BNNs, with recent papers arguing that batchnorm cannot be formulated as a likelihood that factorises across datapoints (Izmailov et al., 202 bb; D' Angelo & Fortuin, 2021: Noci et al., 2021). However, argue that valid likelihoods do not need to factorise across datatpoints. We introduce such a likelihood for BNNs, which behaves like applying batchnorm during training and keeping batchnorm test-time batchnorm gives huge performance improvements in BNNs of 21.99 21.9% for CIFAR- 10-C for the highest corruption level, which is more than twice the performance improvement for ensembles of 10.6%. We then go on to develop ShiftMatch which gives even larger improvements of 24.6% for BNNs and 12.4% for ensembles. While ShiftMatch bears some superficial similarities to test-time batchnorm, it is in reality quite different. In particular, on convolutional neural networks, batchnorm matches only the means and variances for each channel. In contrast, ShiftMatch matches the spatial covariances statistics for each channel separately, allowing us to correct complex, spatially extended corruptions, and to perfectly correct certain simpler corruptions such as Gaussian blur (Sec. 4.6). We find ShiftMatch gives considerable performance improvements over test-time batchnorm (Figs. 4,9).
test-time. Surprisingly, we find that just
There is work suggesting whitening or transforming the feature correlations (Huang et al., 2018), in part to improve corruption robustness (Roy et al., 2019; Choi et al., 2021). These approaches have at least three key differences from ShiftMatch. First, these approaches whiten at train-time. which dramatically slows down training, and prevents the use of pre-trained networks, regardless of being Bayesian or not. In contrast, ShiftMatch training is exactly equivalent to faster standard NB)NN training, allowing the use of pre-trained models. Second, these papers apply images with
3

---

Published as a conference paper at ICLR 2023
C' channels, width W and height H, but whiten only over channels (i.e. using C C matrices). However, this ignores the complex spatial structure present in most corruptions: we find that by transforming spatial correlations we can perfectly remove simple corruptions such Gaussian blur (Sec. 4.6). As such, we instead transform the spatial correlations of each channel separately (see Fig. for variants). Importantly, these two choices interact: it probably is not possible work with spatial correlations at training time, because at training time, we only have a few samples in the minibatch (e.g. Nmini 128) estimate high-dimensional spatial covariances (e.g. for MNIST, these covariance matrices would be HW x HW = 784 x 784). ShiftMatch largely avoids this issue because it computes the spatial correlations once at the end, using a large amount of training and testing data, giving extremely accurate covariance estimates (see Sec. 4.4 for further details).
There is also a line of work (e.g. Sun et al., 2020; Wang al., 2020; Zhang et al., 2022) aiming at improving the robustness of pre-trained model by performing self-supervised training at test time. This approach likely be too computationally expensive our setting, as we need to evaluate on many different HMC samples/ensemble members.
Trinh et al. (2022) propose a method for inducing corruption robustness based on a Bayesian interpretation of dropout (Gal & Ghahramani, 2016). This method ís very different, in that it proposes a modified training objective, and therefore cannot use gold standard HMC samples.
METHODS
4.1 REINTERPRETING EMPCOV AS A TRAINING- -EE---DATA---EEEEEE DEPENDEEN LIKELIHOOD
The EmpCov prior over weights (Eq. 3) induces a distribution over f = Xw E IR,, the value of single output feature for all P inputs (see Appendix A for a full derivation),
PmmpCov (f] = N (f; (f:0, M M.
Xeest
Irain
Maiin)
Jest
(5)
We obtain an equivalent distribution over f by first transforming the inputs using ShiftEmpCov,
1/2
ShiftEmpCov(X X Xust
train
train
In the arguments to Shiftmppov, the "input", X is to the left of the semicolon, and the quan- tities used purely for normalization, Xarain are the right of the semicolon. Computing f from these transformed inputs, f = ShiftEEmpCov(X.a. lest; Xmain) Wr, and using a standard, isotropic prior, P (w) = A' (0, 1),, recovers the EmpCov distribution over f (Eq. 5). (For details of the full equivalence using the function-space viewpoint, Appendix B). Importantly, ShiftEmpCov moves the explicit training- maining-data-depeen dependeee from the prior to the likelihood,
P EmpittEmpCov (Yuaan trann w) Categorical (y wain TShuftEmpCov X (Xurin w)),
(7a)
P EmhootEmpoov (y vesa X W, X = Categorically TShiitEmpCur (Xiest: w)). (7b)
Here, Empp Xuni is a function parameterised by a neural network with ShiftEmpCov at the in- put layer. Critically, this function (in particular, the normalization) depends X giving a training- data dependent likelihood. We can then apply the function to arbitrary input data, (which could be training, X Ximins test, X or something else), and weights, w, giving Tihhin (-, w). Importantly, we are using exactly the same function, Eshi likelihoods: the only unusual thing is that the function itself (i.e. the likelihood) depends on the training data, which is analogous to how the prior depended on the training data in the original EmpCov setup.
w) for the train and test
4.2 DEVELOPING IMPROVED TRAINING- AAAINING-DATA-AA DAAA--AATAEEEEEEA DEPENDENT LIKELIHOODS
Motivated by EmpCov's equivalent training-datte data-dependent likelihoods and priors, considered other
likelihoods that might resolve issues with EmpCov priors. In particular, we designed ShiftMatch as a maining-data-dependee likelihood that is capable of reusing samples from pre-trained networks. The ShiftMatch likelihood is built upon the ShiftMatch transformation
1/2
1/2
ShiftMatch(H H Hast
Hestt)
Hurin
+ M
(8)
Naa
4

---

Published as a conference paper at ICLR 2023
where Hest and Harin mean-subtracted versions of the test and training inputs/features, Muain is the training mean, and we use the symmetric matrix square root, (Huang et al., 2018). ShiftMatch matches the mean and covariance of the test data to that of the training data; for the covariances,
Nest (ShiftMatch Hest H gran) -Maain)
(ShiftMatch) (Hest
H -
H Harin
Nusia
As such, applying ShiftMatch to the training data leaves it unchanged,
ShiftMatch( (aaan H main))
H main + M H
(10)
allowing us to train the network without including the ShiftMatch layers. As the priors over weights also the same isotropic Gaussian priors used in standard BNNs, we are able to use e.g. pre-trained gold-standard samples from HMC-trained BNNs (Izmailov et al., 2021b). A summary of the full ShiftMatch algorithm be found in Appendix E Alg. Overall, ShiftMatch likelihoods have the same basic form as the ShiftEmpCov likelihoods (Eq. 7),
MaiftMatch (Y train Xaain W) Categoricallyunnn uraan MatchtMt (Xurains w))
(1la)
MshiftMatch (Frst X W, Xuain) = Categorical(yusst MattitMatch: (Xrest, w))
(11b)
The key difference is that for the training data, the likelihoods are equal to those in standard BNNs, Mathh (Xrrain w) = 7T (Xtrain: w)
(12a)
Match (y train X train: w) P (y Irain X W)
(12b)
where standard BNN likelihoods and NNs are denoted by P and π without any subscript. This implies that ShiftMatch posteriors are exactly equal to posteriors in a BNN with standard likelihoods, allowing us to use samples of the weights from pre-trained networks,
Match (w) Yuain Xtrain ) X PshiftMatch (Yuain X w) P. (W)
P (Ytrain train X w) P (w) x P (w)ye Vurain Xurin)
(13)
4.3 STRUCTURED COVARIANCE ESTIMATION IN CNNS
The ShiftMatch operation shown in Eq. (8) is appropriate for fully-connected networks with relatively small number of features. However, this approach cannot be applied directly CNNs, because the feature maps are of size CHWW, where C is the number of channels, H is the height and W is the width. Thus, we end up with huge CHW CHW covariance matrices, which presents computational difficulties (e.g. the matrix square roots require O H compute).
Instead, we need to structure our covariance estimates to improve accuracy and speed up matrix computations. We begin by choosing to match the spatial covariance for each channel separately. which would require us estimate C separate HW x HW covariance matrices. Even then, this poses difficulties because images or feature maps can have very large numbers of pixels. Thus, we need to impose further spatial structure, and we choose to treat the covariance along the height and width separately, using Kronecker factored covariance (Martens & Grosse, 2015),
end up with a x H matrix and a W x W matrix for each channel (see Appendix Fig. 7 for other variants).
4.4 COMPARISON WITH BATCHNORM
ShiftMatch does bear some minor resemblance batchnorm, in that matches the feature moments to prespecified values. Indeed, inspired by ShiftMatch, we could develop a "spatial batchnorm", a generalisation of batchnorm that matches spatial covariances of feature maps to learned values. However, even spatial batchnorm would ultimately be quite different from ShiftMatch. In particular, batchnorm is usually applied only at train-time. and turned off test-time (though is sometimes left on at test-time, where it can improve robustness to corruption Nado et al., 2020; Schneider et al., 2020). However, ShiftMatch is the opposite in that is applied only at test-time; at train-time, ShiftMatch is identical to standard NN without extra normalization layers (Eq. 10). Thus, ShiftMatch allows us to decouple test time from train time normalization, which turns out to be a powerful idea. In particular, ShiftMatch has two key advantages covariance matrices are complex, high- dige--imnnsional objects, which require much more than a single minibatch to estimate accurately. With ShiftMatch, it is straightforward to aggregate minibatches,
a spatial batchnorm. First, the required spatial
5

---

Published as a conference paper at ICLR 2023
as we estimate these matrices after training the weights using any available training and test data. Moreover, as we compute the covariances after training, we can use large batchsizes (e.g. we can fit batch of 1000 for ImageNet even on single 2080ti with 11GB of memory), because we do not need to retain feature maps at all layers for use in the backward pass (for further discussion, see e.g. Sander et al., 2021). In contrast, training batches used in spatial batchnorm would need to be much smaller (often 16 examples for ImageNet on smaller GPUs), which can give very inaccurate or even low-rank, non-invertible estimates of the covariance matrices. Second, ShiftMatch trains much faster than a spatial batchnorm. In particular, even after carefully simplifying the covariance matrices, computing the matrix square roots every layer at every gradient step is still very slow. In contrast, in ShiftMatch, we only do the matrix square roots once after training. For instance, it took us around 0.35s for a forward pass without spatial batchnorm, which contrasts with 0.77s for a forward pass with spatial batchnorm using a mini-batch of 128 inputs for CIFAR-10.
4.5 TRAINING DATA AT TEST TIME?
One potential disadvantage of ShiftMatch is the requirement for training data to compute moments, which is sometimes not provided for pre- pre-trained models. However, this requirement is mitigated by three factors. First, if ShiftMatch becomes commonplace, then the necessary training statistics could be released along with model weights. This is in principle exactly the same as capturing the training means and variances in batchnorm layers. Second, you do not need to use the whole training set. You only need enough datapoints to get reasonable estimate of the covariance matrices (see Fig. 5, where performance saturates after using only 10,000 training images, less than I 1% of the full dataset). Third, we may not even need training inputs, as long as we can get proxy inputs that are close enough. Fig. 1, shows only a small performance decrease when we do not use the training dataset to obtain the training statisties, but instead we use the entirely separate CIFAR 10.1 captured by Recht et al. (2019).
4.6 SHIFTMATCH PERFECTLY REMOVES STATIONARY LINEAR CORRUPTIONS
ShiftMatch retains all the useful theoretical properties of EmpCov priors, in suppressing patholo- gies that arise in BNNs when there is low-input variance along a direction in the training but not in the corrupted test data (see Sec. 1 and Izmailov et al., 2021a for more details). ShiftMatch also has considerable theoretical advantages over EmpCov priors, in that certain simpler corruptions can be removed perfectly (Fig. 2). In particular, Shift- Match can perfectly remove linear spatially stationary corruptions such as Gaussian blur (in the limit of in- finitely many samples). Without loss of generality, con- sider X E R Muii x HW as zero-mean training images with a single input channel (e.g. a grayscale image). We have uncorrupted test data, Xuncor E N HW which is corrupted with a stationary linear spatial operation, repre- sented by S I (rWWHH!
Original
Gaussian Blur
Post ShirtMatch
Figure 2: An example of ShiftMatch ap- plied on Gaussian blur corrupted image.
X u E R Mea
(14)
Taking the SVD of Xirain Xuncur and S, we get,
X Irain
U Uirin DV
Xuncor = DV
S VDVD.VT
(15)
Here, we have used the fact that the right singular vectors of stationary signals such as images is Fourier basis, V. and D is the power-spectrum of uncorrupted images (Olshausen & Field, 1996; Hernandez et al., 2018), which is the same for the uncorrupted training data, Xunin and uncorrupted
test data, X Xuncor- Further, the eigenbasis of stationary transformations such as Gaussian blur again the Fourier basis, V. Thus, the SVD of the corrupted test data is,
X = uncorr S U VD, V = UuS DD.
(16)
Critically, applying ShiftMatch to the corrupted test data exactly recovers the uncorrupted test data,
1/2
1/2
Shift (Xoori X = X
Xeor Xcor)
X X
(17)
Nusi
= Una (DD,) V (V(DD.)-vv) (VDVT)- = U DV
= -
6

---

Published as a conference paper at ICLR 2023
1.0
Accuracy
0.8
0.6
0.4
Log re-llkkllhoodd
Identity
Shot Noise
Impulse Noise
Glass Blur
Motion Blur
Shear
Scale
Rotate
Brightness
Translate
Stripe
Fog
Spatter
Dotted Line
Zigrag
Canny Edges
Avg.
HMC
EmpCov
Deep Ens
ShiftMatch
Figure 3: Performance of HMC with isotropic priors (HMC) and EmpCov priors (EmpCov), deep ensembles and ShiftMatch on MNIST-C under 15 different corruption types. The left-most column gives the performance on clean MNIST while the right-most column gives the performance averaged across the 15 corruptions. ShiftMatch performs better overall, but uses additional information from minibatch corrupted test datapoints, which not used by the other methods.
5 EXPERIMENTS
First, we applied ShiftMatch to the HMC samples from Izmailov et al. (2021b) for large-scale Bayesian ResNet trained on CIFAR-10 and tested on CIFAR-IO-C (Hendrycks & Dietterich, 2019), and compare to the baselines in Izmailov et al. (2021b). ShiftMatch gives considerable improvement over using the raw HMC samples and over all baselines (including e.g. ensembles). Second, we show that ShiftMatch performs better than EmpCov priors on small CNNs from Izmailov et al. (2021a) trained on MNIST and tested on MNIST-C(Mu ShiftMatch improves over test-time batchnorm when we use HMC samples from Izmailov et al. (2021b), and when we consider non- -on-aaysiann baselines (SGD and ensembles; Fig. 4). Fourth, we show that ShiftMatch can be applied on a large pre-trained non-Bayesian network, where it improved performance on ImageNet relative to test-time batchnorm. Finally, we evaluate the performance of ShiftMatch
Gilmer, 2019). Third, we find that
OOD detection task and find it has no obvious beneficial or harmful effects.
CIEAR-10-C baselines from Izmailov et al. (2021b). We began by applying ShiftMatch to the gold-standard HMC samples from BNN trained CIFAR-10 from Izmailov al. (2021b), and comparing to the baselines from that paper (Fig. 1, Appendix C, Fig. 6). They used a ResNet-20 with only 40,960 of the 50,000 training samples (in order "evenly share the data across the TPU devices"), and to ensure deterministic likelihood evaluations (which is necessary for HMC), turned off data augmentation and data subsampling (i.e. full batch training), and used filter response normalization (FRN) (Singh & Krishnan, 2020) rather than batch normalization (loffe & Szegedy, 2015). Despite these stringent constraints, their network achieved excellent performance, with greater than 90% accuracy on the clean test set. For further details of how these baselines were generated, please see Izmailov et al. (2021b). We found that ShiftMatch considerably improved over baselines included in that paper, including the raw HMC samples themselves, and deep ensembles.
MNIST-C baselines from Izmailov et al. (2021a). To compare against EmpCov priors, we were forced to use smaller-scale networks and datasets, because that is all that is provided in the EmpCov paper (Izmailov et al., 2021a). Presumably, this was because they wanted gold gold-standard HMC samples, but longer had access the large scale compute in Izmailov et al. (2021b). In particular, they used smaller models such as multilayer perceptron and LeNet-5 (LeCun et al., 1998), which contains 2 convolutional layers followed by fully-connected layers. We used the HMC samples for
7

---

Published as a conference paper at ICLR 2023
0.95
0.90
0.85
Log-Liken like
0.80
Accuracy
0.75
0.70
-15
0.65
0.60
-2.0
0.55
0.50
5
Corruption Intensity
Corruption Intensity
SGD
SGD w/ BatchNorm
SGD w/ ShirtMatch
Deep Ens
Deep Ens w/ BatchNorm
Deep Ens w/ ShiftMatch
HMC
HMC w/ BatchNorm
HMC w/ ShiftMatch
Figure 4: All combinations of SGD (red, left three bars), ensembles (blue, middle three bars) and HMC (green, right three bars) with no normalization, batchnorm and ShiftMatch for CIFAR-10-C. Note that ShiftMatch and test-time batchnorm use information from minibatch of corrupted inputs, while the plain methods do not.
LeNet trained on MNIST under standard Gaussian prior and the EmpCov prior provided by Izmailov al., 2021a to evaluate different methods on MNIST-C (Mu & Gilmer, 2019). Additionally, we used their code to get an ensemble of 50 independently initialised models trained using SGD with momentum as the combined with cosine-annealle we found that ShiftMatch performed better than the other methods. including HMC with standard Gaussian and EmpCov priors, and non-Bayesian baselines such as ensembles (Fig. 3). Note the poor performance of ShiftMatch on translate, which can be understood at least in part because translations are not stationary transformations that can be exactly reversed by a ShiftMatch (Sec. 4.6). That said, robustness to translations is usually induced by data augmentation which was not used when Izmailov et al. (2021a) were training their networks.
learning rate decay and batchsize of 80. Again,
Test-time batchnorm with BNNs, and ShiftMatch with mon-Bayesian methods. Above, focused on combining ShiftMatch with BNNs, and we found that these models performed better than previous approaches used in the BNN literature (Izmailov al., 2021 However, ShiftMatch could also be applied to non- Ban-Baysiian networks, and test-time batchnorm could be applied to BNNs. We therefore considered all possible combinations of ShiftMatch and batchnorm with SGD, ensembles and HMC (Fig. 4). To ensure a fair comparison, we used networks, HMC samples, training code and experimental settings from Izmailov et al. (2021b). To be more specific, for SGD, we used network architecture identical to that of HMC: A ResNet-20 with batchnorm replaced by FRN. The ensemble model was constructed using 50 SGD models with different initialisations. As the underlying network did not have batchnorm layers, we used a variant of ShiftMatch to match the mean and the variance of each channel with the training sample statisties (see Fig. 9 for full batchnorm).
In all cases we found that test-time batchnorm improved over plain models, and ShiftMatch improved over test-time batchnorm. Confirming results from Izmailov et al. (2021b), we found that HMC without either test-time batchnorm or ShiftMatch performed very poorly: worse than ensembles for any level corruption. In contrast, with ShiftMatch, HMC and deep ensembles performed very similarly. Indeed, ShiftMatch gave spectacular improvements in accuracy of 24.6% for HMC, and smaller improvements of 12.4% for ensembles (Appendix Tab. 2).
Scaling to ImageNet with non-Bayesian ShiftMatch ShiftMatch can be applied to Bayesian or non-Bayesian pre-trained networks, and can readily be applied to checkpoints from larger pre-trained models. To demonstrate this, took an ImageNet checkpoint from Keras Applications and applied ShiftMatch. Note that this network included batchnorm in training; to combine ShiftMatch with
hhttps://keras wttpsi//aaapaaaaaa i i cns/
8

---

Published as a conference paper at ICLR 2023
0.7
Log. L--kkeihhood
-1
Accuracy
-2
-3
0.3
0.2
4
5
5
Corruption Intensity
Corruption Intensity
ShiftMatch ShiftMatch 10K
ShiftMatch
BatchNorm
ShiftMatch 1,280K training
Performance of plain SGD, test-time BatchNorm and ShiftMatch on ImageNet-C (Hendrycks & Dietterich, 2019). The number after ShiftMatch indicates the number of training points (out of the 1, 280K training samples) used for estimating the training statistics. All results presented pre-trdined ResNet-50 the backbone.
Figure 5:
Table 1: OOD detection performance of models trained CIFAR-10, taking CIFAR-100 and SVHN as OOD datasets. Note the arrows indicate the "better" direction (e.g. so lower FPR95 is better).
AUROC f
FPR95
OOD dataset SGD
Deep Ens HMC ShiftMatch
SGD
Deep Ens
HMC
ShiftMatch
CIFAR-100
0.780
0.830
0.851
0.848
0.579
0.490
0.457
0.480
SVHN
0.713
0.788
0.874
0.920
0.625
0.453
0.299
0.281
batchnorm, we added ShiftMatch after the batchnorm layer (see Appendix F for further details). We found that ShiftMatch again offers considerable gains over just using batchnorm, especially at the higher corruption intensities (Fig. 5). The test statistics computed using batch of 1,000 images. The training statistics are computed using a random subset of 1,000, 10,000 or 100,000 training images, the full dataset. We slight reduction in performance when using only 1,000 training images, but no differences otherwise.
Effect of ShiftMatch on uncertainty quantification While BNNs have poor performance under corruption, they still provide high-quality uncertainty estimates that can benefit down-stream tasks, such as detecting OOD inputs (Wang & Aitchison, 2021; Kristiadi et al., 2022). We might think that ShiftMatch would be detrimental to OOD detection accuracy, as it would make OOD samples more similar training samples. We consider models trained on CIFAR-10, 10, and asked whether we could identify inputs from CIFAR-100 or SVHN. We used SGD, ensembles and HMC with and without ShiftMatch. For all the models, we used the Izmailov al. (2021b) setting as in Figs. 1-4. We used the maximum softmax probability (Hendrycks & Gimpel, 2017) of the (ensembled) predictive distribution as the score and we used AUROC and FPR95 as the performance metric. found that HMC and ShiftMatch always perform better than ensembles and SGD (Tab. 1). The performance of HMC with and without ShiftMatch was very similar; HMC without ShiftMatch performed slightly better on CIFAR-100, and HMC with ShiftMatch performed slightly better on SVHN. These experiments show that any detrimental effect of ShiftMatch on OOD detection performance is likely to be small.
6 CONCLUSION
We developed new
likelihood for robustness to corruption for BNNs, called ShiftMatch. ShiftMatch has the key limitation that, like keeping batchnorm at test-time but unlike EmpCov priors, it uses information from minibatch of similarly corrupted test points. We found that ShiftMatch indeed offered considerable improvements over past approaches to robustness to corruption in BNNs, both in CIFAR- 10-C (Fig. 1) and in MNIST-C (Fig. 3), and even can give improvements over rest-time batchnorm when used in non- -ayeess settings (Figs. 4.9).

---

Published as a conference paper at ICLR 2023
REFERENCES
Laurence Aitchison. Bayesian filtering unifies adaptive and non-adaptive neural network optimization methods. Advances in Neural Information Processing Systems, 33:18173-18182, 2020a.
Laurence Aitchison. A statistical theory of cold posteriors in deep neural networks. arXiv preprini arXiv:2008.059/1 059/2, 2020b.
Philipp Benz, Chaoning Zhang, Adil Karjauv, and In So Kweon. Revisiting batch normalization for improving corruption robustness. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 494-503, 2021.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International conference on machine learning, pp. 1613-1622. PMLR, 2015.
Mariusz Bojarski, Davide Del Testa, Daniel Dworakowski, Bernhard Firner, Beat Flepp, Prasoon Goyal, Lawrence Jackel, Mathew Monfort, Urs Muller, Jiakai Zhang, et End end learning for self-driving cars. arXiv preprint 160xv:6004.777 2016.
Fabio Maria Carlucci, Lorenzo Porzi, Barbara Caputo, Elisa Ricci, and Samuel Rota Bulo. Autodial: Automatic domain alignment layers. 20/7 IEEE international conference on computer vision (ICCV), pp. 5077-5085. IEEE, 2017.
Sungha Choi, Sanghun Jung, Huiwon Yun, Joanne T Kim, Seungryong Kim, and Jacgul Choo. Robustnet: Improving domain generalization in urban-scene segmentation via instance selec- tive whitening. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1580-11590, 2021.
Francesco D' Angelo and Vincent Fortuin. Repulsive deep ensembles are bayesian. In A. Beygelzimer, Y. Dauphin, Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https: // operrevieeeeeeeeeeeeeee net. fo rum? d=LAKp 1pl.Ap 8.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and bayesian deep learning. Advances in Neural Information
Philipp Hennig. Laplace Processing Systems, 34, 2021.
Andre Esteva, Brett Kuprel, Roberto A Novoa, Justin Ko, Susan M Swetter, Helen M Blau, and Sebastian Thrun. Dermatologist-level classification of skin cancer with deep neural networks. nature, 542(7639):115-118, 2017.
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pp. 1050-1059.
PMLR, 2016.
Adrià Garriga-Alonso, Carl Edward Rasmussen, and Laurence Aitchison. Deep convolutional networks shallow gaussian processes. arXiv preprint arXiv:1808.05588 2018
Alex Graves. Practical variational inference for neural networks. Advances neural information processing systems, 24, 2011.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. arXiv preprint arXiv: /903. /90..//21../1111 2019.
Dan Hendrycks and Kevin Gimpel. A baseline for detecting misclassified and out-of-diitribuation ex- amples neural networks. Proceedings of International Conference on Learning Representations. 2017.
Wilmar Hernandez, Alfredo Mendez, and Türkmen Göksel. Application of principal component analysis to image compression. Statistics- Grriiin Data Sets and Growing Demand for Statistics, 2018.
José Miguel Hemnáneez-Lobtto and Ryan Adams. Probabilistic backpropagation for scalable learning of bayesian neural networks. In International conference on machine learning, pp. 1861-1869 PMLR, 2015.
10

---

Published as a conference paper at ICLR 2023
Lei Huang, Dawei Yang, Bo Lang, and Jia Deng. Decorrelated batch normalization. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 791-800, 2018.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448-456. PMLR, 2015.
Pavel Izmailov, Patrick Nicholson, Sanae Lotfi, and Andrew G Wilson. Dangers of bayesian model averaging under covariate shift. Advances in Neural Information Processing Systems, 34, 202 l0.
Pavel Izmailov, Sharad Vikram, Matthew D) Hoffman, and Andrew Gordon Gordon Wilson. What are bayesian neural network posteriors really like? In International Conference on Machine Learning, pp. 4629-4640. PMLR, 2021b.
Mohammad Emtiyaz Khan and Havard Rue. arXiv:2107.04562, 2021.
The bayesian learning rule.
arXiv preprint
Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. Being a bit frequentist improves bayesian neural networks. In International Conference Artificial Intelligence and Statistics, pp. 529-545. PMLR, 2022.
Balaji Lakshminarayanan. Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 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-2324 1998.
Yanghao Li, Naiyan Wang, Jianping Shi, Jiaying Liu, and Xiaodi Hou. Revisiting batch normalization for practical domain adaptation. arXiv preprint arXiv: 160.0079, 2016.
David JC MacKay. A practical bayesian framework for backpropagation networks. Neural computa- fion, 4(3):448 472,, 1992.
Wesley Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainty in deep learning. Advances in Neural Information Processing Systems, 32, 2019.
James Martens and Roger Grosse. Optimizing neural networks with frotecker-factoreed approximate curvature. In International conference on machine learning, pp. 2408-2417. PMLR, 2015.
Norman Mu and Justin Gilmer. MNIST-C: A robustness benchmark for computer vision. arXiv
preprint arXiv: 02337, 2019.
Zachary Nado, Shreyas Padhy, D Sculley, Alexander D' Amour, Balaji Lakshminarayanan, and Jasper Snoek. Evaluating prediction-time time batch normalization for robustness under covariate shift. arXiv preprint arXiv:2006. 10963, 2020.
Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
Radford M Neal et al. Meme using hamiltonian dynamics. Handbook of markov chain monte carlo, 2(11):2, 2011.
Jeremy Nixon, Michael W Dusenberry, Linchuan Zhang, Ghassen Jerfel, and Dustin Tran. Measuring calibration in deep learning. In CVPR Workshops, 2019.
Lorenzo Noci, Kevin Roth, Gregor Bachmann, Sebastian Nowozin, and Thomas Hofmann. Disentan- gling the roles of curation, augaentnmtntttionn and the prior the cold posterior effect. Advances in Neural Information Processing Systems, 34:12738-12777 2021.
Sebastian W Ober and Laurence Aitchison. Global inducing point variational posteriors for bayesian neural networks and deep gaussian processes. In International Conference on Machine Learning, pp. -259. PMLR, 2021.
11

---

Published as a conference paper at ICLR 2023
Bruno A Olshausen and David J Field. Emergence of simple-cell receptive field properties by learning a sparse code for natural images. Nature, (616653) 1996.
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, David Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshninarayanan, and Jasper Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. Advances in neural information processing systems, 32, 2019.
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize imagenet? In International Conference on Machine Learning, pp. 5389-5400. PMLR, 2019.
Subhankar Roy, Aliaksandr Siarohin, Enver Sangineto, Samuel Rota Bulo, Nicu Sebe, and Elisa Ricci. Unsupervised domain adaptation using feature-whitening and consensus loss. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9471 2019.
Tim GJ Rudner, Zonghao Chen, Yee Whye Teh, and Yarin Gal. Tractable function-space variational inference in bayesian neural networks. In JCML 202/ Workshop on Uncertainty and Robustness in Deep Learning, 2021.
Michael E Sander, Pierre Ablin, Mathieu Blondel, and Gabriel Peyré. Momentum residual neural networks. In International Conference on Machine Learning, pp. 9276-9287. PMLR, 2021.
Steffen Schneider, Evgenia Rusak, Luisa Eck, Oliver Bringmann, Wieland Brendel, and Matthias Bethge. Improving robustness against common corruptions by covariate shift adaptation. Advances in Neural Information Processing Systems, 33:1 1539-11551, 2020.
Saurabh Singh and Shankar Krishnan. Filter response normalization layer: Eliminating batch dependence in the training of deep neural networks. Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 11237-11246, 2020.
Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei Efros, and Moritz Hardt. Test-time training with supersupn for generalization under distribution shifts. In International Conference on Machine Learning, pp. 9229-9248 PMLR, 2020.
Trung Trinh, Markus Heinonen, Luigi Acerbi, and Samuel Kaski. Tackling covariate shift with node-based bayesian neural networks. In International Conference on Machine Learning, pp. 21751-21775. PMLR, 2022.
Ali Unlu and Laurence Aitchison. Gradient regularization as approximate variational inference. Entropy, 23(12):1629) 2021.
Dequan Wang, Evan Shelhamer, Shaoteng Liu, Bruno Olshausen, and Trevor Darrell. Tent: Fully test- time adaptation by entropy minimization. International Conference on Learning Representations, 2020.
Xi Wang and Laurence Aitchison. Bayesian ood detection with alcatoric uncertainty and outlier exposure. In Fourth Symposium on Advances in Approximate Bayesian Inference, 2021.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pp. 681 681-688. Citeseer, 2011.
Arnold Zellner. On assessing prior distributions and bayesian regression analysis with g-prior distributions. Bayesian inference and decision techniques, 1986.
Marvin Mengxin Zhang, Sergey Levine, and Chelsea Finn. Test time robustification of deep models via adaptation and augmentation, 2022. URL https: ptttps://openrevviii netpi/ewwweeeeeeeeeeeeeeeeeeeeee www.morrrmmieeeee J1 D1U0gffbb.. bP.
Ruqi Zhang, Chunyuan Li, Jianyi Zhang, Changyou Chen, and Andrew Gordon Wilson. Cyclical stochastic gradient meme for bayesian deep learning. arXiv preprint 1rXiv:1902.033932 2019.
12

---

Published as a conference paper at ICLR 2023
A DERIVATION OF THE EMPCOV DISTRIBUTION OVER FEATURES
We have f = Xest w. with the EmpCov prior over weights, Eq. (3), Our goal is to compute, the implied distribution, P (eets As W is a linear combination Gaussian w with fixed Xest, we know P (es) that is multivariate Gaussian,
P (f]X)))t N (f; pt, 2).
(18)
Therefore to characterise the full distribution, P X),, we just need to compute the expectation, μ., and covariance, E. The mean is zero,
[[st X E [
(19)
as E [w] (Eq. 3). And the covariance is, www
Σ
Xeett E
X test
Xestt
(X train rrain)
X test
(20)
C.
T as E. [w
Irain-
Substituting la (Eq. 19) and S (Fq. 20) into Eq. (18) gives Eq. (5).
CNumis
train
B EQUIVALENCE OF POSTERIORS UNDER EMPCOV PRIORS AND SHIFTEMPCOV LIKELIHOODS
The key claim in Sec. 4.1 is that models with EmpCov priors and ShiftEmpCov likelihoods have equivalent behaviour, because the distribution of the outputs of a layer with an EmpCov prior and ShiftEmpCov likelihood are the same.
To formalise this notion of equivalence, we need to look function-space viewpoint (e.g. Rudner et al., 2021). This viewpoint works with the distribution over the logits (or neural network outputs) L. The logits, L, depend on the inputs, X and the neural network weights, w. Of course, L is usually deterministic function of X and w. The function space viewpoint (e.g. [1] induces distribution over functions (i.e. logits or network outputs, L) by marginalising the prior over weights,
f dd
P (L[X)= =
(21)
While explicitly computing P XL is usually intractable, this viewpoint very useful for establish- ing theoretical properties. In particular, as the class labels depend only on the logits/network outputs, we can write the full function-space graphical model as,
X y.
(22)
Here, we combine the training and test points, X (X X1es), L (Lurain; Lust) and y (Yuain= Fes).. The prior over logits, P (L is the same in both EmpCov and ShiftEmpCov viewpoints, because the prior distribution of the outputs of each layer is the same. Likewise, the softmax likelihood P (y [L) is again the same in both EmpCov and ShiftEmppov viewpoints. As the priors are exactly equivalent under the two viewpoints, the posteriors over L and Yies must also be equivalent, and we can go further and explicitly write out the equivalent posteriors. The posterior over train and test logits, L, conditioned on all inputs, X and training labels, Ytrain is
P (L|X, Vtrain) x P (Yurain [main) P PL)
(23)
Remember that L = (Turain Lues) and X = (X Xlea). Again, this posterior is the same for the EmpCov and ShiftEmpCov viewpoints, because the prior, P (L(X), is the same, and because the softmax likelihood, P (ytrain Lmrin)) is the same. From the posterior over logits, we can obtain the predictive distribution over test labels,
P (Yreal X, Ytrain) = faL AL PL Viran).
(24)
and this is again equivalent because the posterior over logits is equivalent, and the softmax likelihood is equivalent.
13

---

Published as a conference paper at ICLR 2023
Frosted
Accuracy Logiketheed
-9.5
-1.3
-1.3
Log ikalheed
Intonsity
Intonstty
Intorsity
Intonsity
Intersity
Intensity
Intonstry
Figure 6: Performance of HMC, SGD, Deep Ensemble, SGLD, Mean-field VI (from dark red to yellow) and ShiftMatch (blue) on LEARR-0-CC under 16 different corruption types with intensity levels.
C FURTHER EXPERIMENTAL RESULTS
Here, we present full performance across different corruptions for HMC on CIFAR-10-C and other baselines from (Izmailov et al., 2021b) (Fig. 6). Tab. 2 gives performance differences with and without ShiftMatch, and Tab. 3 gives performance differences with and without test-time batchnorm, for the usual Izmailov et al. (2021b) setting. Note that HMC benefits far more than non-Bayesian methods, with improvements often being
than double those for ensembles.
ShiftMatch did take slightly more time: 1.85s for a minibatch of 10,000 points as opposed to 1.15s without ShiftMatch. This tradeoff is likely acceptable due the very large improvements in performance.
D ABLATION STUDY FOR THE DESIGN OF SHIFTMATCH.
We performed an ablation study to choose the design of ShiftMatch (Fig. 7). EmpCov priors might suggest that applying ShiftMatch the input alone might be sufficient. However, we found applying ShiftMatch only on the input data had performance much worse than other ShiftMatch variants that matched intermediate features. which emphasised the importance of correcting corrupts at all layers. Second, we considered modifying the spatial covariances. In particular, we considered "w/o Kronecker factored covariance". where directly estimate the HW HW covariance matrices, and "channel-joint" where we compute only a single spatial covariance matrix, aggregating information across all channels (we usually have C' separate W HW matrices one for each channel). Both of these modifications performed worse than our approach. Presumably, "w/o Kronecker factored covariance" compromises the ability obtain accurate covariance estimates. and "channel- joint" compromises the model's ability to match account for different spatial statistics in different channels. Finally. we considered applying ShiftMatch on the post-activation feature, which gave very slightly worse performance than the usual pre-activation approach. Nonetheless, note that all of these sub-optimal variants outperform plain HMC across all corruption intensities.
E ALGORITHMS
We summarised the end-to-end inference procedure of ShiftMatch in Alg. 1.
14

---

Published as a conference paper at ICLR 2023
Log- Big-Likelihood
Accuracy
0.8
-0.5
0.7
-1.0
0.6
Corruption Intensily
Corruption Intensity
HMC
ShiftMatch (wis Kronecker factored covariance)
ShirtMatch (post activation)
ShiftMatch (input only)
ShiftMatch (channel-jiint
ShifMaath
Figure 7:
The performance of various variants ShiftMatch (described in the main text).
Table 2: Accuracy improvement given by using ShiftMatch the plain method. Note that the benefits of ShiftMatch under HMC are around double the benefits of ShiftMatch for SGD or deep ensembles.
Corruption Intensity
2
5
SGD
2.8%
5.1%
7.7%
10.4%
13.1% 1%
Deep Ens
1.8%
3.7%
6.2%
9.1%
12.4%
HMC
5.2%
8.9%
12.8%
18.3%
24.6%
Table 3: Accuracy improvement given by using test-time batchnorm over the plain method.
Corruption Intensity
4
5
SGD
1.3%
2.6%
4.2%
5.8%
8.2%
Deep Ens
1.4%
2.8%
4.9%
7.2%
10.6%
HMC
4.5%
7.6%
10.9%
15.8%
21.9%
Table 4: Accuracy improvement given by using ShiftMatch over test-time batchnorm.
Corruption Intensity
2
4
SGD
1.5%
2.5%
3.5%
4.6%
4.9%
Deep Ens
0.4%
0.9%
1.3%
1.9%
1.8%
HMC
0.7%
1.3%
1.9%
2.5%
2.7%
F INTRODUCING SHIFTMATCH INTO A MODEL WITH BATCHNORM
When we add ShiftMatch to network which was trained with batchnorm, we add ShiftMatch after the batchnorm layer. In these networks with ShiftMatch, we then have choice: do we set test-time batchnorm or off? To answer this question, we tried both (Fig. 8). We found that keeping test-time batchnorm on "SGD w/ SM & BN" performed slightly better than turning off test-time batchnorm "SGD w/ SM".
G CIFAR-10 10 WITH NON-BAYESIAN SHIFTMATCH WITH DATA AUGMENTATION AND TRAINING TIME BATCHNORM.
Remember that Izmailov et al. (2021b) used an unusual setting without data augmentation or train- time batchnorm. We therefore additionally evaluated whether ShiftMatch improved robustness to corruption on CIFAR- 10 in the standard setting with data augmentation and train-time batchnorm. In particular, trained ResNet-20 using SGD with momentum, with an initial learning rate of 0.05, a cosine annealing schedule, and weight decay coefficient 10 We augmented the training data through random horizontal flip and random crop. To combine ShiftMatch with batchnorm,
15

---

Published as a conference paper at ICLR 2023
Algorithm I End-to-end procedure of ShiftMatch on pre-trained BNN
Input: Test set X Training set Xtrain
M posterior samples from an L-layer BNN 07............1 (I....,M).
Step 1. Acquire training statisties
for m72 1,2. .M do
> Acquire the training statistics for all M posterior samples.
H X
train
for L = 0, 1, ... L do
> Compute the training statistics M O Compute the mean.
M M
Average (H
H
H train
M
urin
1/2
f.
(Hé
train
T (H (rain)
Compute the matrix square root.
uan
Φeti (H
> Forward pass to the next layer.
train
end for
Save the training statistics f. train
Quam
(ram))e E (O.....L.. L- 1}.
end for
Step 2. Perform Bayesian Model Averaging 0
Initialise the predicted probabilities Bayesian Model Averaging over M posterior samples ) l E {O..... -- 1} from the disk.
for m 1,2, . do
Load (M tsin
fm
H t X
lest
for l 0, 1..
do M lest f- Average(H (H lest)
Match the test-time feature for each layer. Compute the mean of test feature.
H lest f- H M
Lest
1/2
Q
H H Hast)
>> Compute the inverse of symmetric matrix square root.
N
H (++1 t 0t+1
(H Qies Q f.m
+ M 2. trin
ShiftMatch transformation (Eq. (8)).
lest
end for p P (H(a)) where f can be e.g. softmax.
Update the BMA prediction
end for
return p
added ShiftMatch after the batchnorm layer (Appendix F). Again, we found that ShiftMatch offered improvements over simply using batchnorm, especially the higher corruption intensities (Fig. 9).
H ANALYSIS CALIBRATION ERROR
We have analysed expected calibration error (ECE) in Fig. 10. Broadly, SGD has very poor accuracy and ECE, while ShiftMatch performs best in terms of accuracy, and deep ensembles performs best in terms of ECE. We could understand these results as ShiftMatch trading a bit of ECE for alot of accuracy (see Fig. 10 right). However, comparing ECE difficult when accuracy differs so dramatically (c.g. for corruption intensity 5; see (Nixon et al., 2019) for further details). As one example, remember that ECE bins the models confidence (e.g. 3. in the range 0.5-0.6), and asks whether the model's probability of being correct is close to the confidence (in this case, the probability of being correct should be about 0.55). When the model more accurate, more of its probabilities will be close to 0 or close to meaning that there will be fewer samples in the intermediate bins (e.g. Figure in Garriga-Alo Alonno et al., 2018). Having fewer samples in the intermediate bins implies that there will be larger errors in the calibration curves just due to finite samples, and larger errors in the calibration curve of course means a higher ECE. This is one example where changes the ECE may occur not due to changes in the model's actual calibration, but due to finite sampling errors.
16

---

Published as a conference paper at ICLR 2023
0.0
0.95
0.90
0.2
Log- LieLLieelihodd
Accuracy
0.85
0.80
0.75
-0.6
0.70
-0.8
0.65
0.60
1..
3 Corruption Intensity
2
2
3
4
5
Comuption Intensity
SGD
SGD w/ SM
Deep Ens
Deep Ens w/ SM
SGD w/ BN
SGD w/ SM & BN
Deep Ens w/ BN
Deep Ens w/ SM & BN
Figure 8: Various combinations of test-time batchnorm and ShiftMatch. "SGD" implies no Shift- Match, and no test-time batchnorm. "SGD w/ BN" implies test-time batchnorm but no ShiftMatch. "SGD SM" implies ShiftMatch but no test-time batchnorm. "SGD w/ SM & BN" implies Shift- Match combined with test-time batchnorm.
0.0
0.95
0.90
0.2
Accuracy
0.85
Log-Likelihood
0.80
0.75
-0.6
0.70
-0.8
0.65
0.60
1.0
2
3
0
3
4
5
Corruption Intensity
Corruption Intensity
SGD
SGD BatchNorm
SGD w/ ShiftMatch
Deep Ens
Deep Ens w/ BatchNorm
Deep Ens w/ ShiftMatch
Figure 9: All combinations of SGD (red, left three bars), ensembles (blue, right three bars) with batchnorm and ShiftMatch for CIFAR-IO-C with data augmentation and batchnorm at train-time.
1 PERFORMANCE OF SHIFTMATCH ON MIXTURE OF CLEAN AND CORRUPTED DATA
Here, we considered whether ShiftMatch still improves performance when test minibatches are 50-50 mix of uncorrupted and fully corrupted (samples of intensity 1-5). Specifically, we constructed CIFAR-10 test minibatches with a 50-50 split of uncorrupted and fully corrupted (intensity level 1-5) data. Thus, the full combined test set, consisted of 20,000 examples. We trained SGD and deep ensembles using code from Izmailov et al. (2021b), and used HMC samples from Izmailov al. presented in Fig. 11. Test-time batchnorm still gives improvement, and
(2021b). The results ShiftMatch gives a further improvement for all methods tested (SGD, deep ensembles and HMC), except for SGD for the log- likelihoodd Of course, the improvement is less spectacular than when the test minibatches all consist of corrupted data (Fig. and Table. 2).
17

---

Published as a conference paper at ICLR 2023
0.9
0.8
0.3
0.3
Accuracy
0.2
ECE
0.2
0.6
0.5
0.1
0.4
2
2
4
0.6
0.7
0.8
Corruption Intensity
Corruption Intensity
Accuracy
Deep ShiftMatch
Intensity Intensity Intensity
Figure 10: Expected calibration error (ECE) different methods. We begin by showing the accuracy (left) and ECE (middle) separately, with corruption intensity on the x-axis, for SGD, HMC, deep ensembles and ShiftMatch. Note that, for all experiments, the ECE numbers are computed using the same code and setting. The results for SGD is acquired by averaging 5 independent trials of different random seeds as this method had very high variance in the ECE. The trials are conducted using the code of Izmailov et al. (2021b), which used filter response normalization (FRN) (Singh & Krishnan, 2020) rather than batch normalization (Ioffe & Szegedy, 2015) and did not apply data augmentation. The right most figure directly compares ECE and accuracy under different corruption intensities. Results under the same corruption intensity are connected via grey lines, and different methods are presented with different markers and colours. Note the accuracy of SGD is lower than the numbers reported in Figure. 1, which are directly extracted from Fig. 7 in Izmailov et al. (202 We believe that this is due Izmailov al. (2021b) using only one random seed, whereas here we averaged over five random seeds trained using their code.
0.95
0.0
0.90
0.2
0.85
-0.4
Likellhood
Accuracy
0.75
-08
0.70
0.65
-1.2
0.60
0.55
-1.6
0.50
3
4
2
5
Corruption Intensity
Corruption Intensity
SGD
SGD w/ BatchNorm
SGD w/ ShirtMatch
Deep Ens
Deep Ens w/ BatchNorm
Deep Ens w/ ShittMatch
HMC
HMC w/ BatchNorm
HMC w/ ShiftMatch
Figure 11: Performance of plain, test-time batchnorm and ShiftMatch in combination with SGD, deep ensembles and HMC on CIFAR-10 10 where we test on combined minibatches with a 50-50 split of uncorrupted and fully corrupted (intensity level 1-5) data.
18


Reviews:
[{"review_id": "qC236SH3Yrg", "reviewer": "Reviewer_NoHk", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes ShiftMatch – a transformation of the training data that represents likelihood in Bayesian neural networks. This data-dependent likelihood helps dealing with corruptions of the data at test time. It can perfectly remove certain types of corruptions. Interestingly, it does not change the training data, so this transformation can be applied on pre-trained networks. Moreover, it can be applied to normal neural networks, not only Bayesian neural networks.", "summary_of_the_review": "I suggest acceptance of the paper as I believe the strengths outweigh the weaknesses. The only real big weakness is the absence of comparison to other methods other than EmpCov dealing with covariate shift. \n\nI would be interested to know the authors thoughts on the two weaknesses I mentioned: absences of other baselines and runtime discussion – during the rebuttal\n", "strength_and_weaknesses": "**update after rebuttal** I have read the authors’ response and other reviews. The authors did clarify mostly my concerns so I am increasing the score.\n=======================\n\n\n\nStrengths:\n* The paper proposes an interesting way to deal with test time data corruptions without the need to touch the training time. This allows a user to use any of the pretrained networks and apply ShiftMatch only at test time. With the caveat that one would need an access to training data to compute statistics required for transformation. \n* This transformation is neatly incorporated into Bayesian formulation. \n* The paper provides rather extensive empirical evaluation on the method, though check the weaknesses on experiments. \n* The paper is well written mostly and easy to follow.\n\nWeaknesses:\n* Experiments are missing comparison with other methods that deal with corruptions or covariate shift. The only baseline that directly addresses the same problem is EmpCov, but there are other works, e.g., Trinh et al. 2022. Tackling covariate shift with node-based Bayesian neural networks. ICML.\n* (not major) I am missing a bit more discussion how ShiftMatch fits into the Bayesian deep learning. For example, does it help/hurt with cold posterior effect? (Wenzel et al., 2020. How Good is the Bayes Posterior in Deep Neural Networks Really? ICML). Also, there is a discussion on effect (non-decremental as one might think) on OOD detection, but it would be great to see more discussion and experiments on effect of ShiftMatch on uncertainty quantification provided by BNNs.\n* (not major) A bit more discussion and experiments on approximations for covariance estimates would be beneficial. I find results from Appendix C are rather surprising and interesting, but lacking more explanation here. It is unclear for me why ““w/I Kronecker factored covariance” compromises the ability to obtain accurate covariance estimates. \n* Test time runtime discussion is missing. The paper rightly so emphasises that training time is unaffected by ShiftMatch, but how slower does the test time become with ShiftMatch? There is a brief mentioning of runtime for spatial batchnorm (a potential implementation of ShiftMatch).\n* (not major) It would be interesting to see the comparison not only to batchnorm but with decorrelated batchnorm cited in the paper.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "ezD-7K2se3", "reviewer": "Reviewer_BEej", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "", "summary_of_the_paper": "This paper proposes ShiftMatch, a new training-data-dependent likelihood for robustness to corruption in  pre-trained Bayesian neural network. ShiftMatch allows it to use publicly available samples from pre-trained BNNs (Although it needs train data on test phase).", "summary_of_the_review": "While the subject of this paper is interesting and novel, I believe the authors' claims are not adequately supported due to a series of weaknesses. I am more than willing to increase my score if the authors address the aforementioned limitations.", "strength_and_weaknesses": "Strengths :\n\n- It's ingenious to try to see a change in the prior as a change in the likelihood.\n\n- The performance is good compared to the comparison methods (BatchNorm, EmpCov).\n\nWeaknesses :\n\n- Chapter 4.1, which is the motivation of this paper, is hard to follow. More specific, authors claim that\n\n$w \\sim N(0, \\frac{1}{C N_{\\text{train}}}( X_{\\text{train}}^{\\top} X_{\\text{train}}) )$,\n\n$X_{\\text{test}} w \\sim N(0, \\frac{1}{C N_{\\text{train}}} X_{\\text{test}} ( X_{\\text{train}}^{\\top} X_{\\text{train}}) X_{\\text{test}}^{\\top})$\n\ncan be seen as\n\n$\\operatorname{ShiftEmpCov}(X_{\\text{test}} ; X_{\\text{train}}) = X_{\\text{test}} ( X_{\\text{train}}^{\\top} X_{\\text{train}})^{1/2}$,\n\n$w^{\\prime} \\sim N(0, I)$,\n\n$\\operatorname{ShiftEmpCov}(X_{\\text{test}} ; X_{\\text{train}})\nw^{\\prime} \\sim N(0, \\frac{1}{C N_{\\text{train}}} X_{\\text{test}} ( X_{\\text{train}}^{\\top} X_{\\text{train}}) X_{\\text{test}}^{\\top} )$.\n\nThis statement is obviously true, but however, I think these statement is related to prior distribution of $X_{\\text{test}} w$. \nIf we change prior like above, then what can you say about the posterior and predictive distribution? I think more detailed explanation is needed.\n\n- Proposed method can be summarized as follows : Match the first and second moments of the test data distribution with the those of the train data distribution. In other words, the algorithm compares test dataset with train dataset, and make test dataset distribution similar to train dataset distribution by covariate shift. I wonder if it will work well even if half of images are corruption-free and the other half are corrupted by intensity 5. Considering the idea of proposed method, I'm worried that it works well only for one unified corruption intensity.\n\n- Is it much more effective to match the distribution in all layers than to match only the distribution of input data? I wonder what it means to consider all the layers.\n\n- It would have been nice to measure the ‘Expected Calibration Error’, which is commonly used to evaluate Bayesian models on classification problem.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "YiXfYkk7XMH", "reviewer": "Reviewer_MwDb", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents ShiftMatch, a simple method based on a data-dependent likelihood, enhancing the robustness of Bayesian neural networks to corrupted data. Given a (pre)trained BNN, ShiftMatch computes the training data statistics (mean, variances) and uses them to construct a transformation matrix for each layer. The prediction for a test input is then done with the transformation matrix computed from the training statistics. The prediction with the training data-dependent transformation alters the likelihood, but this transformation leaves the likelihood evaluated with the training data invariant; that is, even though the likelihood was changed, it remains the same for the training data, so one can just reuse pretrained BNNs trained with typical data-independent categorical likelihoods. For high-dimensional data with spatial structures, the authors further propose structural covariance estimation schemes that could improve the scalability of the method. The benefit of the proposed method is demonstrated via various image classification benchmarks with corruptions.", "summary_of_the_review": "In summary, I enjoyed reading the paper, and find it to be a novel contribution with promising results.\n\npost-rebuttal: I have read the authors' response which resolved my concerns. I keep my original score.", "strength_and_weaknesses": "Strength\n- The paper is well-written and easy to follow. \n- The problem of interest is indeed an important topic in the literature, so the motivation is clear.\n- I find the idea of the data-dependent likelihood that does not alter the training procedure interesting and novel. I think the design of the ShiftMatch likelihood and its interpretation as a data-dependent likelihood is clever.\n- The experiments are thorough, and the results are convincing. Including the additional experiments in the appendix, the authors provide various ablations and alternatives that can further strengthen the arguments made in the paper.\n\nWeakness\n- If I have to be picky, the method can be considered incremental since it is largely inspired by existing methods including EmpCov and test-time batchnorm. Still, I think the proposed method is valuable, especially the design of ShiftMatch likelihood leaving the training procedure unchanged.\n- As already discussed in the paper, the method requires storing or recovering training data for prediction. However, the arguments in section 4.5 to defend against this is also reasonable, so I won't think this is a critical problem.\n\nHere are some questions and minor comments:\n- EmpCov seems to transform only the input layers, but ShiftMatch seems to be applied for all the intermediate layers. Can you comment on this difference? For instance, would the performance of ShiftMatch be degraded if only applied to the input layers?\n- In page 7, the paper explains the reason why ShiftMatch performs poorly for \"Translate\" corruption. It seems that ShiftMatch also underperforms baselines for \"Scale\", and slightly underperforms deep ensembles for \"Gauss blur\". The latter is quite disappointing since section 4.6 explains how ShiftMatch can \"perfectly\" removes stationary linear corruptions. Can you comment on this?\n- I don't expect the size of the training data to stably estimate the statistics to be large. Having said that, one could think of a bagging-like procedure where multiple subsets of training data are sampled, corresponding statistics are computed, and multiple predictions are computed from the multiple statistics to be ensembled. This may make the interpretation of the data-dependent likelihood more cumbersome, but may enhance the robustness of the method.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster