PaperlayerView exact revision on arXiv

Source-linked AI summary

What Are Bayesian Neural Network Posteriors Really Like?

Pavel Izmailov, Sharad Vikram, Matthew D. Hoffman, Andrew Gordon Wilson

arXiv:2104.14421v1cs.LGstat.ML

TL;DR

Bayesian neural network posteriors are difficult to characterize because they are high-dimensional and multimodal, while practical approximations may not capture the true posterior. The paper uses full-batch HMC on modern architectures to study these questions, finding strong in-domain performance but poor robustness to distribution shift and distinct predictive behavior from cheaper alternatives.

  • Problem

    The central problem is understanding Bayesian neural network posteriors when practical approximations may not capture their multimodal, high-dimensional structure.

  • Method

    The paper uses full-batch HMC to obtain precise posterior samples and investigate BNN posterior geometry, performance, priors, temperature, and approximate inference.

  • Results

    BNNs achieve strong results, often outperforming SGD-trained models and large deep ensembles, while cheaper alternatives produce predictive distributions distinct from HMC.

  • Takeaways & Limitations

    Accurate posterior approximation can improve predictive accuracy and uncertainty calibration, motivating Bayesian neural networks in practical applications within the evaluated settings.

  • Takeaways & Limitations

    HMC is impractical for most realistic BDL scenarios because image-classification experiments require orders of magnitude more computation than the evaluated baselines.

Abstract

from arXiv · show

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.

1. Introduction

The paper uses full-batch HMC to study Bayesian neural network posteriors that are difficult to approximate, evaluating posterior geometry, predictive performance, and approximate inference alternatives.

  • BNN inference requires approximating a multi-million-dimensional, multimodal posterior under severe computational constraints.
  • Common alternatives have limitations: variational inference is typically unimodal, deep ensembles represent posterior modes, and stochastic MCMC may yield biased posterior expectations.
  • The study applies multi-chain full-batch HMC to investigate posterior geometry, BNN performance, priors, approximate inference, and posterior temperature.
  • The implementation parallelizes HMC over hundreds of TPUs and is released as a resource for evaluating scalable approximate inference methods.
  • BNNs can outperform standard training and deep ensembles, while a single long HMC chain can represent the posterior comparably to multiple shorter chains.
  • Posterior tempering is not needed for near-optimal performance, and the reported cold-posterior effect is largely attributed to data augmentation.

2. Background

Bayesian neural networks average predictions over a posterior distribution of parameters, but the required integral is generally intractable and must be approximated by sampling.

  • Bayesian learning infers a posterior p(w|D) over parameters w after observing data D, rather than selecting one best parameter setting.
  • Bayesian model averaging integrates predictive distributions p(y|x,w) over the parameter posterior p(w|D).
  • The BMA integral is compelling because modern neural-network posteriors can contain complementary solutions represented by different parameter settings.
  • Because neural-network BMA cannot generally be evaluated in closed form, high-dimensional posterior structure makes approximate inference necessary and difficult.
  • MCMC approximates BMA by averaging predictions from samples drawn from the posterior; HMC produces asymptotically exact samples given the unnormalized posterior and its gradient.

3. Related work

Related work emphasizes scalable approximate inference, but direct evidence that these methods match true BNN posteriors on practical models remains limited; this paper instead uses precise HMC sampling.

  • Bayesian deep learning has concentrated on scalable approximate methods, including stochastic variational inference, dropout, and Laplace approaches.
  • Many approximate methods improve predictions or uncertainty estimates without being directly evaluated for matching the true posterior on practical architectures and datasets.
  • The paper constructs a highest-quality posterior approximation while ignoring computational practicality to evaluate approximate inference procedures and foundational questions.
  • Stochastic-gradient MCMC methods are computationally convenient but fundamentally biased because they omit Metropolis-Hastings correction and perturb the stationary distribution through subsampling noise.
  • Unlike earlier full-batch HMC studies using short trajectories and small problems, this work uses practical architectures and datasets with up to 10^5 leapfrog steps per iteration.
  • The paper finds strong performance at temperature 1, extends analysis beyond temperature scaling, and reports that deep ensembles can more closely approximate BMA than standard approximate inference.

