Improving Neural Network Accuracy and Calibration Under Distributional Shift with Prior Augmented Data 
Authors: Jeffrey Ryan Willette
Juho Lee
Sung Ju Hwang
Venue:  (2021)
Keywords: Bayesian
Calibration

Abstract: Neural networks have proven successful at learning from complex data distributions by acting as universal function approximators. However, neural networks are often overconfident in their predictions, which leads to inaccurate and miscalibrated probabilistic predictions. The problem of overconfidence becomes especially apparent in cases where the test-time data distribution differs from that which was seen during training. We propose a solution to this problem by seeking out regions in arbitrary feature space where the model is unjustifiably overconfident, and conditionally raising the entropy of those predictions towards that of the Bayesian prior on the distribution of the labels. Our method results in a better calibrated network and is agnostic to the underlying model structure, so it can be applied to any neural network which produces a probability density as an output. We demonstrate the effectiveness of our method and validate its performance on both classification and regression problems by applying it to the training of recent state-of-the-art neural network models.

Under review as a conference paper at ICLR 2021
IMPROVING NEURAL NETWORK ACCURACY AND CALIBRATION UNDER DISTRIBUTIONAL SHIFT WITH PRIOR AUGMENTED DATA
Anonymous authors Paper under double-bling blind review
ABSTRACT
Neural networks have proven successful learning from complex data distribu- tions by acting as universal function approximators. However, neural networks are often overconfident in their predictions, which leads inaccurate and mis- calibrated probabilistic predictions. The problem of overconfidence becomes es- pecially apparent in cases where the test-time data distribution differs from that which was seen during training. We propose solution to this problem by seeking out regions in arbitrary feature space where the model is unjustifiably overcon- fident, and conditionally raising the entropy of those predictions towards that of the Bayesian prior on the distribution of the labels. Our method results in a better calibrated network and is agnostic to the underlying model structure, it can be applied to any neural network which produces probability density as an output. We demonstrate the effectiveness of our method and validate its performance both classification and regression problems by applying it to the training of recent state-of-the-an neural network models.
INTRODUCTION
While deep neural networks have achieved success on many diverse tasks due their ability to learn highly expressive task-specific representations, they are known to be overconfident when presented with unseen inputs from unknown data distributions. Probabilistic models should be accurate in terms of both accuracy and calibration. Accuracy measures how often the model's predictions agree with the labels in the dataset. Calibration measures test the accuracy of the uncertainty around probabilistic output. For example, event predicted with 10% probability should be the empirical outcome 10% of the time. The probability around rare but important outlier events needs to be trustworthy for mission critical tasks such autonomous driving.
Bayesian neural networks (BNN) and ensembling methods are popular ways to achieve a predictive distribution for both classification and regression models. Since Gal & Ghahramani (2015) showed that Monte Carlo Dropout acts as a Bayesian approximation, there have been numerous advances in modeling predictive uncertainty with BNNs. As laid out by Kendall & Gal (2017), models need to account for sources of both aleatoric and epistemic uncertainty. Epistemic uncertainty arises from uncertainty in knowledge or beliefs in system. For parametric models such BNNs, this presents as uncertainty in the parameters which are trained to encode knowledge about data distribution. Aleatoric uncertainty arises from irreducible noise in the data. Correctly modeling both forms of uncertainty is essential in order to form accurate and calibrated predictions.
Accuracy and calibration are negatively impacted when the data seen during deployment varies substantially from that seen during training. It has been shown that when test data has undergone a significant distributional shift from the training data, one can witness performance degradation across all models (Snoek et al., 2019). A recurring result from Snock et al. (2019) is that Deep Ensembles (Lakshminarayanan et al., 2017) show superior performance on shifted test data. Previ- work has also shown that BNNs fail to accurately model epistemic uncertainty, regions with sparse amounts of training data often lead to confident predictions even when evidence to justify such confidence is lacking (Sun al., 2019). Bayesian non-parametric models such as Gaussian

---

