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: This paper presents numerous empirical facts about the distribution of the weights after training a neural network. The authors considered one FCNN, one CNN, and one ResNet, and computed several statistics on their weights.  The authors also propose, in a Bayesian setup, to compare the performance of trained BNNs with different priors.