Why Cold Posteriors? On the Suboptimal Generalization of Optimal Bayes Estimates
Abstract
Recent works have shown that the predictive accuracy of Bayesian deep learning models exhibit substantial improvements when the posterior is raised to a 1 /T power with T <
In this work, we explore several possible reasons for this surprising behavior
1. Introduction
Many diﬀerent approaches have been suggested to integrate neural networks with the toolbox of Bayesian inference (MacKay, 1992; Hinton and Camp, 1993; Neal, 1994; Graves, 2011;
Welling and Teh, 2011; Ritter et al., 2018).
In such Bayesian neural networks, we have, after training, not a single set of parameters w (weights), but an (approximate) posterior distribution over w .
Such a posterior distribution is very useful for many potential applications.
For example, it enables uncertainty estimates over the network output; selection of hyper-parameters and models; and guided data collection (active learning).
More recently, there is a growing interest in understanding the properties of tempered posteriors, given by
[Equation]
where ( X , Y ) is the observed dataset (input, output), p ( Y | X , w ) is the likelihood, p ( w ) is the prior, and T is a ‘temperature’ parameter.
This is motivated by the cold posterior eﬀect, studied recently by Wenzel et al.
and observed empirically in numerous previous works on Bayesian deep learning (e.g. Li et al.
; Zhang et al.
Ashukha et al.
The observed eﬀect is that in the posterior predictive distribution on input x
[Equation]
the use of T < 1 (with full tempering) commonly outperforms the standard, optimal
Bayesian estimator with T = 1 in terms of predictive test accuracy.
Note that the 1 /T power in (1) with T < 1 artiﬁcially sharpens the posterior around models with high posterior probability by overcounting the data by a factor of 1 /T (apart from the prior sharpening in the full tempering case).
This result is surprising because, although many works argued that tempered posteriors improve posterior inference for misspeciﬁed models (arguably the case for neural network models), all these works required T > 1 (Jansen, 2013; Grünwald et al., 2017; Miller and Dunson, 2018).
In this work, we examine several possible reasons for the cold posterior eﬀect.
After reviewing related works in Section 2, we consider ﬁrst the possibility of a mismatch between the true prior and the one we use in practice.
We explore two types of prior mismatch.
First, in Section 3 we argue that good priors should be input-dependent, and show how this can lead to the observed cold posterior.
In Section 4 we examine a prior mismatch due to depth, when the assumed model (student) is deeper than the actual model generating the labels (teacher).
Finally, we demonstrate in Section 5 that the cold posterior eﬀect is feasible, even without model mismatch in the case of a single teacher (the standard supervised learning setting) in some relevant cases, e.g. with heavy-tailed posterior distributions.
Existing empirical evidence suggests the hypothesis of input-dependent priors to be most likely.
2. Related works
Recently Wenzel et al.
(2020) conducted an extensive empirical and theoretical study of the cold posterior eﬀect, ruling out several candidate reasons, such as an inaccurate inference method, or the non-formal likelihood functions used in deep learning models (e.g. data augmentation, batch normalization).
The work by Adlam et al.
(2020) studied the cold posterior eﬀect in Gaussian processes (GP) classiﬁcation and regression, and argued that the high quality of the labels in academic benchmarks is not reﬂected in the high observation noise (called aleatoric uncertainty) assumed by the GP model.
The eﬀect of the 1 /T power would be to reduce the observation noise of the model, thus adequating the latter to the data.
A problem in this explanation for the regression case is that it requires an 1 /T factor not only to aﬀect the observation noise, but also the GP kernel — and it remains unclear why this should be the case.
We provide an explanation of this simultaneous rescaling of noise and GP kernel in
Section 4 and in Appendices B.3 and B.4, via a depth mismatch between the data-generating network and the one assumed for the model.
Anonymous (2021) formalizes the argument of Adlam et al.
(2020) for classiﬁcation by showing that the data overcounting implied by the power of 1 /T in the likelihood in (1) is consistent with a training dataset that only includes examples for which 1 /T human data labelers have agreed on the same label.
However, to validate this proposal it remains to be seen if there is a cold-posterior eﬀect when only a single labeler is used to create the dataset.
In Section 3 we provide an alternative explanation to the issue of likelihood normalization, that still allows a single labeler to generate datasets with seemingly cold posteriors.
Additionally, note that all the above works explain the cold posterior eﬀect as the result of model misspeciﬁcation.
However, in Section 5 we show that the eﬀect is also possible without model misspeciﬁcation.
3. Prior mismatch due to input-dependent prior leads to cold posterior
The notion of input-dependent priors is natural if we assume that the input data X (the regressors) inﬂuence our prior knowledge of the model before any observations Y are made.
In such a case, the natural form of Bayes rule is
[Equation]
While in Bayesian deep learning it is common to assume that weights and inputs are independent i.e. p ( w | x ) = p ( w ) , we argue that there is much to gain by allowing a data-dependent prior.
For example, in some cases, it was shown that learning the dependence of w on x with non-predictive losses (using e.g. information theory objectives) reduces (Grandvalet and Bengio, 2005) or obviates (Ji et al., 2019) the need to learn from the targets y .
In this section, we show that if the likelihood function is assumed of the form
[Equation]
then a particular X -dependent prior leads naturally to a tempered posterior.
This form of likelihood function was suggested in Wilson and Izmailov (2020);
Wenzel et al.
However, with the standard Gaussian prior on the weights, this likelihood does not lead to good performance.
We observe training data Y = (cid:2) y 1 , · · · , y N (cid:3) , X = (cid:2) x 1 , · · · , x N (cid:3) , where y i are one-hot vectors representing categorical data.
We denote by y i the index of the single element of y i which is equal to one.
Let { u i ( x , w ) }
Ki =1 be the logits of the output of a neural network with parameters vector w and input data x .
Then the likelihood (4) for K-class classiﬁcation with class label y is
[Equation]
Assuming the following input-dependent prior distribution
[Equation]
and applying Bayes’ rule
[Equation]
the resulting log posterior is
[Equation]
Now we will show that (8) is equivalent to the tempered posterior distribution in the case of classiﬁcation.
In the case of classiﬁcation the commonly used likelihood function is
[Equation]
and the commonly used prior is p ( w ) =
N ( w | 0 , Σ w )) .
The fully tempered log posterior is
[Equation]
so (8) and (10) are equivalent.
Note the particular input-dependent prior we propose in (6), which has a sum over the K indices of y i .
For T < 1 it has a similar ﬂavor as the information theory objectives mentioned above (Grandvalet and Bengio, 2005;
Ji et al., 2019), encouraging the network to be conﬁdent in its most likely prediction when no label is observed.
In
Appendix A we introduce the full generative model.
4. Prior mismatch due to depth leads to cold posterior
In this section, we show that the cold posterior eﬀect also occurs when the network used as a model (‘student network’) is deeper than the data-generating network (‘teacher network’).
To avoid inexact inference methods, we focus on models with closed-form solutions for the Bayesian optimal estimator to demonstrate this eﬀect.
We focus here ﬁrst on deep and wide linear models, and extend the results to neural network Gaussian process (NNGP) and neural tangent kernel (NTK) with ReLU activation functions in Appendices B.3 and B.4.
We observe training data ( x i , y i )
∈ R d × R for i = 1 , . . .
, N from a model
[Equation]
where W l ∈ R n l × n l − 1 , n 0 = d , n L = 1 , and the samples are drawn independently.
We consider the case of an inﬁnite width network.
The prior distribution of W lij is N (cid:16) 0 , σ 2 w n l − 1 (cid:17) i.i.d.
In the linear case, the neural network is equivalent to a linear predictor.
Thus, to calculate the Bayesian predictor we only need the posterior of z = (cid:81)
Ll =1 W l .
To calculate the Bayesian estimator of z we only need to ﬁnd its prior distribution.
According to the central limit theorem (see Appendix B.1 for additional details) the prior distribution of the equivalent linear model z is given by N (cid:16) 0 , σ 2 Lw d I (cid:17) .
Thus, we can calculate the fully tempered posterior distribution of z
[Equation]
The posterior distribution of z is equivalent to the posterior of Bayesian linear regression with prior variance T L σ 2 Lw d and noise variance T σ 2 , therefore the estimator is
[Equation]
As can be seen from (11), changing the temperature T is equivalent to changing the variance of the prior distribution.
We demonstrate the eﬀect of the prior mismatch on the optimal temperature using a wide linear neural network with L t layers as a teacher network, and a wide linear neural network with L s layers as a student network.
Therefore, we get prior mismatch if L t (cid:54) =
L s .
In this case, the tempered posterior can compensate for the prior mismatch, and the optimal temperature is (see Appendix B.2 for the full derivation)
[Equation]
The average MSE is presented in ﬁgure 1 (see Appendix B.5 for implementation details).
The results demonstrate that for σ w > 1 (as used in Wenzel et al.
we get the cold posterior eﬀect when using a student network deeper than the teacher network.
In
Appendix B.3 and B.4 we show that this phenomenon also occurs in non-linear models (neural network Gaussian process (NNGP) and neural tangent kernel (NTK) with ReLU activation functions).
[Figure]
l op t i m a l ( b )   6 Temperature T 1. Lo ss T heo r e t i c a l op t i m a l ( c )
The hypothesis in this section that the cold posterior eﬀect originates from prior mismatch due to depth (the student network is deeper than the teacher network) implies a simple prediction: in a model which does not use bias, changing the variance of the prior is equivalent to changing the depth of the model.
Assuming the hypothesis is correct, we should observe less cold posterior eﬀect when we change the variance of the prior distribution (we do not expect much diﬀerence when we do use bias term).
However, Wenzel et al.
(2020) (in section 5.2) examines the eﬀect of prior variance on the cold posterior eﬀect.
The experiment shows that the cold posterior eﬀect is present for all tested choices of the prior variance.
This seems to contradict the hypothesis’s prediction.
On the other hand, this type of prior mismatch is likely to occur since in Bayesian deep learning we usually use over-parameterized models along with Gaussian prior.
5. Cold posterior eﬀect without model mismatch
In this section, we will demonstrate that when the prior and the posterior distributions are heavy-tailed, using tempered posterior with T < 1 is optimal with high probability, in the setting of supervised learning.
In this setting, we use a given dataset (e.g. CIFAR-10,
ImagNet) to estimate the label for a new unlabeled data, which is equivalent in the Bayesian setting to have a ﬁxed X and only one sample of the teacher network to generate y .
To make a clear demonstration we consider an extreme case of heavy-tailed distribution where for every sampled dataset T < 1 is outperforming the optimal Bayesian estimator.
Let consider the case where the prior and the posterior distributions of the parameter are equal to the distribution of an absolute value of a Cauchy random variable.
In the case of linear model, the Bayesian estimator is ˆ y
= ˆ wx , where
[Equation]
Therefore, for a single dataset the MSE is
[Equation]
[Equation]
Meaning that for every sampled dataset the Bayesian estimator with T < 1 outperforms the optimal Bayesian estimator.
In Appendix C.1 we demonstrate that heavy-tailed prior and posterior distributions exist in the case of narrow and deep neural network and show that with high probability T < 1 is optimal.
In Appendix C.3 we also demonstrate similar results for a bi-modal distribution.
However, if we sample correctly from such posterior distributions (which have rare events with high error) then we expect that only a few predictions would have a high error, while most would be correct.
In contrast, when we examined the SG-MCMC samples (see Appendix C.4) obtained in Wenzel et al.
we found that all of the predictions of the Bayesian ensemble had high error.
This seems to contradict the hypothesis in this section, and suggest that some model misspeciﬁcation is necessary to generate the cold posterior eﬀect.
6. Conclusions
In this work, we suggested and examined several hypotheses to explain the cold posterior eﬀect.
In the case of classiﬁcation we can construct a new likelihood function with tempered soft-max (Hinton et al., 2015) using temperatures T < 1 .
We suggested a speciﬁc inputdependent prior which leads naturally to a tempered posterior.
Next, we considered a prior mismatch due to a deeper student network than the teacher network.
We show that prior mismatch due to network depth results in the cold posterior eﬀect, both for linear and non-linear models.
Lastly, we demonstrated that in the setting of supervised learning (with only one sample of the teacher network) the cold posterior eﬀect is feasible.
We conclude, based on existing empirical observations, that the most likely explanation is the hypothesis of input-dependent prior.
Appendix A. Derivations for section
In this section, we present a generative model which implies an input-dependent prior.
Deﬁning r ( x ) as some distribution (of a single sample x in the input space), and
[Equation]
we ﬁnd the following generative model implies the input dependent prior:
[Equation]
Note that though the marginal prior p ( w ) is dependent on the number of training data N , such dependence was considered in previous works (e.g. Mattingly et al.
Appendix B. Derivations and implementation details for sectionB.1. Calculation of the prior distribution of z - wide linear network
According to an extension of the Central Limit Theorem (CLT), derived in Theorem 4 in
Matthews et al.
(2018), the pre-activation of each layer (for any linear or sub-linear activation function) converge in distribution to a multivariate Gaussian as the layer widths n 1 , ..., n L −are taken to inﬁnity in any order (with n 0 , n L ﬁnite).
For linear activations, this also implies also that z is Gaussian.
The ﬁrst and second moments of z are:
[Equation] [Equation]
Therefore, according to the CLT theorem the prior distribution of z is
[Equation]
B.2. Calculation of the optimal temperature - prior mismatch in wide linear network
The MSE of the Bayesian optimal estimator is
[Equation]
where
[Equation]
The posterior distribution of z is
[Equation]
26) therefore,
[Equation]
As can be seen, the optimal temperature is
[Equation]
B
3. Neural Network Gaussian Process (NNGP)
At initialization, neural networks are equivalent to NNGP inﬁnite-width limit (de G. Matthews et al., 2018).
Due to the non-linearity of NNGP we can no longer obtain the optimal temperature analytically.
However, we numerically demonstrate that prior mismatch due to deeper prior than the true prior leads to cold posterior eﬀect.
We observe training data ( x i , y i )
∈ R d × R for i = 1 , . . .
, N from a model
[Equation]
where (Cho and Saul, 2009; de G. Matthews et al
[Equation] [Equation]
and the samples are drawn independently.
We deﬁne the feature matrix X and the label vector Y as follows
[Equation]
we can write the joint distribution
[Equation]
[Equation]
So the predictive distribution is given by
[Equation]
where
[Equation]
The equivalent of cold posterior in GP is obtain using tempered kernel and tempered noise in the training set (without tempered noise in the test data point) and we obtain
[Equation]
Similar to the previous subsection, we demonstrate the eﬀect of the prior mismatch on the optimal temperature using a Gaussian process with neural network kernel with L t layers as a teacher network, and a Gaussian process with neural network kernel with L s layers as a student network.
The MSE of the Bayesian optimal estimator is
[Equation] [Equation]
The average MSE is presented in ﬁgure 2 (See Appendix B.5 for implementation details).
Similar to the case of a wide linear neural network, the results demonstrate that when using a student network which is deeper than the teacher network we get the cold posterior eﬀect.
2.5 Temperature T 1.
Lo ss
E m p
i r i c a l op t
i m a l ( a ) 0.
2.5 Temperature T2.
Lo ss
E m p
i r i c a l op t
i m a l ( b )   6 Temperature T 2.
Lo ss
E m p
i r i c a l op t
i m a l
[Figure]
B.4. Neural Tangent Kernel (NTK)
Neural networks trained using continuous gradient descent on MSE loss are equivalent to
NTK in the inﬁnite-width limit (Jacot et al., 2018).
Similar to the case of NNGP we numerically demonstrate that prior mismatch due to deeper prior than the true prior leads to cold posterior eﬀect.
In the case of NTK, the derivation of the optimal Bayesian estimator and the MSE is similar to the case of NNGP except for the kernel function.
In this case, the kernel function is given by 3 (Cho and Saul, 2009;
Jacot et al., 2018; Bietti and Mairal, 2019)
[Equation] [Equation]
The average MSE is presented in ﬁgure 3 (see Appendix B.5 for implementation details).
Similar to the case of NNGP, the results demonstrate that when using a student network which is deeper than the teacher network we get the cold posterior eﬀect.
B.5. Implementation details B.5.1. DEEP AND WIDE LINEAR NETWORK
We use a wide linear network with L teacher as a teacher network and a a wide linear network with L student as a student network.
For both we use prior of N (cid:0) 0 , σ 2 w (cid:1) with σ w = 1 . 4 .
2.5 Temperature T 1.
Lo ss
E m p
i r i c a l op t
i m a l ( a ) 0.
2.5 Temperature T 14.
Lo ss
E m p
i r i c a l op t
i m a l ( b )   6 Temperature T
Lo ss
E m p
i r i c a l op t
i m a l
[Figure]
We use as an train set N = 100 samples of ( x , y ) where x ∈ R d and d = 100 , in addition x ∼ N (cid:0) 0 , 1 N I (cid:1) .
The additive noise sampled from N (cid:0) 0 , σ 2 (cid:1) where σ = 1 .
B.5.2. GAUSSIAN PROCESSES
For both NNGP and NTK we use a Gaussian process with kernel of L teacher layers as a teacher and a Gaussian process with neural network kernel of L student layers as a student.
For both teacher and student we use prior of N (cid:0) 0 , σ 2 w (cid:1) with σ w = 2 . 6674 and σ b = 0 .
We use as an train set N = 100 samples of ( x , y ) where x ∈ R d and d = 100 , in addition x ∼ N (cid:0) 0 , 1 N I (cid:1) .
The additive noise sampled from N (cid:0) 0 , σ 2 (cid:1) where σ = 1 .
Appendix C. Derivations for sectionC.1. Scalar Neural Network
Let consider the case of scalar neural network model with an absolute value activation function.
[Equation]
where w ∈ R L and L (cid:29) 1 .
The prior distribution of w is
[Equation]
To calculate the Bayesian predictor we only need the posterior of z = (cid:81)
Lj =1 | w j | .
According to the central limit theorem (CLT), the distribution of z is approximately (see Appendix
C.2 for additional details)
[Equation]
where
[Equation]
For simplicity, we assume that no training data is observed.
In the case of tempered posterior the the Bayesian optimal estimator is
[Equation]
The sampled parameter is most likely to be around the mode of the distribution
[Equation]
Thus, with high probability for a single dataset the optimal temperature is
[Equation]
C.2. Calculation of the prior distribution of z - scalar network
For large enough L the distribution of the sample average of ˜ w j = log ( | w j | ) is approximately
[Equation]
[Equation]
Therefore the distribution of z is approximately
[Equation]
where
[Equation]
C.3. Bi-modal distribution
Another example of distribution where with high probability T < 1 is outperforming the optimal Bayesian estimator is the Bernoulli distribution with p ≈ 1 .
Let consider the following model
[Equation]
In addition, we assume that the prior distribution and the posterior distribution of the parameter are Bernoulli distribution.
Therefore, the tempered posterior distribution is
[Equation]
where
[Equation]
In this case the Bayesian optimal estimator is
[Equation]
Therefore, for a single data set we obtain
[Equation]
For example, if we have a prior distribution with p ( w = 1) = 0 .
99 then the sampled parameter is most likely to be w ∗ = 1 .
Therefore, with high probability the optimal temperature is T = 0 since then q T = 1 and MSE ( w ∗ )
=
σ 2 .
C.4. SG-MCMC samples for T =
In this subsection, we present the SG-MCMC samples obtained in the experiment of ResNet20 on CIFAR-10 (Wenzel et al., 2020).
As can be seen from ﬁgure 4, all of the predictions of the optimal Bayesian ensemble ( T = 1 ) have high error compared to the Bayesian ensemble with T = 0 .
01 and the SGD baseline.
[Figure]
[Figure]
[Figure]