Under review as a conference paper at ICLR 2021
processes (GP) also model epistemic uncertainties, but suffer from limited expressiveness and the need to specify a kernel a priori which may not be feasible for distributions with unknown structure.
In this work, we propose generated samples from a variational distribution which augments the natural data to seek out areas of feature space for which the model exhibits unjustifiably low amounts of uncertainty. For those regions of features, the model is encouraged to predict uncertainty closer to that of the Bayesian prior belief. Our method can be applied to any existing neural network model during training, in any arbitrary feature space, and results in improved accuracy and calibration on shifted test data.
method for achieving accurate and calibrated models by providing
Our contributions in this work are follows:
We propose
method of data augmentation, which we dub Prior Augmented Data (PAD) that seeks to generate samples in areas where the model has an unjustifiably low level of epistemic uncertainty.
We introduce method for creating OOD data for regression problems, which to our knowledge has not been proposed before.
We experimentally validate our method shifted data distributions for both regression and classification tasks, on which it significantly improves both the accuracy and the cali- bration of a number of state-of-the-art Bayesian neural network models.
2 BACKGROUND
For regression tasks, we denote set of features x E Rd and labels R which make up a dataset of i.i.d. samples D = {(x,. (n)) with := {x, 1. Let fe be a neural network which is parameterized by weights 0. Let the output of fa be probability density po(y)x) which is either in the form of a Gaussian density NNH, or) for single task regression or a categorical distribution in the case of multi-class classification. Let 96 be generative model which generates pseudo inputs for fo- We assume that both So and are iteratively trained via mini-bateh stochastic gradient descent with updates to generic model parameters given by the update rule T 7: V.,C. with & representing a loss function which is differentiable w.r.t generic parameters т. We refer to an out-of out-of-distribution (OOD) or distributionally shifted dataset D as one which is drawn from different region of the distribution than the training dataset D. This distributional shift can occur naturally for multiple reasons including temporal modal shift, or an imbalance in training data which may come about when gathering data is more difficult or costly in particular regions of D.
2.1 BAYESIAN NEURAL NETWORKS
BNNs are neural networks with distribution over the weights that can model uncertainty in the weight space. In practice, this often done by introducing a variational distribution and then min- imizing the Kullback-Leibler divergence (Kullback, 1997) between the variational distribution and the true weight posterior. For further discussion of this topic, we refer the reader to existing works (Kingma & Welling, 2013; Blundell et al., 2015; Gal & Ghahramani, 2015). During inference, BNNs make predictions by approximating the following integral with Monte Carlo samples from the variational distribution (()
8
ltd. p(u|x, D) = po(v)x))(011))00 = po. (0l 01....,0s q(0[D). (1)
2.2 MISPLACED CONFIDENCE
A problem arises when neural networks do not accurately model the true posterior over the weights p(0[D) given in (2). Our conjecture is that a major factor contributing generalization error in both likelihood and calibration is a failure to revert to the prior p(0) for regions of the input space with insufficient evidence to warrant low entropy predictions. Bayesian non- parameetee models such as Gaussian processes (GP) solve this through utilizing a kemel which makes pairwise comparisons between all datapoints. GP's come with the drawback of having to specify kernel a priori and
2

---

Under review as a conference paper at ICLR 2021
Dropost
Deep
PottO 4.17
popt -1 01
19 P0100-2-2
log pornt 38678
ug pMMM 0.97
Dropout FAD
Ensembbe
SWAG PAD
Pank One PAD
Cal 125
CAL DAP 3:33
CAL Em. 5.66
CAL CALEm 0.37
Figure Performance of different models on OOD test data. The top row contains baseline models, the bottom row contains baseline PAD. It can he seen that PAD effectively increases the epistemic uncertainty in regions with sparse training data and exhibits more predictable reversion to the prior, much like a GP (top right). Baseline BNN models in the top row show unpredictable hehavior in OOD regions of input space.
generally outperformed by deep neural networks (DNNs) which are known to be more expressive than GP with common RBF kernel.
p(0 p(D10)p(0) (p('d')
The problem of misplaced confidence of neural network was first studied by Guo et al. (2017) which showed that modern DNNs have poor correlation between their confidence and actual accuracy, and often overconfident. Snoek et al. (2019) came to the general conclusion that Deep Ensembles (Lakshminarayanan et al., 2017) tend to be the best calibrated model on OOD test data. Since then, there have been number of newly proposed BNN models such SWAG (Snoek et al., 2019) Multi-SWAG (Wilson & Izmailov, 2020), and Rank One Bayesian Neural Networks (RIBNN) (Dusenberry al., 2020), each of which utilize different strategies for modeling p))
To illustrate the problem of failing to revert to the prior in underspecified regions of input space, have provided toy example in (figure 1). The true function given by x + + sin(4( + <)) + sin (13(x + e)), where ~ AO0 0.03). We then sample 100 points in the range of [0, 0.4] and 100 points in the range of [0.8, 1.0], which leaves gap in the training data. One can observe in (Figure 1) how different neural network models tend to make confident predictions in regions where they have not observed any data, and exhibit unpredictable behavior around the outer boundaries of the dataset. Our method encourages a reversion to the prior in uncertain regions, similar to the behavior of a GP. addition to covering the gap between the datapoints, our model exhibits more predictable behavior around the outer boundaries of the data when compared to other BNN models.
METHOD
To encourage reversion towards the prior in uncertain regions, we learn to generate pseudo OOD data which leads to a better calibrated model by raising the entropy of the OOD predictions. An important design choice that we make is for our OOD generator network (-)(.) to take dataset X as input to produce distributions of OOD data. The goal of the OOD generator is to fill the "gaps" between the training data, using all available current knowlege the training data itself.
Once trained, the model fe should predict more uncertainty for data generated from gs- Raising the entropy of uninformative noisy psuedo inputs may be solution, but also could be too distant from the natural data, and therefore provide no useful information for fo to learn from. Ideally, we want realistic OOD data that are not too distant from the training data and still predicted with more uncertainty by fo. To achieve this, we employ an adversarial training procedure similar to generative adversarial nets (Goodfellow et al., 2014a)
train gs find where fo may be overconfident,
3

---

Under review as a conference paper at ICLR 2021
and at the same time train fo to defend against this by predicting higher levels of uncertainty in those regions. In the next section, we explain the objectives 9s.
3.1 THE OUT-OF-DISTRIBUTION SAMPLE GENERATOR NETWORK 96
g6 takes dataset X {x: and produces distribution of an equally sized pseudo dataset X = [x,]==1 First, each x; is encoded via a feedforward network Dene to construct a set of representations Z {z, Gene (x,,) ) jn=1- Then, pick subset size K E [1. [N/2]] for each mini-batch, and for each n = 1, I,....N.
1
Do (X) Idec
2m
K
mEnng
where nn (m) denotes the set of K-nearest neighbors of 2n and Idec is another feedforward net- work. The distribution for the generated data X then defined as
Lo (X|)=
II go (Xn
As stated previously, our goal is to generate psuedo-OOD OOD data, but we cannot be sure where such data will arise from. The only thing we can be certain of is that the training data exists, and 2) there exists some distribution which is OOD to the training data. Therefore, our only option is to learn directly from representations of training data order find likely regions of OOD data.
3.2 TRAINING OBJECTIVES
Training 9o Given a batch of data XB )x,) we first construct the OOD distribution
(x(x)|))) via Jø. The training loss for X defined as
(x)))
Σ
Bas((n))
H/po (1po(2/X2n) (v/xx)]
Σ
Xm
BK
== mmmm (n)
(5)
and H(p] == f log p(x)dx is the entropy, and nn (m) in C is the set of K K-nearest neighbors of X, among Xg. The A term trains 9s to fool fo by seek- ing regions where fo(X,) makes low entropy predictions, i.e., where it may be overconfident on OOD data. The B term encourages diversity of the generated samples by maximizing the entropy of (&,). Without B, the generated samples would be prone to mode collapse and become homoge- neous and uninformative. The final term. C. minimizes the average pairwise distance between the real data and the generated data it is conditioned on, where x.m. XB. As the C term minimizes this distance, the generated data is likely to exist the "gap"" regions of the natural training data, Com- bined with the training loss for fo which we will describe next, we aim to train Je to produce X that are not too distant from the training data but still differ enough to warrant an increase in uncertainty. We provide ablation study to see the effects of the terms A, B, and C, in tables 4 and 5.
where qo (xx) == 90 (Xn199
(x))) (X))
Given the training data D, we optimize the expected loss Lo = Ex, (X B)) over subsets of size B obtained from X. In practice, at each step, sample single mini-batch X to compute the gradient. Also, we choose qo (X,,. to be a reparameterizable distribution, and approximate the expectation over X, via single sample X, (x,)).
For regression tasks we found the constraint given by C 5 was too strong of a constraint given that the generation happens directly in the input space and the datasets generally of low dimension- ality. For regression, we therefore only started to penalize the C term distance when it exceeded a certain boundary threshold. We used ||]d]], where d is the dimensionality of the inputs.
4

---