4. HMC for deep neural networks

The paper develops a scalable full-batch HMC setup for modern networks, using architecture choices and trajectory tuning to obtain well-mixing posterior samples.

  • Implementation: HMC computation is parallelized over 512 TPUv3 devices using synchronized, data-sharded copies of the Markov chain state.
  • Neural network architectures: The main models are ResNet-20-FRN for CIFAR-10 and CNN-LSTM for IMDB; FRN and smooth Swish activations support independent likelihood terms and smoother posterior geometry.
  • Datasets and Data Augmentation: The experiments omit data augmentation because random augmentations would introduce stochasticity into posterior log-density and gradient evaluation.
  • Trajectory length: Longer HMC trajectories reduce correlation between samples and improve accuracy and log-likelihood, although they increase gradient-evaluation cost.
  • Trajectory length: The recommended trajectory length performs well across accuracy, log-likelihood, and calibration, supporting tens of thousands of gradient evaluations per iteration.
  • Step size: Larger step sizes sharply reduce acceptance probabilities, while higher acceptance rates correspond to higher BMA log-likelihoods in the reported ablation.
  • Number of chains: With a fixed compute budget, using two or three chains is only slightly better than using one, suggesting HMC encounters relatively few obstructive energy barriers.

5. How well does HMC mix?

HMC mixes substantially better in function space than weight space, producing diverse predictions while some parameter directions remain poorly mixed. Its chains navigate complex, apparently multimodal posterior regions, with burn-in requirements differing across datasets.

  • 5.1. R diagnostics: HMC mixes surprisingly well in function space, better than in parameter space, while exploring connected posterior basins with high functional diversity.A single chain can capture prediction diversity nearly as well as multiple chains, despite poor mixing in some weight directions.
  • 5.1. R diagnostics: About 91% of CIFAR-10 and 98% of IMDB posterior-predictive probabilities have R less than 1.1.Function-space R is computed from test-set softmax predictions, whereas weight-space R is computed from raw weights.
  • 5.1. R diagnostics: Some weight parameters have very large R values, indicating that chains fail to mix along particular directions in weight space.Most weight-space diagnostics are small, but the large values show that chains can sample different weight-space distributions.
  • 5.2. Posterior density visualizations: HMC samples occupy complex, high-density regions that appear multimodal in two-dimensional posterior cross-sections.Samples from a single chain fall in seemingly isolated modes, showing that the chain is not restricted to one convex Gaussian-like mode.
  • 5.2. Posterior density visualizations: Samples from independent HMC chains produce qualitatively different posterior surfaces, suggesting imperfect weight-space mixing.The independent-chain visualization appears more regular and symmetric than the same-chain visualization.
  • 5.3. Convergence of the HMC chains: A burn-in of 50 iterations is used for both CIFAR-10 and IMDB because CIFAR-10 performance rises slowly before saturating.IMDB shows no visible performance trend, whereas CIFAR-10 requires a longer burn-in than the approximately 10 iterations suggested by its diagnostic.
  • 5.4. UCI regression: HMC typically matches or outperforms all regression baselines in test RMSE and log-likelihood, often by a significant margin.On Boston, HMC has slightly higher average RMSE than subspace inference and SWAG but significantly better log-likelihood.

6. Evaluating Bayesian neural networks

Using full-batch HMC at temperature T = 1, the paper evaluates BNNs across regression, classification, OOD detection, and distribution shift. HMC generally performs strongly in-distribution but is surprisingly less robust under corruption.

  • 6.1. Regression on UCI datasets: HMC typically outperforms baselines on UCI regression in test RMSE and log-likelihood, except for slightly higher RMSE on Boston.On Boston, HMC still significantly outperforms subspace inference and SWAG in log-likelihood.
  • 6.2. Image Classification on CIFAR: HMC outperforms deep ensembles on OOD detection for CIFAR-100 and SVHN, using predictive confidence and AUC-ROC.It is competitive with ODIN on near-OOD CIFAR-100 but underperforms baselines on far-OOD SVHN.
  • 6.2. Image Classification on CIFAR: HMC BNNs outperform SGD, SGLD, MFVI, and deep ensembles in accuracy and log-likelihood on CIFAR-10, CIFAR-100, and IMDB.The methods are broadly competitive on ECE, while MFVI is particularly well-calibrated on CIFAR-100.
  • 6.2. Image Classification on CIFAR: HMC is the least robust method on CIFAR-10-C, where deep ensembles and SGLD outperform it and a single SGD model performs better at high corruption intensities.The evaluation covers 16 corruption types at intensities from 1 to 5; HMC also shows poor robustness across many noise types.
  • 6.2. Image Classification on CIFAR: For posterior temperature, T = 1 is optimal for IMDB accuracy and log-likelihood, while colder posteriors slightly improve ECE.At T = 1, all three metrics outperform the SGD baseline and a deep ensemble of 10 models.

