[Paper Review] Bayesian Neural Network Priors Revisited

Vincent Fortuin, Adrià Garriga-Alonso+6 more|arXiv (Cornell University)|Feb 12, 2021

Adversarial Robustness in Machine Learning82 references24 citations

TL;DR

The paper empirically analyzes SGD-trained neural networks to identify weight distribution patterns and proposes architecture- and task-dependent priors (heavy-tailed for FCNNs; correlated Gaussian for CNNs/ResNets) to improve Bayesian NN performance and address the cold posterior effect.

ABSTRACT

Isotropic Gaussian priors are the de facto standard for modern Bayesian neural network inference. However, it is unclear whether these priors accurately reflect our true beliefs about the weight distributions or give optimal performance. To find better priors, we study summary statistics of neural network weights in networks trained using stochastic gradient descent (SGD). We find that convolutional neural network (CNN) and ResNet weights display strong spatial correlations, while fully connected networks (FCNNs) display heavy-tailed weight distributions. We show that building these observations into priors can lead to improved performance on a variety of image classification datasets. Surprisingly, these priors mitigate the cold posterior effect in FCNNs, but slightly increase the cold posterior effect in ResNets.

Motivation & Objective

• Motivate Bayesian priors that reflect true weight distributions in modern neural networks.
• Characterize empirical weight distributions across FCNNs, CNNs, and ResNets trained with SGD.
• Propose and evaluate alternative priors (heavy-tailed, correlated Gaussian) informed by empirical findings.
• Assess how different priors affect predictive performance and the cold posterior effect.
• Provide a reusable PyTorch library to facilitate Bayesian NN inference with diverse priors.

Proposed method

• Compare isotropic Gaussian priors with heavy-tailed (Laplace, Student-t) and correlated Gaussian priors.
• Empirically analyze weight distributions of SGD-trained FCNNs, CNNs, and ResNets on MNIST and CIFAR-10.
• Fit distributions to SGD solutions and inspect tail behavior and spatial correlations within CNN filters.
• Use SG-MCMC with gradient-guided Monte Carlo and cyclical learning rates to approximate posteriors under different priors.
• Evaluate predictive performance, calibration, and OOD detection across priors and architectures.
• Provide a publicly available PyTorch library for priors in Bayesian NN inference.

Experimental results

Research questions

• RQ1 Do SGD-trained FCNNs, CNNs, and ResNets exhibit non-Gaussian weight distributions that motivate alternative priors?
• RQ2 Are there architecture-specific weight correlation structures (e.g., spatial correlations in CNNs) that should be encoded in priors?
• RQ3 Do heavy-tailed or correlated priors improve Bayesian NN performance compared to isotropic Gaussian priors?
• RQ4 How do chosen priors influence the cold posterior effect across FCNNs, CNNs, and ResNets?
• RQ5 Does data augmentation interact with prior misspecification to affect posterior inferences?

Key findings

• FCNN weights tend to be heavy-tailed, better captured by heavy-tailed priors for improved performance.
• CNN and ResNet weights show significant spatial correlations within convolutional filters.
• Heavy-tailed priors improve Bayesian FCNN performance and reduce the cold posterior effect in FCNNs.
• Correlated Gaussian priors improve Bayesian CNNs and ResNets performance relative to isotropic priors, though effects on the cold posterior vary by architecture and task.
• In FCNNs, heavy-tailed priors can mitigate the cold posterior effect, while in ResNets the effect may be amplified under correlated priors in some datasets.