Under review as a conference paper at ICLR 2021
spit
shifted split
lue
UC: Housing dataset
Figure 2: Left: A random train/test split. The test data is randomly chosen from the whole distribution. Clusters resulting from spectral clustering. Right: Shifted data distributions chosen from clusters as train/test split for our experiments on OOD data.
Middle:
Training fo For fa. we minimize the expected loss,
La =
+ Exg [X),,
(6)
where the first term is the negative log-likelihood of the natural data, and the second term is regularizer which encourages more uncertainty on OOD data,
1 B
B
(Ke) (1-eex) exxp
(_minm minm [R, Xm 1²)) (=Lp)(()(
ro (KB)
Σ
E
where
scaling parameter. The role of the regularizer is to encourage predictions on the OOD data to be closer to the prior ply). Note that the KL term decays towards zero as x becomes closer to x. This allows fo the freedom to make confident predictions in regions where real data exists. We draw mini-batch XB from X, approximating ro (XB) with a single sample Xg go (X The loss is then approximated by (8),
B
La R
log Po xx.) + (XB)
(8)
==1
In terms of added complexity, PAD optimizes one additional set of parameters 6 and requires an alternating pattern of training akin to that of GAN's (Goodfellow et al., 201 In our experiments, we utilize a shallow network with single hidden layers for ő. For further information about how we implement the KL term in we we refer the reader to section 8 for details and algorithms.
4 EXPERIMENTS
4.1 EXPERIMENTAL SETUP AND DATASETS
For regression, we use UCI datasets (Dua & Graff, 2017) following Hernández-Lobato & Adams (2015). The base network a multi-layer perceptron with two hidden layers of units with ReLU activations. The generator uses single hidden layer of 50 units for both the encoder and decoder and uses a permutation invariant pooling layer consisting of [mean(x), max(=)). The encoding of X and generation of X are done directly in the input space. We train the models for a total of 50 epochs on 10 random splits of train/test data. In order to create shifted test set, we first run spectral clustering algorithm to get 10 clusters on each dataset. We then randomly choose test clusters until we have training. We repeat this process 10 times for each dataset. A TSNE visualization of a dataset created this way is given in (figure 2). We do this to ensure there significant shift between training and testing data. We report both negative log likelihood (NLL) and calibration error proposed by Kuleshov et al. (2018) measured with 100 bins on the cumulative distribution function (CDF) of the density Pa (y: xx). We show calibration plots in (figure 3), plotting the expected outcome frequency versus the empirical frequency. For baseline models, we tune hyperparameters with a randomly chosen 80/20 training/validatton
set which > 20% of the total dataset size. then use the remaining clusters for
split. For PAD models, selecting half of the previously made clusters for each fold. For all models and datasets, we tune the hyperparameters for each individual train/test split.
do k-fold cross validation with K 2,
5

---

Under review as a conference paper at ICLR 2021
Negative log likelihood UCI regression datasets. GP's and FVBNN's included for reference. bold entries contain the best result for the base model. underlined entrics are those with large difference of 1 (log scale) between methods.
Table I:
Model
Housing
Concrete
Energy
Kin8nm
Naval
Power
Wine
Yacht
3,81+0.23
4,44++0.88
6.7447744
-0.50+0.04
-3,55+22.07
0,14+0.336 1.22++0.05 4.31 +0221
FVBNN
1.62+1108
4,80+t.50
3.16::1.02
0.01
-3.19-0.400
0.18 1,30+0.00 0.26+0.00
DUN
1.66
4.78++0.3
4.41::1.22
0.85+1.49
-1.16:0.77
4.34+1.22
3.61
6.37+3.
DE
6.83:12.92
8.34:14.66
10.73:1:3.25
-1.03.1..23
1..42:1:.92
3.12:10.18
2.18:10.67
2.39:10.64
DE PAD
3.61=0022
5.13++0.88
3.24++0.62
-0.38+0.10 10
-2.81=00.6
3.12++0.08
1.24==0.08
0.78+
RIBNN
4.18++0.70
5.17+11.2 1.27
9.22-+-9.74
-0.86+0.19
1..3+-10.400
2.95+-0.12
1.72++1.34
3.29-+0.64
RIBNN PAD
3.24
4.30
3.78++0.18
0.09:0022
-3.25-+0.20
3.13:1:.05
1.27:10.10
4.14
SWAG
5,02:12.26
1.85+t1.00
3.73:12.83
-0.96:10118
0.37
3.03:10.14
1.41=10.39
4.25:10.71
SWAG PAD
3.80:10.58
4.47:10.29
3.511:1.69
-0.5810.09
-2.46.10.88
3.18:10.12
1.20:
3.13
MC Drop
1.37+1142 3.88+1.799
4.04+-4.23
+0,82+0.23
4.93+7.55
3.09--0.15
1.70+-0.33
3.21++2.05
MC Drop PAD
1.92 0.98
3.34+0086
+0.53+0.09
-0.96
3.18++0.88 1.22++0.07 3.45+ 14
Table Calibration error on UCI regression datasets. GP's and FVBNN's are included for reference, bold entries contain the best result for the base model. underlined entries are those with large difference of 2 5 between methods.
Model
Housing
Concrete
Energy
Kin8nm
Naval
Power
Wine
Yacht
8.24+2273
3.34
4.88+3.14
2.17 6.65+7.39 4.30:14.72 1.52++..07 5.10+3..0
FVBNN
7.10=7.76
7.79:15.7
5.90=8862
4.46:5.006 19.5#23.9
3,23+2.41
2.06#:.47 2.10::0.00 2.46::0.00
DUN
15.8+21.5
14.3+10.9
16.5
22.415.4
19.3120.7
18.7122.4
18.9116.9
DE
15.95+66.97
18.83+8.66
8.94
1.34
8.06+
2.56+22.53
2.81+1.14
6.61+6.12
DE PAD
5.12::3.22
12.87-16.66
6.34
1.88
6.1913.93
1.60
1.53
2.37:10.18
RIBNN
6.42+15.45
14.18:88.77
4.82+19.37
1.63
6.41+4.93
1.07
1.29
2.82++0.61
RIBNN PAD
6.04+33.93
2.94
4.42+:3.3
4.63115.06
3.77+22.5
1.38:00.1
2.14
2.18::0.09
SWAG
12.43-+-.84
10.16:6..99
8.74
1.62+
2.98+-2.36
1.57+22.19
1.85+-2.26 4.91
SWAG PAD
6.73.16.44
8.72:14.33
6.05118.0
2.15
4.59
1.91
0.74110.51
2.761105
MC Drop
11.48+6.17
14.92+99.0
5.98+8.94
2.26++2.31 5.75+2288
2.334:217
1.46+00.73 6.74+66.95
MC Drop PAD
8.38-1832
12.78
6.30 2.30 5.52133.9 1.90 1.39.11.11 3.58
Classification experiments are done on both MNIST (LeCun et al., 2010) and CIFAR-10 (Krizhevsky et al., 2009) datasets. We use an architecture consisting of 4 and 5 convolutional layers respectively, followed by fully connected layers. We provide extra information regarding the exact architecture in the appendix (section 8). Instead of clustering create shifted test distributions, corruptions as were used by Snoek et al. (2019). As PAD works in any arbitrary feature space, we apply our method in the latent space between the last convolutional layer and the fully connected layers. We report classification accuracy, negative log likelihood, and expected calibration error (ECE) (Guo et al., 2017) over runs for all models.
image
4.2 BASELINES
We compare PAD against a number of Bayesian models and neural networks including Gaussian Processes, Functional Variational Bayesian Neural Networks (FVBNN) (Sun al., 2019), Monte Carlo Dropout (MC Drop) (Gal & Ghahramani, 2015), Deep Ensembles (DE) (Lakshminarayanan et al., 2017), SWAG (Maddox et al., 2019), and Rank One Bayesian Neural Networks (RIBNN) (Dusenberry al., 2020) and Depth Uncertain Networks (DUN) (Antorán et al., 2020).
4.3 ANALYSIS
For regression it can be seen in (table 1) that PAD generally improves the likelihood in the scenario where the shifted data distribution caused the baseline model to perform poorest. The underlined

---

Under review as a conference paper at ICLR 2021
Classification accuracy, NLL,, and calibration error OOD test for classification. Each row contains baseline model, with Mixup and PAD variants, Left: Accuracy, Middle: ECE, Right: NLLI
Table 3:
ACCURACY
ECE
NLL
Model
MNIST
CIFAR
Model
MNIST
CIFAR
Model
MNIST
CIFAR-I0
MC Drup
0.6110001
0.45=00.1
MC Drop
0.311:0.02
0.38=0.02
MC Drop
0.36::.226
3.23:10.00
MC Drop PAD
0.62
0.05
MC Drop PAD
0.07::0.07
0,05 0.01
MC Drop PAD
1,20:00.11
0.69
DE
0.68:20.00
0.52=00.00
DE
0,24:0.01
0.26=0.01
DE
1,52=0.11
2.06=0.03
DE
0.60:10.01
0.48--0.01
DE
0.18:-0.02
1.35:-0.07
RIBNN
0.69:10.03
0,43=00.01
RIBNN
0.10::0.05 0.06=0.02
RIBNN
1.47=0.58
1.76=0.01
RIBNN PAD
0.69:10.04
0.45:00.00
RIBNN PAD
0.07::0.04
0.07=0.04
RIBNN PAD
0.98=00.07
1.67=0.01
SWAG
0.68:10.01
0.59=0.00
SWAG
0.22::0011
0.10=0.00
SWAG
1.75=0.18 1.28=0.01
SWAG PAD
0.69:10.01
0.59:00.00
SWAG PAD
0.11::0.02
0.17=0.01
SWAG
1.10=0.08
1.3320.02
Housing
GF
MC Crop EAD
podlctt
132N3
ShC No
Crop F4B
MC Drep 1o
reqquecc
expected frequency
espected frequency
"espected trequency
Naval
peted
perkc:
Doposren Fredu Frequorrr
Mc
Mc Eirp AND
peted
0000m MD
No. Drop
frequercy
expected frequarcy
expected frequency
expected frequency
Figure 3:
Calibration curves for each model on each dataset. Base models are shown with dotted lines. Our models (model + PAD) are shown with solid lines. GP's are included with solid line for reference.
entries in (table 1) are those which the negative log likelihood differs by > 1 on the log scale. It can seen that the majority (8/9) of these increases in likelihood are achieved by PAD models. In terms of single experiments, PAD beat the baseline models 21/32 s 66% of the time. For calibration, in (table 2), the only difference that the underlined entries are those for which differ by 2 5. For calibration, it can be seen that PAD models contain all cases where large difference between methods is present. terms of single experiments, PAD beat the baselines 22/32 = 69% of the time.
In (figure 3) we show calibration curves for MC Drop and RIBNN models on each dataset (other models included in figure 9). Baseline models are shown with dotted line, while PAD models shown with a solid line. It can be seen that in datasets where the baseline model is poorly calibrated, such as Housing and Concrete, PAD results in curves which are closer to the center line than the baseline models. On datasets where the baseline models tend to be well calibrated, such as Power and Kin8nm, PAD does not result large deviation from the baseline. We attribute this to the fact that we tune the lengthscale parameter in such way that if no improvement can be made over the baseline model on the validation set, the lengthscale parameter and the second term (6) will have little to no effect on the resulting model parameters. However, the opposite is also true, which leads to large improvements in negative log likelihood and calibration error the worst case.
It is interesting to note that PAD does not give uniform performance gains across all datasets. To further investigate this, we used TSNE (Hinton & Roweis, 2002) to embed both the natural data X, and the generated data X for all datasets (figures 6 and 7). Datasets which naturally form a single dense cluster, such as Kin8nm and Power consistently the lowest performers when paired with PAD, where pad shows the best performance in 1/8 (NLL) and 2/8 (calibration) experiments. For datasets which form clusters manifolds, such as Naval and Yacht, PAD gives the best results

---

Under review as a conference paper at ICLR 2021
MNNTT
Dro AAD
HC DregiMpD
McDrep
Accurory
Interiity
Incensity
inter sirtityy
CIFAR10
NC NCorop
0c
06
MC Dragpppp
Crop AWO
Accurse
Intensity
Inteesity
Shift intersity
Figure 4: Model performance varying degrees of shift intensity for MNIST-C and CIFARIO- C. represents the original test set while 6 represents the most extreme level of shift. Models which augmented with PAD show comparable performance intensity increases, PAD models exhibit superior performance in terms of negative log likelihood and calibration error. Other models included in figure 10
the natural test set. As shift
4/8 (NLL) and 7/8 (calibration) experiments. PAD shows the strongest performance on datasets which naturally form separable clusters such Housing, Concrete, Energy, and Wine, for which it performs best on 16/16 16 (NLL) and 13/16 16 (calibration) experiments. These results suggest that PAD would be most useful in situations where the data naturally form loose, separable clusters where it is likely that a new cluster may form outside of the region of clusters which are known during training.
4.4
ABLATION STUDY
Table 4: Equation 5 Ablation: NLL
Table 5: Equation 5 Ablation: Calibration Error
Dataset
Regular
Without A Without B
Without AB
Dutaset
Regular
Without A
Without B
Without AB
Housing
4.32+1192
4.62+:2.31
4.69+42.30
4.70++2.31
Housing 8.38-18.32
9.43:8822
9,6911833
9,70=8.35
Concrete
4.98:1.79
4.96:10.3
5.054:0.7
5.05:t0.74
Concrete
12.78-66.71
14.04:47.44
14.53+:7.31
14.53-17.31
Energy
3.34=10.86
3.26=1100
3.34::1.3
3.34
Energy 6.30:17.71
6.53418.9
6.52-8.30
Kin8nm
-0.533+0.09
-0.53++0.09
-0.48+0.111
-0.48+0..11
Kin8nm
2.30+f-1.29
2.24
2,47+11.11
2.47=1661
Naval
-0.96-1..39
0,21
0.314:3.69
0.224:3.20
Naval
5.524:3.96
5.4433.08
5.88:13.94
5.65=3.71
Power
3.18:10.08
3.23:10110
3.23:1020
3.23110.0
Power 1.90::1.66
1.9721.56
1.83+11766
1.83=1822
Wine
1.22:1.07
1.3420.57
1.72411.1
1.70:11666
Wine 1.3911..11
0.5010.31
0.50-10.31
0,48=028
Yacht
3.4520.14
3.45#0.11
3.4111.113
3.4110.133
Yacht 3.58+0081
3.42:10.74
3.44-10.82
3.44-0.82
In order to understand the effect of each term in equation 5, we performed an ablation each term on the MC PAD (MC Dropout PAD) variant. It can be seen that as more terms are removed from the equation, the model performance degrades in both NLL and calibration. The effect is more pronounced for NLL, as the full equation contains the majority the best performances.
5 RELATED WORK
Bayesian Methods for Deep Learning Functional Bayesian Neural Networks (FVBNN) (Sun et al., 2019) use functional prior from fully trained GP (Rasmussen, 2003). This approach somewhat similar to PAD, but PAD accomplishes this in a way which adversarially generates difficult samples for fo and does not require specifying GP kernel and then fitting data to it with
8

---

Under review as a conference paper at ICLR 2021
the limited expessivity of a GP. (Maddox et al., 2019) recently proposed SWAG, which keeps moving average of the weights, and subsequently builds a multivariate Gaussian posterior p() which can be sampled for inference. Deep ensembles (Lakshminarayanan et al., 2017) simply trains an ensemble of N' independent models, and has become a strong standard baseline in the calibration literature. The resulting ensemble of networks are sampled on new inputs and combined in order to achieve set of highly probable solutions from the posterior ((w)D). While simple and effective, it can require excessive computation for large models and datasets. (Dusenberry et al., 2020) proposed RIBNN, which uses shared set of parameters 0 along with multiple sets of rank one parameters r; and S; which can combine via an outer product 0 rs" to create parameter efficient ensemble. Other recent advances include gathering ensemble members by treating depth as random variable and ensembling the latent features through a shared output layer (Antorán et al., 2020). Recently, there has been resurgence of interest in distance sensitive RBF networks (LeCun et al., 1998) for modeling predictive uncertainty (Liu et al., 2020; van Amersfoort et al., 2020) with both Bayesian and deterministic methods, highlighting the importance of distance in uncertainty modeling.
Data Augmentation The data augmentation literature is vast, and as such we will only mention some of the most relevant works here. Adversarial examples (Goodfellow al., 2014b) utilize the gradient of the input with respect to the loss which is then used to create slightly perturbed input which would cause an increase to the loss, and further training with that input. The resulting network should be more robust to such perturbations. (Zhang et al., 2017) proposed Mixup, which creates linear interpolations between natural inputs with the hopes of creating a robust network with linear transitions between classes instead of having hard decision boundary. Subsequent work by (Thulasidasan et al., 2019) showed that mixup training and the soft decision boundary has the effect of improving network calibration.
Set Encoding Methods which operate as a function of sets have been an active topic in recent years. Deep Sets (Zaheer et al., 2017) first proposed model f((x;, X ... xn)) + Rd which passes the input set through a feature extractor, before being aggregated with permutation invariant pooling function, and then decoded through a DNN which projects the set representation to the output space Rd (Lee et al., 2019) proposed to use attention layers to create the set transformer. Our method builds on the idea of set based methods by using the sets towards a new objective of identifying and generating data in underspecified regions where the model is likely overconfident.
6 CONCLUSION
We have proposed a new method of increasing accuracy and calibration of Bayesian neural net- work models which we named Prior Augmented Data (PAD). Our method works by encouraging a Bayesian reversion to the prior beliefs of the labels for inputs in previously unseen or sparse regions of the known data distribution. We have demonstrated through various experiments that our method achieves improvement in likelihood and calibration under shifting data conditions, creating expo- nential improvements in log likelihood in those conditions which the baseline models tend to make the poorest predictions. One interesting direction of future work could be investigate a principled way of doing cross validation when training a network with the objective of being robust to OOD data. The O0D data is by definition unknown, so selecting the best hyperparameter settings is not trivial task. We performed k-fold cross validation using clusters in our regression experiments, but we feel that this topic deserves more study into possible better solutions. Another possible avenue of future research would be look into the effect that different sizes of X may have on the resulting uncertainty predictions. It may be the case that the specific cluster sizes and manifold geometries may require more/less attention and therefore some performance and efficiency gains could be had by giving these regsions the proper amount of attention.
REFERENCES
Javier Antorán, James Allingham. and José Miguel Hernández-Lobate Depth uncertainty in neural networks. Advances in Neural Information Processing Systems, 33, 2020.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural networks. arXiv preprint 155 2015.

---

Under review as a conference paper at ICLR 2021
Dheeru Dua and Casey Graff. UCI machine learning repository, 2017. URL http://arccie ics. uci, edu/m1.
Michael W Dusenberry, Ghassen Jerfel, Yeming Wen, Yi-an Ma, Jasper Snoek, Katherine Heller, Balaji Lakshminarayanan, and Dustin Tran. Efficient and scalable bayesian neural nets with rank- factors, arXiv preprint wrxiv:2005.07/880 2020.
Yarin Gal and Zoubin Ghahramani. Dropout as bayesian approximation: Representing model uncertainty in deep learning, arXiv preprint arXiv:1506.02142: 2015.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural infor- mation processing systems, pp. 2672-2680, 2014a.
lan J.
Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples, 2014b.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In Proceedings of the 34th International Conference on Machine Learning-Voien Vollue 70, pp. 1321-1330. JMLR. org, 2017.
José Miguel Hernández-Lobate and Ryan Adams. Probabilistic backpropagation for scalable learn- ing of bayesian neural networks. In International Conference on Machine Learning, pp. 1861- 1869, 2015.
Geoffrey E Hinton and Sam Roweis. Stochastic neighbor embedding. Advances in neural informa- tion processing systems, 15:857-864, 2002.
Alex Kendall and Yarin Gal. What uncertainties do we need bayesian deep learning for computer vision? In Advances in neural information processing systems, pp. 5574-5584, 2017.
Diederik P Kingma and Max Welling. 13X1v:1312.61 2013.
Auto-encoding variational bayes. arXiv preprint
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Volodymyr Kuleshov, Nathan Fenner, and Stefano Ermon. Accurate uncertainties for deep learning using calibrated regression. arXiv preprint arXiv: 2018.
Solomon Kullback. Information theory and statistics. Courier Corporation, 1997.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predic-
tive uncertainty estimation using deep ensembles. In Advances in neural information processing systems, pp. 6402-641 64133 2017.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):228--2324 1998.
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann. www.sennneecc..... wtt.uum.com//xx///iiss.... 2, 2010.
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-basee permutation-invane invariant neural networks. In Interna- tional Conference on Machine Learning, pp. 3744-3755 3754-3753, 2019.
Jeremiah Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax Weiss, and Balaji Lakshmi- narayanan. Simple and principled uncertainty estimation with deterministic deep learning via distance awareness. Advances in Neural Information Processing Systems, 33, 2020.
Wesley J Maddox. Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainty in deep learning. In Advances in Neural Information Processing Systems, pp. 13153-13164 2019.
10

