Disentangling the Roles of Curation, Data-Augmentation and the Prior in the Cold Posterior Effect

1 Introduction . Deep neural networks have achieved great success in predictive accuracy for supervised learning tasks . Unfortunately , however , they still fall short in giving useful estimates of their predictive uncertainty , i.e . meaningful confidence values for how certain the model is about its predictions ( Ovadia et al. , 2019 ) . Quantifying uncertainty is especially crucial in real-world settings , which often involve data distributions that are shifted from the one seen during training ( Quionero-Candela et al. , 2009 ) . Bayesian deep learning combines deep learning with Bayesian probability theory . Bayesian neural networks ( BNNs ) learn a distribution over model parameters or equivalently sample an ensemble of likely models given the data , promising better generalization performance and principled uncertainty quantification ( robust predictions ) ( Neal , 1995 ; MacKay , 1992 ; Dayan et al. , 1995 ) . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In Bayesian deep learning we either learn a distribution q ( θ ) over models compatible with the data , i.e . q ( θ ) ' p ( θ|D ) , or we sample an ensemble of models θ1 , . . . , θK ∼ p ( θ|D ) from the posterior over likely models p ( θ|D ) ∝ p ( D|θ ) p ( θ ) , ( 1 ) where in the i.i.d . setting p ( D|θ ) = ∏n i=1 p ( yi|xi , θ ) is the likelihood , relating the model we want to learn to the observations D = { ( xi , yi ) } ni=1 , and p ( θ ) is a proper prior , e.g . a Gaussian density . BNN predictions involve model averaging : rather than betting everything on a single point estimate of the parameters , we predict on a new instance x by averaging over all likely models compatible with the data , p ( y|x , D ) = ∫ p ( y|x , θ ) p ( θ|D ) dθ ' K∑ k=1 p ( y|x , θk ) , where θk ∼ p ( θ|D ) . ( 2 ) Equation 2 is also known as the posterior predictive or Bayesian model average . Note that , in practice , solving the integral exactly is impossible . However , we can approximate it via Monte Carlo sampling using an ensemble of models θk ∼ p ( θ|D ) , see Section C in the Appendix for further details . The two main inference paradigms to learn a distribution over model parameters q ( θ ) ' p ( θ|D ) respectively to sample from the posterior θ1 , . . . , θK ∼ p ( θ|D ) are Variational Bayes ( VB ) ( Hinton and Van Camp , 1993 ; MacKay et al. , 1995 ; Barber and Bishop , 1998 ; Blundell et al. , 2015 ) and Markov Chain Monte Carlo ( MCMC ) ( Neal , 1995 ; Welling and Teh , 2011 ; Chen et al. , 2014 ; Ma et al. , 2015 ) . We will focus on MCMC methods as they are simple to implement and can be scaled to large models and datasets when used with stochastic minibatch gradients ( SG-MCMC ) ( Welling and Teh , 2011 ; Chen et al. , 2014 ; Li et al. , 2016 ) . In recent years , the Bayesian deep learning community has developed increasingly accurate and efficient approximate inference procedures for deep BNNs ( cf . references in the paragraph above ) . Despite this algorithmic progress , however , important questions surrounding BNNs remain unanswered to this day . A recent and particularly prominent one concerns the “ cold posterior effect ” ( CPE ) , which describes the observation that the predictive performance of BNNs can be significantly improved if the Bayes posterior is artificially sharpened p ( θ|D ) 1/T using a temperature parameter T < 1 ( Wenzel et al. , 2020 ) . Such cold posteriors sharply deviate from the Bayesian paradigm but are commonly used as heuristics in practice , see Section 2.3 in ( Wenzel et al. , 2020 ) . The CPE is problematic in theory and practice , and since the effect was identified many researchers have proposed hypotheses to explain the phenomenon . There has been an ongoing debate questioning the roles of isotropic Gaussian priors ( Wenzel et al. , 2020 ; Zeno et al. , 2020 ; Fortuin et al. , 2021 ) , the likelihood model ( Aitchison , 2020 ) , inaccurate inference ( Izmailov et al. , 2021 ; Wenzel et al. , 2020 ) , and data augmentation ( Izmailov et al. , 2021 ; Fortuin et al. , 2021 ) . However , despite this intensive research effort the effect remains poorly understood . Contributions : We provide novel and nuanced evidence relevant to existing explanations for the cold posterior effect , disentangling the roles of curation , data augmentation and the prior : • The dataset curation hypothesis of Aitchison ( 2020 ) : we show empirically that the CPE does not arise in a real curated data set but can be produced in a controlled experiment with varying curation strength . • The data augmentation hypothesis of Izmailov et al . ( 2021 ) and Fortuin et al . ( 2021 ) : we show empirically that data augmentation is sufficient but not necessary for the CPE to be present . • The bad prior hypothesis of Wenzel et al . ( 2020 ) : we use a simple experiment evaluating the relative importance of the prior and the likelihood , strongly linking the CPE to the prior . Our results demonstrate how the CPE can arise in isolation from synthetic curation , data augmentation , and bad priors . In fact , we are able to reproduce the cold posterior effect with each of the three factors alone . Cold posteriors observed “ in the wild ” are therefore unlikely to arise from a single cause ; as a result , we do not expect a simple “ fix ” for cold posteriors . 2 Cold Posteriors : Background & Related Work . The “ cold posterior effect ” ( CPE ) states that among all temperized posteriors p ( θ|D ) 1/T the best posterior predictive performance on holdout data is achieved at temperature T < 1 ( Wenzel et al. , 2020 ) . Formally , tempering the posterior corresponds to a 1/T -scaling of the potential energy function U ( θ ) , p ( θ|D ) 1/T ∝ exp ( − 1 T U ( θ ) ) , where U ( θ ) = − n∑ i=1 log p ( yi|xi , θ ) − log p ( θ ) , ( 3 ) i.e. , both the log-likelihood and the log-prior are scaled by 1/T . For T = 1 we have the Bayes posterior , whereas for T → 0 we obtain a sequence of distributions which have their mass more and more confined around the MAP mode of the distribution ( Leimkuhler et al. , 2019 ) . We can thus think of the T → 0 limit of posterior inference as MAP estimation . Note that , besides tempering the posterior as in Equation 3 , one can also temper only the likelihood resp . scale only the log-likelihood , as is commonly done in VB , see Section 2.3 in ( Wenzel et al. , 2020 ) . It is worth pointing out though that both variants are practically equivalent if the prior variance is multiplicative in log-prior pre-factors ( such as for Gaussian priors ) and if sufficiently many prior variances are grid-searched over as part of the inference pipeline , such that for the best performing posterior-tempered model there is a corresponding likelihood-tempered model with variance scaled by 1/T in the grid and vice versa . There are three main building blocks in Bayesian deep learning that may be at fault for the CPE to emerge : ( i ) model misspecification : the likelihood model could be misspecified ( Wenzel et al. , 2020 ; Adlam et al. , 2020 ; Aitchison , 2020 ; Zeno et al. , 2020 ) , ( ii ) bad priors : the priors currently used in deep BNNs may be inadequate ( Wenzel et al. , 2020 ; Fortuin et al. , 2021 ) , or ( iii ) inaccurate inference : the inference method might not yield an accurate enough approximation to the true posterior ( Wenzel et al. , 2020 ; Adlam et al. , 2020 ; Fortuin et al. , 2021 ; Izmailov et al. , 2021 ) . Next we review some of the most prominent hypotheses for the emergence of the CPE . Inaccurate Inference Hypothesis : While recent works on the CPE have taken great care to ensure that their SG-MCMC based inference procedure yields as accurate an approximation to the true posterior as possible , it remains difficult to definitively assess the approximation accuracy without having access to the ( intractable ) true posterior . To rule out obvious problems with the inference mechanism , Wenzel et al . ( 2020 ) proposed a set of diagnostics , based on comparing ensemble statistics to their theoretically known values . We have closely monitored these diagnostics in our experiments , however , despite our own extensive efforts to ensure accurate inference , we can not exclude the possibility that our inference may be inaccurate even though the diagnostics match . To investigate the inference hypothesis further , together with other foundational questions in Bayesian deep learning , Izmailov et al . ( 2021 ) recently applied full-batch Hamiltonian Monte Carlo ( HMC ) 1 , which is considered to be the gold-standard in terms of inference accuracy , to the models of Wenzel et al . ( 2020 ) , showing that with HMC inference and when data augmentation is turned off the CPE disappears . While these results can easily be misconstrued as evidence that the approximate inference is at fault in the CPE , there are good reasons to think otherwise . For instance , Izmailov et al . ( 2021 ) show , using the code of Wenzel et al . ( 2020 ) , that turning off data augmentation alone is sufficient to remove the CPE ( cf . Table 7 in Izmailov et al . ( 2021 ) Appendix G ) . We too can confirm that the CPE does not arise with SG-MCMC based inference applied to Wenzel et al . ( 2020 ) ’ s models when data augmentation is turned off , cf . Figure 1 in Section 4 . From this we can already conclude that either SG-MCMC inference is accurate enough for this specific setting , or inaccurate inference is not necessary for the CPE to emerge2 . A more direct counter argument follows from ( Adlam et al. , 2020 ) , which have demonstrated that there can be a CPE in Gaussian Processes ( GP ) regression , where the posterior has a closed form solution , provided that the aleatoric uncertainty is overestimated . Hence , the CPE can arise in a 1Izmailov et al . ( 2021 ) parallelized the HMC computation over hundreds of Tensor Processing Units ( TPUs ) 2We can also conclude that data augmentation is sufficient for the CPE to arise ( although it is not necessary , as we will see later ) . setting where exact inference is possible . They also provide experimental evidence that the CPE can arise in classification tasks in the infinite-width neural network Gaussian process ( NNGP ) limit . Thus , while inaccurate inference may be sufficient , it does not appear to be necessary for the CPE to emerge . Note that a similar observation was made in ( Grünwald , 2012 ; Grünwald et al. , 2017 ) , which demonstrated benefits of tempering , with T > 1 in their setting , in the context of exact inference . Curation Hypothesis ( model misspecification ) : Aitchison ( 2020 ) devises a theory that attributes the effectiveness of tempering to the fact that standard benchmark datasets such as CIFAR-10 are carefully curated and that we should take this curation into account by tempering BNNs . In Aitchison ( 2020 ) ’ s model of curation ( actual dataset curation may differ ) , a datapoint x is added to the dataset if and only if all S labellers independently agree on the label ys to be assigned to x , while the datapoint is discarded if at least one pair of labellers s , s′ disagrees ys 6= ys′ . The main argument put forward by Aitchison ( 2020 ) is that if we a priori know that the dataset is curated in the sense described above , we should take this into account in our likelihood model . The proposed likelihood to be used in the case of curation should then be of the following form p ( y , x|θ ) ∝ p ( { ys = y } Ss=1|x , θ ) = ∏ s p ( ys = y|x , θ ) . ( 4 ) Thus , assuming that the labellers are i.i.d. , the probability of consensus on label y is p ( y , x|θ ) ∝ p ( y|x , θ ) S , ( 5 ) which corresponds to a cold posterior where only the log-likelihood is re-scaled while the prior is not ( cf . discussion at the beginning of this Section ) . Aitchison ( 2020 ) argues that we should observe S ≈ 1/T , i.e. , the optimal temperature should roughly3 be inversely proportional to the total number of labellers involved in the curation of a datapoint . Note that Aitchison ( 2020 ) ’ s model of curation only includes a data point if all labellers agree , but in practice , datasets are often collected with some tolerance of labeller disagreement , e.g . in that a datapoint is included if a certain fraction of labellers agree . However , a more realistic ( weaker ) model of curation , that filters out fewer “ hard ” instances from the pool of uncurated datapoints , would only give rise to a weaker , less pronounced CPE : if we do not observe the CPE for Aitchison ( 2020 ) ’ s simplistic model of curation , which is the most extreme form of curation imaginable , we would not expect to observe it under more realistic models of curation either . See Section B in the Appendix for additional details on how popular datasets like CIFAR-10 or SVHN were collected . Finally , we would like to point out that Adlam et al . ( 2020 ) made a similar , albeit somewhat more general argument regarding curation resp . mismatch of aleatoric uncertainty . Specifically , they show that the CPE can arise if the model overestimates the aleatoric uncertainty , which is naturally reduced when the dataset is curated . Data Augmentation Hypothesis ( model misspecification ) : Current deep learning practices use a number of techniques , including data augmentation , that technically do not obey the likelihood principle ( see Appendix K in ( Wenzel et al. , 2020 ) for an in-depth discussion of so-called “ dirty likelihoods ” ) . The data augmentation hypothesis specifically says that the CPE is largely an artifact of using data augmentation and that turning off data augmentation is sufficient to remove the CPE . The hypothesis has recently gained traction with both Izmailov et al . ( 2021 ) and Fortuin et al . ( 2021 ) pointing out that turning off data augmentation is sufficient to remove the CPE in Wenzel et al . ( 2020 ) ’ s ResNet CIFAR10 setting ( cf . Table 7 in Izmailov et al . ( 2021 ) Appendix G and Figure A.11 in Fortuin et al . ( 2021 ) Appendix A ) . On the other hand , Wenzel et al . ( 2020 ) ’ s CNN-LSTM IMDB model already had a clean likelihood function and still gave rise to a CPE . From these observations we can already conclude that data augmentation is sufficient but not necessary for the CPE to arise . As the performance of deep neural networks is often significantly better when using some form of data augmentation , it is not really an option to just turn it off in BNNs , while properly accounting for it in Bayesian inference does not seem trivial either . On the one hand , data augmentation affects the data points that enter the likelihood function . However , while data augmentation may increase the 3To obtain an exact correspondence , we would need access to the discarded datapoints in order to be able to marginalize them out , cf . “ marginalise over unknown latents ” in Section 3 in ( Aitchison , 2020 ) for how to account for the discarded images in a proper Bayesian sense amount of data seen by the model , that increase is certainly not equal to the number of times each data point is augmented ( after all , augmented data is not independent from the original data ) . On the other hand , considering data augmentation as a form of regularization ( constraining the classification functions to be invariant to certain transformations ) , one can argue that it should be represented in the prior ( Wilk et al. , 2018 ) . It remains an interesting open problem how to properly account for data augmentation in a Bayesian sense . Bad Prior Hypothesis : Isotropic Gaussian priors are the de facto standard for modern Bayesian neural network inference ( Fortuin et al. , 2021 ) . However , it is questionable whether such simplistic priors are optimal and whether they accurately reflect our true beliefs about the weight distributions . The bad prior hypothesis says that the CPE may only be an epi-phenomenon of a misspecified prior . The underlying argument is as follows : in classic Bayesian learning the number of parameters remains small and the prior is quickly dominated by the data . In contrast , in Bayesian deep learning the model dimensionality is typically on the same order if not larger than the dataset size ( Kaplan et al. , 2020 ) . For such large models the prior will not be dominated by the data and will continue to exert an influence on the posterior . Hence the prior is critical . The hypothesis has already been put forward by Wenzel et al . ( 2020 ) , who reported that the CPE becomes stronger with increasing model dimensionality . It recently got additional empirical support by Fortuin et al . ( 2021 ) , who found that the CPE can be partially alleviated by using heavy-tailed non-Gaussian priors . More specifically , they find that for fully connected neural networks ( FCNNs ) , heavy-tailed priors can both improve predictive performance and alleviate the CPE . For convolutional neural networks ( CNNs ) , the CPE can also be removed with heavy-tailed priors , however , the resulting performance gains are less striking . On the other hand , the performance of CNNs can be improved with correlated priors , although they no longer appear to alleviate the CPE . Finally , we note that the prior that ultimately matters is the prior over functions that is induced when a prior over parameters is combined with the functional form of a neural network architecture ( Wilson and Izmailov , 2020 ; Izmailov et al. , 2021 ) . Still , this does not render the prior over parameters irrelevant , as innocent-looking priors may inadvertently be highly informative , for instance placing large prior mass on undesirable functions .

Summary: This paper studies the so-called cold-posterior effect (CPE), which refers to the poorly understood previous empirical observation that Bayesian posteriors sharpened by a temperature parameter T<1 perform better that the "correct" Bayesian posteriors corresponding to T=1. To me, the paper makes two nice contributions: first, a thorough discussion summarizing all previous explanations for the CPE from the literature; second, an insightful empirical analysis (or re-analysis) of the putative explanations and their impact on the CPE. The main finding was that the CPE does not have a single cause, but can be caused by any of: 1) dataset curation, 2) data augmentation, 3) bad priors. In particular, I think the paper implies that data augmentation and especially the prior have a larger role than previously appreciated, while the impact of dataset curation may not be that significant (although it's non-negligible in some setups, too).