[Paper Review] Dangers of Bayesian Model Averaging under Covariate Shift

Pavel Izmailov, Patrick K. Nicholson+2 more|arXiv (Cornell University)|Jun 22, 2021

Gaussian Processes and Bayesian Inference53 references7 citations

TL;DR

This paper identifies that Bayesian model averaging (BMA) in Bayesian neural networks (BNNs) can severely degrade performance under covariate shift due to unregularized weights corresponding to linearly dependent input features—such as dead pixels in MNIST—being sampled from the prior and amplifying test-time noise. The authors propose novel low-variance priors for these orthogonal directions and sum-constrained priors for non-zero mean corruptions, significantly improving OOD generalization while preserving in-distribution performance.

ABSTRACT

Approximate Bayesian inference for neural networks is considered a robust alternative to standard training, often providing good performance on out-of-distribution data. However, Bayesian neural networks (BNNs) with high-fidelity approximate inference via full-batch Hamiltonian Monte Carlo achieve poor generalization under covariate shift, even underperforming classical estimation. We explain this surprising result, showing how a Bayesian model average can in fact be problematic under covariate shift, particularly in cases where linear dependencies in the input features cause a lack of posterior contraction. We additionally show why the same issue does not affect many approximate inference procedures, or classical maximum a-posteriori (MAP) training. Finally, we propose novel priors that improve the robustness of BNNs to many sources of covariate shift.

Motivation & Objective

• To investigate why high-fidelity Bayesian neural networks (BNNs) with Hamiltonian Monte Carlo (HMC) exhibit poor out-of-distribution (OOD) generalization under covariate shift.
• To explain the root cause of this failure, particularly the role of linear dependencies in input features and lack of posterior contraction in BMA.
• To contrast the robustness of Bayesian model averaging with classical MAP training and stochastic optimization under distributional shift.
• To develop novel priors that improve BNN robustness to multiple types of covariate shift, including data corruption, domain shift, and spurious correlations.
• To demonstrate that the issue is fundamental to BMA in non-linear parameter models with linearly dependent features, not limited to specific architectures or datasets.

Proposed method

• Analyze the impact of linear dependencies in input features (e.g., dead pixels in MNIST) on BNN posterior contraction and prediction robustness under covariate shift.
• Compare Bayesian model averaging (BMA) via HMC with maximum a-posteriori (MAP) training and SGD optimization, focusing on weight initialization and regularization effects.
• Propose a novel prior that constrains the variance of first-layer weights in directions orthogonal to the data manifold, reducing sensitivity to test-time noise.
• Introduce a second prior that constrains the sum of first-layer weights to mitigate performance degradation under non-zero mean corruptions.
• Evaluate the proposed priors on CIFAR-10-C and MNIST under various corruptions, comparing BNNs with HMC to MAP and SGD baselines.
• Conduct loss surface visualization and analysis to assess the flatness of minima and their robustness to data corruption.

Experimental results

Research questions

• RQ1 Why do Bayesian neural networks with high-fidelity HMC inference underperform classical MAP training under covariate shift?
• RQ2 How do linear dependencies in input features—such as dead pixels—lead to poor generalization in BNNs?
• RQ3 Why is Bayesian model averaging more vulnerable than MAP or SGD to test-time distributional shifts?
• RQ4 Can targeted priors improve the robustness of BNNs to multiple types of covariate shift without degrading in-distribution performance?
• RQ5 What role does weight initialization and optimization dynamics play in the robustness of BNNs under distributional shift?

Key findings

• On the pixelate-corrupted CIFAR-10-C dataset, a ResNet-20 BNN with HMC underperforms a MAP solution by 25 percentage points (44% vs. 69% accuracy), despite outperforming MAP on in-distribution data.
• BNNs exhibit poor robustness to covariate shift because weights corresponding to linearly dependent features (e.g., dead pixels) are sampled from the prior and amplify test-time noise, while MAP regularization suppresses such weights.
• The MAP solution produces structured, near-zero weights for dead pixel features, whereas BNN samples show high noise in these weights, making predictions sensitive to input corruption.
• Smaller weight initialization (e.g., U(−0.001, 0.001)) improves SGD’s robustness to covariate shift, suggesting that small initial weights help avoid overfitting to spurious features.
• The proposed low-variance prior for orthogonal directions and sum-constrained prior for non-zero mean corruptions significantly improve OOD generalization in BNNs.
• Loss surface analysis confirms that flat minima (found via SGD) are less affected by corruption than random solutions, and the optimal mode remains stable under increasing corruption intensity.