---

Under review as a conference paper at ICLR 2021
Carl Edward Rasmussen. Gaussian processes in machine learning. In Summer School Machine Learning. pp. 63-71. Springer, 2003.
Jasper Snoek, Yaniv Ovadia, Emily Fertig, Balaji Lakshminarayanan, Sebastian Nowozin, D Scul- ley, Joshua Dillon, Jie Ren, and Zachary Nado. Can you trust your model's uncertainty? eval- uating predictive uncertainty under dataset shift. In Advances in Neural Information Processing Systems, pp. 13969-13980, 2019.
Shengyang Sun, Guodong Zhang, Jiaxin Shi, and Roger Grosse. Functional variational bayesian neural networks. arXiv preprint urXiv:1903.05777 2019.
Sunil Thulasidasan, Gopinath Chennupati, Jeff A Bilmes, Tanmoy Bhattacharya, and Sarah Micha- lak. On mixup training: Improved calibration and predictive uncertainty for deep neural networks. In Advances Neural Information Processing Systems, pp. 11888-13999 2019.
Joost van Amersfoort, Lewis Smith, Yee Whye Teh, and Yarin Gal. Uncertainty estimation using single deep deterministic neural network. In International Conference on Machine Learning, 2020.
Andrew Gordon Wilson and Pavel Izmailov. Bayesian deep learning and probabilistic perspective of generalization. arXiv preprint arXiv:2002.0879), 2020.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander Smola. Deep sets. In Advances in neural information processing systems, pp. 3391- 3401, 2017.
Hongyi Zhang. Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint 17 2017.

