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: The paper presents a series of experiments on the cold-posterior effect, and argues that it is broadly prior-dominated,  rather than arising from e.g. the data curation hypothesis of Aitchison (2021).  While it is an interesting series of experiments that clearly deserve to be published, I have a number of concerns about the interpretation of the experiments.  It should be possible to deal with these during the rebuttal phase, at which point I'd be happy to increase my score.