7. Do we need cold posteriors?

The paper finds that Bayesian neural networks do not need cold posteriors for near-optimal performance, and attributes the reported cold-posterior effect largely to data augmentation.

  • The analysis contrasts with earlier work reporting poor performance at T = 1 and benefits from decreasing temperature.
  • Cold posteriors are not needed for near-optimal Bayesian neural network performance and may hurt performance.
  • At temperature T = 1, the CNN-LSTM BNN outperforms both SGD and a 50-model deep ensemble on test accuracy and log-likelihood.
  • Tempering may still improve robustness to noise under distribution shift, despite not being necessary for standard predictive performance.
  • Role of data augmentation: Data augmentation is identified as the key factor behind the cold-posterior effect; turning it off removes the effect.

8. What is the effect of priors in Bayesian neural networks?

BNN performance is robust across a broad range of prior scales and several prior families, suggesting that architecture matters more than parameter-prior details for specifying useful function-space priors.

  • High-variance Gaussian priors produce strong BNN performance, with results robust to the prior scale.
  • The experiments use single HMC chains with 40 samples for prior-scale comparisons and 80 samples per prior for non-Gaussian comparisons.
  • All considered prior scales outperform SGD training, although very small prior variances can over-regularize and hurt performance.
  • Logistic and mixture-of-Gaussians priors perform relatively similarly to Gaussian priors, with the heavier-tailed logistic prior slightly better on IMDB.
  • Importance of Architecture in Prior Specification: The authors argue that architecture has greater influence than parameter-prior details because it induces the prior over functions.

9. Do scalable BDL methods and HMC make similar predictions?

HMC serves as a reference for comparing predictive distributions: scalable alternatives can generalize well, but their predictions remain distinct from HMC, with deep ensembles and SGLD closer than MFVI.

  • Deep ensembles and SGMCMC can provide good accuracy and calibration while producing predictive distributions that differ from HMC.
  • Deep ensembles are about as close to HMC as standard SGLD and closer than standard variational inference.
  • Comparing the predictive distributions: Agreement measures matching top-1 predictions, while total variation compares class-probability differences between predictive distributions.
  • Advanced SGHMC variants fit HMC's predictive distribution better, but do not necessarily improve accuracy and are less robust to CIFAR-10-C corruptions.
  • Deep ensembles outperform SGLD and SGHMC in total variation on CIFAR-10 and agreement on CIFAR-10-C.
  • MFVI is particularly far from the HMC predictive distribution, whereas SGLD and SGHMC-CLR-Prec more closely match HMC entropy and calibration patterns.

10. Discussion

The paper presents precise posterior sampling as a tool for understanding Bayesian neural networks, while separating its scientific value from HMC’s impracticality for routine use. Its discussion emphasizes both strong in-domain performance and serious limitations under distribution shift and computational cost.

  • Convergence: HMC does not mix perfectly in weight space, although the authors find no evidence of non-mixing in prediction space.Theoretically, undiscovered posterior regions could still contain most of the posterior mass.
  • Should we use Bayesian neural networks?: Precise posterior sampling reveals that Bayesian neural networks can improve predictive accuracy and uncertainty calibration over SGD in varied settings.The authors report that Bayesian neural networks achieved the best results on most problems considered.
  • Should we use HMC in practice?: HMC is generally impractical for realistic Bayesian deep-learning applications because it requires orders of magnitude more computation than the evaluated baselines.The authors frame HMC primarily as a research instrument and call for accurate, scalable approximate inference methods.
  • Implications for approximate inference: Deep ensembles may represent the Bayesian predictive distribution more faithfully than widely used approximate inference methods.The paper cautions against treating deep ensembles simply as a non-Bayesian alternative to variational methods.
  • Distribution shift: Higher-fidelity Bayesian inference surprisingly appears to suffer more under covariate shift, despite approximate Bayesian methods being popular in that setting.This conclusion contrasts with the expectation that Bayesian prediction and uncertainty averaging should improve robustness.