---

Under review as a conference paper at ICLR 2021
7 APPENDIX
8 IMPLEMENTATION DETAILS
8.1 KL DIVERGENCES
For regression tasks use an analytic calculation of (xL))p(((( by assuming a prior of (0, 1). In practice, we only want to raise the uncertainty of the prediction while keeping the expressive generalization properties of DNNs, so we use the output of po the mean with the standard deviation set to 1.
(x) == (1) 1)]
(9)
For classification tasks, the prior is uniform categorical distribution. In practice, we found that directly raising the entropy of the output led to a degradation accuracy, so we instead add another output parameter the base model such that instead of outputting C' class logits, output C + and treat the extra output logit as log and use it as a temperature scaling parameter (z/tt) where is the softmax function. We then implement the KL divergence by setting hyperparameter controlling the maximum temperature, and conditionally raising the temperature by the following function.
T (7"h t)²
(10)
where is the weight given by the exponential term before the KIL divergence in 6. In this way, the minimum temperature of 1 is enforced when w = 0, and the maximum temperature of is enforced when 1. Importantly, do transform the logits by when calculating the likelihood loss (first term in equation 8) because we found that this led to larger generalization eπrоr in practice.
8.2 CONVOLUTIONAL ARCHITECTRE DETAILS
Table 6: Convolutional architecture used MNIST experiments
Layers
Conv2d(1, 32) - BatchNorm LeakyReLU MaxPool Dropout
Conv2d(32, 64) - BatchNorm - LeakyReLU - MaxPool - Dropout(O.00)
Conv2d(64, 128) - BatchNorm -> LeakyReLU -> MaxPool - Dropout(O.05)
Conv2d(128, 256) - BatchNorm + LeakyReLU + AvgPool + Dropout (0.05)
FC(128, 128) ReLU - Dropout(O.1)
FC(128, 128) - ReLU - Dropout(O.1)
FC(128, 10)
Table 7: Convolutional architecture used for CIFAR-10 experiments
Layers
Conv2d(3, 32) + BatchNorm LeakyRelLU MaxPool Dropout(O.05)
Conv2d(32, 64) - BatchNorm -> LeakyReLU -> MaxPool - Dropout(O.05)
Conv2d(64, 128) -> BatchNorm - LeakyReLU + MaxPool - Dropout(O.OS)
Con22d(128, 256) -> BatchNorm - LeakyReLU - AvgPool - Dropout(O.0S)
Conv2d(256, 256) BatchNorm LeakyReLU AvgPool Dropout(O.05)
FC(128, 128) ReLU + Dropout(0.1)
FC(128, 128) + ReLU + Dropout(O.1)
FC(128, 10)
12

---

Under review as a conference paper at ICLR 2021
Algorithm 1: PAD for input space
Algorithm 2: PAD for latent features
for DDl = (X) do optimize O...
for X do
optimize 0...
XB 9s [x B XB).
In = [o(Xn) for n = 1, B.
On - fa(Xn) for n2 = 1,..., B.
Zn fo (X,n) for 1, B.
In = fx( for 12 = ,....,B. B.
Zg 90 (Zg ZB).
0++1= c Or- VoLo (DB)
On = fmm) for n2 1,..., B.
optimize o...
Bc+1 t Bl VoLo (DB)
X ~ qá (x B [ nz).
optimize o...
On = (x() for = 1,..., B.
2n So (xm) for 1 ..... B.
@t+1 t Φi- VoLoi
Zg qo (Z ZB).
U. So (Zm:) for n 1, 1,... B.
Φc+1 t Φ: VoLoi
Figure 5: Training steps for fo and go. PAD can work in both the input space (left) and the latent space (right)
8.3 NEW YORK REAL ESTATE
To apply PAD real world problem where the dataset shift is not synthetically created, applied it to regression on New York real estate data. The dataset consists of 12,000 sales records spanning over 12 years. Each instance has total of 667 features including real valued and one-hot categorical features. We used the same base models as outlined in section 4. The regression labels are the price that the house sold for. We train the model to predict log(y) to account for the log-normal distribution of prices and report results based on the log-transformed label. We used the years of 2008-2009 as training/validation
data and evaluated the performance on all following years until 2019. It can be seen that with real temporal distributional shift, PAD models exhibit strong performance. PAD models achieve the best negative log likelihood 35/40 cases. Similar performance can be seen in terms of calibration error, where PAD models show the best calibration error in 35/40 cases.
Negative Log likelihood on 10 years of NY real estate data. The timeline constitutes a real-world temporal datasct shift
Table 8:
Model
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
MC Drop
.0.8
4.570.1
36.90199833
29.55::75.7
20.22:151.8
9.203::20.6
1.0.0:117.0
3.647
0.791
0.856-00.9
MC Drop
0.54910.26
0.66711.00
109610.11
1.16710.30
1.494:0061
1.552:10.59
1.25110.54
0.62611023
0.654_0.26
DE
-0.14+0.42
0.628++1.00
3.9734-4,74
3,633--3.73
4,189-:4.24
4.148-14.46
2.367442.25
1.773+11.57
0,750+0.53
0.735-0.51
DE PAD
0.540::0.04
0.659:00.06
0.971:0.08
1.043:0.10
1.061=0.21
1.217=0.26
0.859:10.11
0.754:20.0
0.656:2007
0.643=0.07
RIBNN
-0.2511005 0.367.10.46
3,941
4.668.12.66
7.047:14.10
6.619.15.01
5.927-12.97
4.548.13.02
1.520
1.085_ 1.66
RIBNN PAD
0.814200001
0.85510.03
1006110.07
1.03020.06
1.020:0.05
1.228:0.09
0,934110.02
0.893:2001
0.844:0001
0.81-00.01
SWAG
4.384-33.41 4.349433.44
4.426H33.4
4,423+-1.39
4.566-3344
4.781#3.466
4-3.5
4.789+33.55
4.818+33511
4.722-1.56
SWAG PAD
0.811:00006
0.891:1005 1.163:0.10
1.254:0.21
1.280=00.30
1.570:20.26
1.148::0.40
1.092:20.77
0.325:10.66
1.006
Table 9: Calibration error across 10 years of NY real estate data. The timeline constitutes a real-world temporal
dataset shift.
Model
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
MC Drop
1.244+1.83
7.8574-5.87
10.34-13.77
12.61-:44.09
10:06=::.05
6.270:43.83
3.186-+2.15
3.039::3.20
2.318:+3.50
3.382=6.17
MC Drop PAD 2.212::1443
1.55
3,58013.03
4,687=3371
5.000:24.31
3.076::3.00
1.40::1441
1.716:1.63
1,149:0,59
0.54
DE
1.713:1089 5.761
10.61:15.45
13.78:16.19
11.98.14455
8.402-12.39
3.696:11.73
5.277
1.85
1.138 0.99
DE PAD
3.408810.48
3.188:10.8
5.13110.70
6.546:11553
6.761:12.40
5.260::2.19
2.64
3.227:1.115
2.757
0.690-0.68
RIBNN
3.041:3.11 7.311--8.33
12.17466.78
14. 19+-5.00
12.04-:4622
5.771
66 4.436+11.55
4.499+11.36
2.387+11.6
0.698 0.82
RIBNN PAD
4.513:50,3
3,463-20.41
4.080:4073
4,449=0.33
4.074=0.42
3.255=9952
2.267=#0.37
2.847=0.34
3,301
3.604=0.25
SWAG
8.207.12.08 7.966.1.208
7.785.11.92
8.100 7.887.1.05
7.604.12.32
7.88311.74 8 1.192
1.42
1.15
SWAG PAD
3,444:10.51
3,20110087 4,385
5,4451:2.21
5.146:3314
4.210::2.80
2.607:11.00
3.20912.34
3,343
3,136=1.27
13

---

Under review as a conference paper at ICLR 2021
Housing
0.1
0.5
Concrete
0.5
i 1.0
Energy
2011
0.5
1.0
Kin8nm
Figure 6: TSNE embeddings for the first half of the UCI regression datatsets used in exper- iments. Both in-distribution D and out-of-distribution data D are shown. The color gradient of D corresponds to different settings for l
14

---

Under review as a conference paper at ICLR 2021
Naval
0.5
1.0
1.5
Power
Wine
1.8
1.5
Yacht
t=
V 0.5
+2 1.0
1= 1
Figure 7: TSNE embeddings for the second half of the UCI regression datatsets used in our experiments. Both in-distribution D and D corresponds different settings for in
data shown. The color gradient of
MC Dropout
Deep Ensembte
SWAG
Rank One BNN
plylk:
plyixi:
log plyjei:
plytat Col. 0.91
MC Dropoct PAD
Deep Ensemble
PAD
SWAG + PAD
Rack Orne BNN PAD
plrlet:
plyjxi:
log plyjvi--11
plyte: Cal. 0,96
Toy regression experiments on an independent run of the toy experiment. These included to show the unpredictable behavior of the un-PAD'd d models which have an undefined behavior outside of the known regions of training data. Comparing the baselines and PAD with figure 1, it can be seen that the baseline models exhibit unpredictable performance in OOD regions of the input space
Figure 8:
15

---

Under review as a conference paper at ICLR 2021
Housing
Concrete
Krrdnm
1o00n Frodu Proqucncyy
SMR SMAGGAND
1885
Y000 PAD
1850
Ireqpercy
frequercy
espected frequency
espected frequency Yscht
Naval
Power
Wine
frequenty
Iouanno paustoe
SWANG Svard PAC
SING
swis PXD
500G
Winn PAD
expected Fequarcy
frequercy
expected frequency
Figure 9: Calibration curves for other models not shown in figure 3.
SMAS
Ceans Ruans
MNSS
SWYG
Acourses
Intensity
Intensity
Inbensity
RUDW
CIFAR1O
SMAS
5AAG
07
1.7
Accuraty
Intensity
Intenstty
Intensity
Figure 10: Model performance varying degrees of shift intensity for MNIST-C and CIFARI0- C. represents the original test set while 6 represents the most extreme level shift. Models which augmented with PAD show comparable performance intensity increases, PAD models exhibit superior performance in terms of negative log likelihood and calibration error. Other models included in figure
the natural test set. As shift
16


Reviews:
[{"review_id": "vLCKXzHYjb1", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Overview:\n\nThe authors propose a data augmentation scheme that generates samples out of distribution and helps with uncertainty estimates. Comparisons are to various bayesian methods in uci regression and mnist/cfar for classification. PAD seems to give some improvements in out of distribution uncertainty quantification.\n\nThe major concern is that the gains seem relatively small, and the objective is ad-hoc. It would be nice to see either more substantial, uniform gains (so that the authors can justify the procedure on the results alone) or more solid conceptual motivation of the method, especially from the Bayesian side. It seems like the motivation and intro is clear, and section 3 onwards becomes very ad-hoc and loses much of this. It would be nice to be convinced that there are a set of assumptions and conditions under which this is the right way to do uncertainty quantification. \n\nPositives:\n\nThe evaluations are extensive, and it's commendable that they include both positive and negative results in their regression evaluations. \n\nUncertainty estimation out of distribution is an important and timely problem.\n\nNegatives:\n\nMinor: I'm not sure why there is a claim that the problems with uncertainty estimation comes from p(theta|D) and not p_theta(y|x). The fact that non-bayesian methods have similar issues with uncertainty quantification would suggest that the latter is certainly an issue.\n\nFigure 1 doesn't seem like a compelling argument for the narrative in the paper. MC dropout and deep ensembles both have decent behavior outside the support (x<0, x>1) but suffer in the gap between 0.25 to 0.75 which is arguably due to overaggressive interpolation.\n\nTerm A in equation 5 is justified as \"generating data where f is overconfident\" but I don't see how this is true. It's just generating data where there's low prediction entropy... this includes areas where it is confident for the right reasons.\n\nSomewhat minor, but the sum of term A and term B seems a bit problematic, since A will be in terms of discrete entropy in classification, and term B is going to be differential entropy in general. Rescaling X also seems like it would arbitrarily shift the weights between AB and C? The weird thresholding on the C penalty for regression problems does not inspire confidence.\n\nOverall, equation 5 gives a sense of a fairly ad-hoc criterion. I'd like to be convinced that this is actually the right way of doing things, especially from a Bayesian perspective.\n\nLooking at equation (7) it seems like learning the distribution of tilde X is alot of work to regularize the KL towards the marginal with a squared-exponential penalty away from the training data. Is it really not possible to post-process the model distribution to achieve the same thing? \n\nThe experiments are extensive, but a bit mixed. The dataset construction for regression seems like it would naturally favor PAD-type methods, because the clustering occurs on the basis of feature distances, and PAD enforces uncertainty based on feature distances (via the squared-exp term in equation 7). In terms of results, I think overall PAD gives gains, but it's not uniform and in cases like SWAG on table 2 seems to hurt more than it helps. \n\nThe corruptions in MNIST / CIFAR must also be pretty aggressive, as the accuracy numbers are quite low for both. Does PAD do similarly well on milder or no distribution shift settings? I am slightly concerned that the evaluations here focus so much on the large distribution setting and that PAD is tuned to that case.\n\nMinor:\n\nInline equation involving sines is missing a closing parenthesis.\n\nNotation. g_phi is a 'generative model' in section 2 but it seems to be the output of an autoencoder in section 3.1 q_phi seems to be the actual generative model?", "extra_scores": {}, "extra_text": {}}, {"review_id": "kEtCHAJFQtW", "reviewer": "AnonReviewer1", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "I think the paper is interesting and well-written. I agree with the mis-calibration can be caused by out-of-distribution data, even though it is still commonly observed without such discrepancy. Addressing OOD data is an important direction, and I think the author proposed a reasonable approach to prevent models from overfitting on data points that are rarely observed during training. However, I believe there are some limitation of the method at the current stage, and the experiments did not fully convince me.\n\nStrength:\n1. The paper addresses an important question of models being over-confident on out-of-distribution data. The method is practical for applications where uncertainty estimation is needed. \n2. The adversarial generation of OOD data is intersting, and the rationale is well explained.\n3. The authors included a good selection of datasets and experiments. The PAD-based methods are also compared to a good variety of baselines.\n\nWeakness:\n1. To determine whether a data point is out-of-distribution, both the adversarial model and main model rely on the L2 distance and the length scale parameter \\ell. My concern here is about 1) the heterogeneity in different dimensions, and 2) \\ell seems particularly important and difficult parameter to tune. I would like the authors to give more details on how it is chosen.\n2. I believe the approach here is to revert to prior when evaluated data points are far away from the observed ones. Thus the difference in accuracy depends on how good the starting prior is, and how much bias the baseline model learned from OOD data. Table 3 shows some of the trade-off, but I think it also be good to show the difference when there's no OOD data, because it's not necessarily known in advance.\n3. Looking at figure 3, I don't think the PAD method shows significant improvement in most of the datasets. Housing seems to be the only one here. In figure 5, it looks like the OOD data are mostly in the convex hull of observed data (at least in this low dimensional embedding). It is unclear how to differentiate those from the region where models should be interpolating. Moreover, all the OOD data are artificially constructed. I think it would be more convincing to test the methods on some OOD data that arises naturally. One such source could be temporal data where distribution could shift over time.\n\nMinor comments:\n1. In section 5 \"excessive computation for large models an datasets\", an->and.\n\n-----------\n\nThank the authors for lot of these responses. I'm still around neutral for this paper, but I will raise my score to marginally above acceptance.\n\n\n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Jg4ahMpSR2", "reviewer": "AnonReviewer2", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "I have read the authors' responses to all reviews and ultimately elected to leave my score as it is (weak accept). I think the empirical results are strong, and while I am not as troubled by the motivation and framing of the work as reviewers 3 and 4, I think their more conceptual and methodological critiques have merit, dampening my enthusiasm for the submission.\n\n-----\n\nThis submission proposes a model-driven data augmentation strategy that aims to improve the calibration and reduce the over-confidence of a variety of Bayesian neural network architectures when dealing with out-of-distribution (OOD) samples. It involves adding a generator network that aims to generate plausible OOD samples during training, along with an objective term that tries to force the predictor network to make high entropy (low confidence) predictions for these samples. The paper does a fairly thorough empirical comparison with ten datasets (eight regression, two image classification) and half a dozen baselines, most of which can be combined with PAD. The results indicate that PAD usually improves both calibration and accuracy by at least a small amount.\n\nThis is a solid paper: the proposed method seems sensible (if pretty complex) and appears to be modestly effective in the included experimental results. The introduction summarizes the paper's contributions as:\n1. It proposes a model-driven data augmentation technique aimed at improving calibration and reducing over-confidence for OOD samples.\n2. It adapts and extends the technique to regression problems, which the paper argues is unprecedented.\n3. It demonstrates empirically that the proposed approach improves the OOD accuracy and calibration of four different strong Bayesian neural net models.\n\nI lack the broad familiarity with the data augmentation literature required to verify claim (2.). I suspect that if this simple claim is true, then it may be trivially so: it's hard to believe that _no one_ has applied data augmentation to regression tasks, so perhaps folks haven't bothered to publish it. The authors can always modify or remove this claim, if needed. The other two contributions seem supported, although the empirical improvements are for the most part small (and probably not statistically significant?). I lean weakly toward acceptance: I would not oppose its inclusion in the ICLR 2021 proceedings, but I wouldn't enthusiastically endorse it. I'll explain below.\n\nThe paper's motivation as laid out in Sections 1 and 2 is sound: calibration and proper quantification of uncertainty are increasingly important in a wide range of applications where machine learning has real world consequences for safety, fairness, etc. What is more, existing techniques based on neural networks (increasingly widespread) do seem to suffer significant flaws, especially exhibiting overconfidence when they should not. The paper offers a diagnosis in the form of a conjecture (Section 2.2): \"failure to revert to the prior $p(\\theta)$ for regions of the input space with insufficient evidence to warrant low entropy predictions.\" Figure 1 effectively visualizes this phenomenon in a toy setting, but no further proof is offered. Further, the assumption that prior reversion is the correct thing to do isn't examined (though that's a basic tenet Bayesian modeling, so we'll set that aside).\n\nThe proposed technique seems sensible, if complicated: add a generator network to produce OOD pseudo-samples during training and penalize the prediction network for making high confidence (low entropy) predictions on these pseudo-samples. We can consider this a form of model-driven (vs. heuristic) data augmentation. The generator loss, given in Equation (5), looks correct to my non-expert eye, and I suspect it's immediately comprehensible to readers familiar with GANs, VAEs, and Bayesian neural nets. The intuition for the OOD samples resonates with me: they should be close enough to real data to be plausible but far enough away that the predictor would be unjustified in assigning a high confidence or departing from the prior.\n\nThe regularization term in Equation (7) is a bit more arcane at first glance, but it's intuitive: the conditional prediction distribution should be close to the prior for pseudo-data points far from the real training distribution. The derivations of the K-L term for regression and categorical classification are given in the appendix, but these aren't critical details for judging the significance of the paper (they're quite straightforward).\n\nThe design of the experiments is sound: they simulate OOD settings by clustering each dataset and using distinct clusters for training and test splits and measure both accuracy and calibration. I don't have an opinion about the choice of the Kuleshov metric for calibration. The chosen baselines look strong, but I am not up-to-date on the relevant literature so I would not be able to identify a non-obvious missing baseline.\n\nThe experimental results are promising. A PAD variant is usually (but not always) the best for each task and metric (exceptions include GP for Naval/accuracy and R1BNN for Power/calibration). Perhaps more important PAD does generally seem to improve both accuracy and calibration across both variants (DE, MC, SWAG, etc.) and datasets. So in other words, if a modeler chooses to use one of the compatible Bayesian neural networks, in most cases they should also use PAD.\n\nThe work and manuscript have a few weaknesses that prevent me from more strongly recommending acceptance. For one, some of the exposition around training is unclear, in particular, how the objectivees in Equations (5) and (8) are combined during training.\n\nI praised the results above, but I think the manuscript's interpretation of its results (Section 4.3) is still more generous than mine. PAD does consistently improve accuracy and calibration, but the margin is sometimes small, raising the question aboout whether the added complexity is worthwhile in all cases. The paper argues that PAD consistently improves the calibration curves in Figure 3, at least for poorly calibrated models, but that does not seem obvious to me. This might be because the curves are somewhat cluttered, but I see a number of exceptions where the PAD look potentially worse: Energy, Naval (SWAG), and Yacht.\n\nI think perhaps the real problem is not the results themselves, which overall are strong, but rather the manuscript's rather cursory discussion of the results and its failure to offer any insights or guidance about, e.g., when PAD should be expected to help (based on task or dataset) or which baselines it works best with.\n\nOne last note: I don't want to over-index on a toy figure included for illustration purposes, but I don't find the results in Figure 1 convincing! Perhaps I am misunderstanding what behavior we desire (if so, please correct me). I agree that the PAD distribution does a better job of capturing the uncertainty in the central low-data area, but at the left- and righthand ends, the baselines actually look preferable in that the uncertainty is often appropriately wider and contains or at least follows the true function. It looks like PAD might be over-regularizing things in these cases.\n\nHere are some actionable suggestions for improvements:\n- Clarify how the objectives are combined and how training proceeds. Consider adding, e.g., an \"Algorithm\" summary figure.\n- Expand the results discussion beyond simply restating the results (which are displayed in the Tables). For the raw accuracy and calibration numbers, perhaps you could compute some summary statistics for the baseline vs. PAD differences, so readers could get a quick sense of whether PAD usually beats the baseline. Maybe also some counts for how often a PAD variant has the best performance.\n- Also in the discussion, try to distill out some illustrative patterns that could be turned into insights or practical guidance.\n- For the calilbration curve plots (Figure 3), consider reducing the clutter by removing redundant curves. For most of the tasks, most baselines (and corresponding PAD variants) are quite similar so perhaps you could show a representative subset for each task (and then put the complete figures in the appendix).\n\nHere's a laundry list of questions:\n- How does training proceed? Is it a typical alternating adversarial optimization, i.e., optimize generator, then discriminator, repeat?\n- What additional computational complexity does PAD introduce during training?\n- Under which conditions PAD should be expected to help most: type or distribution of data, task structure, baseline model, etc.?", "extra_scores": {}, "extra_text": {}}, {"review_id": "Iu3I51wOKSh", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a data augmentation scheme (named PAD) to improve accuracy and calibration of NNs. The idea is to generate OOD data, close to the training data, where the model is overconfident, and force a higher entropy for their corresponding predictions.\n\nThis topic is very relevant to the ICLR community, the paper is clear, and I was excited with the goal in a first place. However, the paper as it is has major drawbacks.\n\n1. The biggest drawback is that the proposed approach is ad-hoc, a heuristic with no guarantees that it will work as desired. In fact, recent work has shown that Data augmentation on top of Ensembles can be harmful, the authors should discuss this in the paper (see [Wen et.al, 2020: Combining Ensembles and Data Augmentation can Harm your Calibration). For this paper to be accepted, the authors should explore the properties of the proposed approach with careful controlled toy scenarios, and bring further insights on when the approach is expected (ideally, guaranteed) to work. \n2. Using PAD on top of other probabilistic approaches destroys the probabilistic interpretation.\n3. Experimental results are extensive, but not convincing: Figure 1 lacks the GP reference, and shows bad performance on the left extreme; the Ablation study suggest that Equation (5) could be simplified; finally results in Table 2-4 suggest that the proposed approach hurts in high-dimensional scenarios (Energy and Kin8nm datasets), the reported numbers also strongly depend on model selection and tuning of PAD and other baselines, information which is currently missing.\n4. The authors do not compare nor mention recent advances on calibrating DNNs, for example:\n\t* (Antoran et.al, 2020) Depth Uncertainty in Neural Networks\n\t* (Liu et.al, 2019) Simple and principled uncertainty estimation with deterministic deep learning via distance awareness\n\t\n\nMore comments:\n\n* the proposed model does not seem to scale to high-dimensions, as \"filling the gaps\" with the OOD data generator becomes infeasible (this is reflected in the tables, where both accuracy and calibration are systematically worse for the high-dim Kin8nm dataset). Up to how many dimensions would this approach be useful?\n\n* the OOD dataset produce an \"equally sized pseudo dataset\". Yet, one might think that the amount of data needed to robustify uncertainty would depend on the manifold geometry.\n* location of OOD samples is chosen as an interpolation of latent representations for the observed data. That means that many generated datapoints will NOT bee out-of-sample.\n\n* From the ablation study (Tables 4 and 5), \"without AB\" gives similar results to Regular (always within the reported error bars of \"regular\". That seems to indicate that terms A and B are not that relevant. Am I missing something?\n\n* Figure 1: the authors should include one column for the GP behavior, since the authors claim that the observed behavior is similar to that. Otherwise, it is unclear by eye what is best. In particular, PAD\n\n* Could the proposed approach suffer from the opposite issue, i.e., deliver too high uncertainty in the augmented OOD data? How do you avoid this issue?\n* How does the proposed approach compare to a DNN whose last layer is GP or Bayesian RBF network? (see http://www.gatsby.ucl.ac.uk/~balaji/udl2020/accepted-papers/UDL2020-paper-009.pdf)\n* The proposed method encourages a reversion to a *specific* prior (0 mean functions)\n\nMinor:\n\n* the authors mention limited expressiveness of GPs, but this is subject to a simple kernel. If the kernel is complicated enough, then GPs are as expressive as we would like to (see equivalences between DNN and GPs in [Neil, 1997] and [Lee et.al, 2017]). Please clarify this statement.\n* Figure 3 is hard to read, I suggest to highlight the PAD curves by changing the color scheme).", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies the problem of uncertainty estimation under distribution shift. The proposed approach (PAD) addresses this under-estimation issue, by augmenting the training data with inputs that the network has unjustified low uncertainty estimates, and asking the model to correct this under-estimation at those augmented datapoints. Results show promising improvement over a set of common benchmark tasks in uncertainty estimation, with comparisons to a number of existing approaches.

All the reviewer agreed that the experiments are well conducted and the empirical results are very promising. However, they also had a shared concern on the justification of the approach. Reviewers are less willing to accept a paper merely for commending its empirical performance.

I share the above concern as the reviewers, and I personally found the presentation of the approach a bit rush and disconnected from the motivation. For example, the current presentation feels like the method is motivated by BNNs but it is not clear to me how the proposed objective connects to the motivation. Also no derivation of the objective is included in either main text or appendix. 

In revision, I would suggest a focus on improving the clarity and theoretical justification of the proposed objective function.
Decision: Reject