[Paper Review] What Are Bayesian Neural Network Posteriors Really Like?

Pavel Izmailov, Sharad Vikram+2 more|arXiv (Cornell University)|Apr 29, 2021

Gaussian Processes and Bayesian Inference66 references71 citations

TL;DR

The paper uses full-batch Hamiltonian Monte Carlo (HMC) on modern architectures to study true Bayesian neural network posteriors, showing BNNs can outperform standard training and ensembles, with nuanced insights on priors, tempering, domain shift, and comparisons to SGMCMC and deep ensembles.

ABSTRACT

The posterior over Bayesian neural network (BNN) parameters is extremely high-dimensional and non-convex. For computational reasons, researchers approximate this posterior using inexpensive mini-batch methods such as mean-field variational inference or stochastic-gradient Markov chain Monte Carlo (SGMCMC). To investigate foundational questions in Bayesian deep learning, we instead use full-batch Hamiltonian Monte Carlo (HMC) on modern architectures. We show that (1) BNNs can achieve significant performance gains over standard training and deep ensembles; (2) a single long HMC chain can provide a comparable representation of the posterior to multiple shorter chains; (3) in contrast to recent studies, we find posterior tempering is not needed for near-optimal performance, with little evidence for a "cold posterior" effect, which we show is largely an artifact of data augmentation; (4) BMA performance is robust to the choice of prior scale, and relatively similar for diagonal Gaussian, mixture of Gaussian, and logistic priors; (5) Bayesian neural networks show surprisingly poor generalization under domain shift; (6) while cheaper alternatives such as deep ensembles and SGMCMC methods can provide good generalization, they provide distinct predictive distributions from HMC. Notably, deep ensemble predictive distributions are similarly close to HMC as standard SGLD, and closer than standard variational inference.

Motivation & Objective

• Investigate whether true Bayesian posteriors offer advantages over standard training and deep ensembles.
• Assess how well multi-chain and single long HMC chains approximate the posterior.
• Examine the role of posterior temperature (cold vs warm posteriors) in BNN performance.
• Evaluate robustness to domain shift and compare HMC to cheaper inference methods.
• Provide practical guidance on deploying full-batch HMC and share resources for further study.

Proposed method

• Apply full-batch Hamiltonian Monte Carlo to sample from BNN posteriors on architectures like ResNet-20-FRN and CNN-LSTM.
• Parallelize sampling across hundreds of TPU devices in an SPDM setup to handle full-batch gradients.
• Tune HMC hyperparameters (trajectory length, step size, number of chains) to achieve good mixing and acceptance rates.
• Visualize and analyze posterior geometry in weight and function space to understand mixing and mode connectivity.
• Compare HMC with SGLD, MFVI, SGD, and deep ensembles across classification and regression benchmarks.
• Evaluate predictive performance, log-likelihood, calibration error, and out-of-distribution detection metrics.

Experimental results

Research questions

• RQ1 Does a single long HMC chain provide a comparable posterior representation to multiple shorter chains?
• RQ2 Do Bayesian neural networks with true posteriors outperform standard training and deep ensembles in accuracy and calibrated uncertainty?
• RQ3 Is posterior tempering (cold/posterior temperature T<1) necessary for near-optimal performance?
• RQ4 How robust are BMA predictions to different priors (diagonal Gaussian, mixture of Gaussians, logistic) and prior scales?
• RQ5 How do BNNs trained with HMC generalize under domain shift and in out-of-distribution settings compared to alternatives?

Key findings

• BNNs obtained via full-batch HMC can outperform standard training and deep ensembles in accuracy and log-likelihood across CIFAR-10 and IMDB.
• A single long HMC chain can provide a comparable posterior representation to multiple shorter chains for predictive performance.
• Posterior tempering is not required for near-optimal performance at temperature T=1; little evidence for a cold posterior effect under their setup.
• Bayesian model averaging is robust to prior choice and scale; diagonal Gaussian, mixture of Gaussian, and logistic priors yield similar BMA performance as controlling the architecture more than the prior.
• BNNs show strong in-domain performance but surprisingly poor generalization under covariate shift; deep ensembles and SGMCMC can generalize well but yield distinct predictive distributions from HMC.
• Deep ensembles and SGLD can provide predictive distributions that are relatively close to HMC, outperforming standard variational inference in some settings.