[Paper Review] Combining Model and Parameter Uncertainty in Bayesian Neural Networks

Aliaksandr Hubin, Geir Storvik|arXiv (Cornell University)|Mar 18, 2019

Gaussian Processes and Bayesian Inference57 references4 citations

TL;DR

This paper introduces a scalable variational inference method for Bayesian neural networks (BNNs) that jointly models structural (model) uncertainty and parameter uncertainty by incorporating latent binary variables for weight inclusion. By reparameterizing marginal inclusion probabilities, the approach enables Bayesian model averaging and selection, leading to significantly sparser networks without sacrificing predictive performance, while providing a principled probabilistic interpretation of binary dropout.

ABSTRACT

Bayesian neural networks (BNNs) have recently regained a significant amount of attention in the deep learning community due to the development of scalable approximate Bayesian inference techniques. There are several advantages of using Bayesian approach: Parameter and prediction uncertainty become easily available, facilitating rigid statistical analysis. Furthermore, prior knowledge can be incorporated. However so far there have been no scalable techniques capable of combining both model (structural) and parameter uncertainty. In this paper we introduce the concept of model uncertainty in BNNs and hence make inference in the joint space of models and parameters. Moreover, we suggest an adaptation of a scalable variational inference approach with reparametrization of marginal inclusion probabilities to incorporate the model space constraints. Finally, we show that incorporating model uncertainty via Bayesian model averaging and Bayesian model selection allows to drastically sparsify the structure of BNNs.

Motivation & Objective

• To address the lack of scalable techniques that combine model (structural) and parameter uncertainty in Bayesian neural networks.
• To develop a variational inference framework that enables Bayesian model averaging and selection in BNNs while maintaining scalability.
• To provide a probabilistic interpretation of binary dropout by linking marginal inclusion probabilities to dropout rates.
• To demonstrate that joint modeling of model and parameter uncertainty leads to significantly sparser yet accurate BNNs.
• To evaluate the method on benchmark datasets (MNIST, FMNIST) and compare it with existing BNN inference approaches.

Proposed method

• Introduces latent binary variables to represent inclusion or exclusion of individual weights, enabling explicit modeling of model uncertainty in BNNs.
• Adapts a reparameterization trick to marginal inclusion probabilities, allowing efficient and scalable stochastic variational inference in the joint model-parameter space.
• Uses a fully factorized variational approximation to handle high-dimensional parameter spaces efficiently.
• Applies Bayesian model averaging (BMA) and median probability model (MPM) inference to marginalize over models and parameters.
• Links the resulting inclusion probabilities to binary dropout rates, providing a principled probabilistic foundation for dropout regularization.
• Employs a spike-and-slab-like prior structure on weights, with hierarchical priors on inclusion probabilities to encourage sparsity.

Experimental results

Research questions

• RQ1 Can model uncertainty be effectively and scalably incorporated into Bayesian neural networks alongside parameter uncertainty?
• RQ2 How does joint inference over models and parameters affect network sparsity and predictive performance?
• RQ3 Can marginal inclusion probabilities from the model provide a principled probabilistic interpretation of binary dropout?
• RQ4 How does the proposed method compare to existing BNN inference techniques in terms of sparsity and uncertainty calibration?
• RQ5 Does the method maintain strong predictive performance while achieving significant weight pruning?

Key findings

• The proposed method achieves significant sparsification of BNNs through joint model and parameter uncertainty, reducing the number of active weights without loss of predictive accuracy.
• On MNIST, the method produces sparser models than standard BNNs and Concrete Dropout, while maintaining comparable predictive performance.
• On FMNIST, the method maintains strong uncertainty calibration, with entropy distributions close to uniform in out-of-domain settings, indicating well-calibrated uncertainty estimates.
• The approach outperforms Concrete Dropout in uncertainty calibration on out-of-domain data, showing better handling of predictive uncertainty.
• The marginal inclusion probabilities from the model provide a direct probabilistic interpretation of dropout rates, linking Bayesian inference to a widely used regularization technique.
• The method demonstrates that Bayesian model averaging and median probability model inference lead to meaningful, sparse, and generalizable BNN architectures.