D. Additional Posterior Visualizations

Additional visualizations show that HMC chains can produce similar predictive distributions and navigate complex posterior geometry. On IMDB, the prior dominates the posterior scale, while temperature changes sharpen or soften likelihood cross-sections.

  • Posterior geometry: On IMDB, the posterior log-density is dominated by the prior because the CNN-LSTM has many more parameters than data points.By contrast, the posterior on ResNet-20 is dominated by the likelihood.
  • Posterior geometry: HMC samples fall into three isolated modes in two-dimensional posterior cross-sections, demonstrating navigation of complex geometry.On IMDB, single-chain and three-chain visualizations are qualitatively similar, suggesting better parameter-space mixing than on CIFAR-10.
  • Temperature: At T = 10, the lowest likelihood values have scale 10^3, compared with 10^6 at T = 0.1.Lower temperature produces sharper likelihood cross-sections, while higher temperature makes them softer.
  • Visualization method: The visualizations construct a two-dimensional affine subspace from three HMC samples, orthogonalize its basis, and evaluate a parameter-space grid.This procedure turns selected high-dimensional parameter locations into posterior surface cross-sections.
  • Synthetic regression: Two independent HMC chains produce almost identical predictive distributions on synthetic regression.Both chains show high uncertainty away from observed data, low uncertainty near it, and contain the true function within the 3σ predictive region.

F. BNNs are not Robust to Domain Shift

The appendix finds that HMC-based Bayesian neural networks generalize poorly under corrupted inputs, although lowering posterior temperature can improve robustness in a small MNIST network.

  • CIFAR-10-C: HMC shows poor robustness in accuracy and log-likelihood across CIFAR-10-C corruptions, often losing even to a single SGD solution.Across the corruptions, ECE results are mixed, while SGLD provides the best average robustness.
  • MNIST: At noise scale σ = 3, SGD reaches near 60% accuracy while an individual T = 1 HMC sample achieves around 20%.The comparison uses MNIST corrupted with Gaussian noise and a small fully-connected network.
  • MNIST: Ensembling HMC samples is strong on original test data but cannot close the gap to SGD after Gaussian noise is added.The authors relate this behavior to individually inferior HMC samples produced at a suboptimal training-loss level.
  • Temperature: Lowering HMC temperature can significantly improve robustness on MNIST, but its effect on CIFAR-10-C remains a hypothesis for future work.The appendix explicitly identifies further understanding of BNN robustness as an open direction.

G.1. What causes the difference with Wenzel et al. (2020)?

The comparison with Wenzel et al. attributes the cold-posterior discrepancy primarily to experimental configuration, especially data augmentation. It also highlights that SGLD’s predictive behavior depends on hyperparameters and is inherently biased.

  • SGLD differences: SGLD with a non-vanishing step size samples a perturbed posterior because it omits Metropolis-Hastings correction and includes minibatch noise.Its predictive distribution therefore depends strongly on choices such as batch size and learning rate.
  • Experimental comparison: The authors reproduced strong T = 1 CNN-LSTM performance, obtaining accuracy 0.855 and cross-entropy 0.35.These values exceed the 0.81 accuracy and 0.45 cross-entropy reported from Figure 1 of Wenzel et al.
  • Data augmentation: Turning off data augmentation removes the cold-posterior effect regardless of whether batch normalization or filter response normalization is used.The comparison uses ResNet-20 on CIFAR-10 and temperatures T = 1 and T = 0.1.
  • Predictive confidence: Individual HMC samples can be confidently correct, confidently wrong, or unconfident across input images, while their ensemble is well calibrated.The figure visualizes predicted probabilities assigned to the true class by 100 HMC samples on eight images.
Loading 2104.14421v1…