Bayesian Neural Network Priors Revisited

1 INTRODUCTION . In a Bayesian neural network ( BNN ) , we specify a prior p ( w ) over the neural network parameters , and compute the posterior distribution over parameters conditioned on training data , p ( w|x , y ) = p ( y|w , x ) p ( w ) /p ( y|x ) . This procedure should give considerable advantages for reasoning about predictive uncertainty , which is especially relevant in the small-data setting . Crucially , to perform Bayesian inference , we need to choose a prior that accurately reflects our beliefs about the parameters before seeing any data ( Bayes , 1763 ; Gelman et al. , 2013 ) . However , the most common choice of prior for BNN weights is the simplest one : the isotropic Gaussian . Isotropic Gaussians are used across almost all fields of Bayesian deep learning , ranging from variational inference ( e.g. , Hernández-Lobato & Adams , 2015 ; Louizos & Welling , 2017 ; Dusenberry et al. , 2020 ) , samplingbased inference ( e.g. , Neal , 1992 ; Zhang et al. , 2019 ) , and Laplace ’ s method ( e.g. , Osawa et al. , 2019 ; Immer et al. , 2021b ) , to even infinite networks ( e.g. , Lee et al. , 2017 ; Garriga-Alonso et al. , 2019 ) . It is troubling that no alternatives are usually considered , since better choices likely exist . Indeed , despite the progress on more accurate and efficient inference procedures , in some settings , the posterior predictive distribution of BNNs using Gaussian priors still leads to worse predictive performance than a baseline obtained by training the network with standard stochastic gradient descent ( SGD ) ( e.g. , Zhang et al. , 2019 ; Heek & Kalchbrenner , 2019 ; Wenzel et al. , 2020a ) . Surprisingly , these issues can largely be fixed by artificially reducing posterior uncertainty using “ cold posteriors ” ( Wenzel et al. , 2020a ) . The cold posterior is p ( w|x , y ) 1T for a temperature 0 < T < 1 , where the original Bayes posterior would be obtained by setting T = 1 ( see Eq . 1 ) . Using cold posteriors can be interpreted as overcounting the data and , hence , deviating from the Bayesian paradigm . This should not happen if the prior and likelihood accurately reflect our beliefs . Assuming inference is working correctly , the Bayesian solution , T = 1 , really should be optimal ( Gelman et al. , 2013 ) . Hence , it raises the possibility that either the prior ( Wenzel et al. , 2020a ) or likelihood ( Aitchison , 2020b ) ( or both ) are misspecified . In this work , we study empirically whether isotropic Gaussian priors are indeed suboptimal for BNNs and whether this can explain the cold posterior effect . We analyze the performance of different BNN priors for different network architectures and compare them to the empirical weight distributions of standard SGD-trained neural networks . We conclude that correlated Gaussian priors are better in ResNets , while uncorrelated heavy-tailed priors are better in fully connected neural networks ( FCNNs ) . Thus , we would recommend these choices instead of the widely-used isotropic Gaussian priors . While these priors eliminate the cold posterior effect in FCNNs , they slightly increase the cold posterior effect in ResNets . This provides evidence that the cold posterior effect arises due to a misspecification of the prior ( Wenzel et al. , 2020a ) in FCNNs . In ResNets , it is difficult to draw any strong conclusions about the cold posterior effect from our results . Our observations are compatible with the hypothesis that the cold posterior effect arises in large-scale image models due to a misspecified likelihood ( Aitchison , 2020b ) or due to data augmentation ( Izmailov et al. , 2020 ) , but there could of course be a prior that we did not consider that improves performance and eliminates the cold posterior effect . We make our library available on Github1 , inviting other researchers to join us in studying the role of priors in BNNs using state-of-the-art inference . FC NN Weight histogram 1 0 1 em p. w Gaussian Q-Q Laplace Q-Q 0.25 0.00 0.25 CN N 1 0 1 1 0 1 em p. w 1 0 1 1 0 1 weight value ( w ) Re sN et e ar ly 3 0 3 theoretical w 3 0 3 em p. w 3 0 3 theoretical w Figure 1 : Empirical marginal weight distributions of a layer of FCNNs and CNNs trained with SGD on MNIST , and an early layer of several ResNets trained on CIFAR-10 . We show weight histograms ( left ) and quantile-quantile ( Q-Q ) plots with different distributions ( right ) . The empirical weights are clearly heavier-tailed than a Gaussian ( green line ) , and better fit by a Laplace ( orange line ) . 1.1 CONTRIBUTIONS . Our main contributions are : • An analysis of the empirical weight distributions of SGD-trained neural networks with different architectures , suggesting that FCNNs learn heavy-tailed weight distributions ( Sec . 3.1 ) , while CNN and ResNet weight distributions show significant spatial correlations ( Sec . 3.2 ) . • Experiments in Bayesian FCNNs showing that heavy-tailed priors give better classification performance than the widely-used Gaussian priors ( Sec . 4.2 ) . • Experiments in Bayesian ResNets showing that spatially correlated Gaussian priors give better classification performance than isotropic priors ( Sec . 4.3 ) . • Experiments showing that the cold posterior effect can be reduced by choosing better , heavy-tailed priors in FCNNs , while the cold posterior is slightly increased when using better , spatially correlated priors in ResNets ( Sec . 4 ) . 2 BACKGROUND : THE COLD POSTERIOR EFFECT . When performing inference in Bayesian models , we can temper the posterior by a positive temperature T , giving log p ( w|x , y ) 1T = 1 T [ log p ( y|w , x ) + log p ( w ) ] + Z ( T ) ( 1 ) for neural network weights w , inputs x regression targets or class-labels y , prior p ( w ) , likelihood p ( y|w , x ) , and a normalizing constant Z ( T ) . Setting T = 1 yields the standard Bayesian posterior . The temperature parameter can be easily handled when simulating Langevin dynamics , as used in molecular dynamics and MCMC ( Leimkuhler & Matthews , 2012 ) . In their recent work , Wenzel et al . ( 2020a ) have drawn attention to the fact that cooling the posterior in BNNs ( i.e. , setting T < 1 ) , often improves performance . Testing different hypotheses for potential problems with the inference , likelihood , and prior , they conclude that the BNN priors ( which were Gaussian in their experiments ) are misspecified—at least when used in conjuction with standard neural network architectures on standard benchmark tasks—which could be one of the main causes of the cold posterior effect ( c.f. , Germain et al. , 2016 ; van der Wilk et al. , 2018 ) . Reversing this argument , we can hypothesize that choosing better priors for BNNs may lead to a less pronounced cold posterior effect , which we can use to evaluate different candidate priors . 1https : //github.com/ratschlab/bnn_priors . MIT licensed . 3 EMPIRICAL ANALYSIS OF NEURAL NETWORK WEIGHTS . As we have discussed , standard Gaussian priors may not be the optimal choice for modern BNN architectures . But how can we find more suitable priors ? Since it is hard to directly formulate reasonable prior beliefs about neural network weights , we turn to an empirical approach . We trained fully connected neural networks ( FCNNs ) convolutional neural networks ( CNNs ) and ResNets with SGD on various image classification tasks to obtain an approximation of the empirical distribution of the fitted weights , that is , the distribution of the maximum a posteriori ( MAP ) solutions reached by SGD . If the distributions over SGD-fitted weights differ strongly from the usual isotropic Gaussian prior , that provides evidence that those features should be incorporated into the prior . Hence , we can use our insights by inspecting the empirical weight distribution to propose better-suited priors . Formally , this procedure can be viewed as approximate human-in-the-loop expectation maximization ( EM ) . In particular , in expectation maximization , we alternate expectation ( E ) and maximization ( M ) steps . In the expectation ( E ) step , we infer the posterior p ( w|x , y , θt−1 ) over the weights , w , given the parameters of the prior from the previous step , θt−1 . In our case , we approximately infer the weights using SGD . Then , in the maximization step , we compute new prior parameters θt , by sampling weights w from the posterior computed in the E step , and maximizing the joint probability of sampled weights and data . As y is independent of the prior parameters if the weights are known , the M-step reduces to fitting a prior distribution to the weights sampled from the posterior , that is , Lt ( θ ) = Ep ( w|x , y , θt−1 ) [ log p ( y|x , w ) + log p ( w|θ ) ] = Ep ( w|x , y , θt−1 ) [ log p ( w|θ ) ] + const ( 2 ) θt = arg maxLt ( θ ) . ( 3 ) Intuitively , this procedure allows the prior ( and therefore the posterior ) to assign more probability mass to the SGD solutions , which are known to work well in practice . This is also related to ideas from empirical Bayes ( Robbins , 1992 ) , where the ( few ) hyperparameters of the prior are fit to the data , and to recent ideas in PAC-Bayesian theory , where data-dependent priors have been shown to improve generalization guarantees over data-independent ones ( Rivasplata et al. , 2020 ; Dziugaite et al. , 2021 ) . While such approaches introduce a certain risk of overfitting ( Ober et al. , 2021 ) , we would argue that standard BNNs are typically thought to be underfitting ( Neal , 1996 ; Wenzel et al. , 2020a ; Dusenberry et al. , 2020 ) and that we do not directly fit the prior parameters , but merely draw inspiration for the choice of prior family from the qualitative shape of the empirical weight distributions . We begin by considering whether the weights of FCNNs and CNNs are heavy-tailed , and move on to look at correlational structure in the weights of CNNs and ResNets . Note that in the exploratory experiments here , we used SGD to perform MAP inference with a uniform prior ( that is , maximum likelihood fitting ) . This avoids any prior assumptions obscuring interesting patterns in the inferred weights . These patterns inspired a choice of novel priors , and we evaluated these priors by showing that they improved classification performance ( see Sec . 4 ) . 3.1 FCNN WEIGHTS ARE HEAVY-TAILED . We trained an FCNN ( Fig . 1 , top ) and a CNN ( Fig . 1 , middle ) on MNIST ( LeCun et al. , 1998 ) . The FCNN is a three layer network with 100 hidden units per layer and ReLU nonlinearities . The CNN is a three layer network , with two convolutional layers and one fully connected layer . The convolutional layers have 64 channels and use 3 × 3 convolutions , followed by 2 × 2 max-pooling layers . All layers use ReLU nonlinearities . Networks were trained with SGD for 450 epochs using a learning rate schedule of 0.05 , 0.005 and 0.0005 for 150 epochs each . We can see in Figure 1 that the weight values of the FCNNs and CNNs follow a more heavy-tailed distribution than a Gaussian , with the tails being reasonably well approximated by a Laplace distribution . This suggests that “ true ” BNN priors might be more heavy-tailed than isotropic Gaussians . Next , we did a similar analysis for a ResNet20 trained on CIFAR-10 ( Krizhevsky , 2009 ) ( Fig . 1 , bottom ) . Since this network had many layers , we quantified the degree of heavy-tailedness by fitting the degrees of freedom parameter ν of a Student-t distribution . For ν →∞ , the Student-t becomes Gaussian , so large values of ν indicate that the weights are approximately Gaussian , whereas smaller values indicate heavy-tailed behavior ( see Sec . 4.1 ) . We found that at lower layers , ν was small , so the weights were somewhat heavy-tailed , whereas at higher layers , ν became much larger , so the weights were approximately Gaussian ( Fig . 2a ) . These results are perhaps expected if we assume that the filters have ( using neuroscience terminology ) “ localized receptive fields ” , like those in Olshausen & Field ( 1997 ) . Such filters contain a large number of near-zero weights outside the receptive field , with a number of very large weights inside the receptive field ( Sahani & Linden , 2003 ; Smyth et al. , 2003 ) , and thus will follow a heavy-tailed distribution . As we get into the deeper layers of the networks , receptive fields are expected to become larger , so this effect may be less relevant .

Summary: The work studies prior distributions for Bayesian CNNs. The work report that conventionally used priors e.g., Gaussian poorly fit empirical distributions of the trined weights. The empirical distributions appear to be heavy-tailed and correlated. Thus the paper proposes to use "heavy-tailed priors" for FCNNs and correlated Gaussian priors for CNNs that tend to improve classification performance.