Approximate Message Passing for Bayesian Neural Networks

TL;DR

This paper introduces a novel message passing framework for Bayesian neural networks that effectively models the predictive posterior, handles convolutional architectures, and improves uncertainty calibration, addressing key limitations of existing methods.

Contribution

It presents the first message passing approach for convolutional neural networks that avoids double-counting data and enhances uncertainty quantification in Bayesian neural networks.

Findings

Competitive performance with SOTA baselines on CIFAR-10

Strong correlation (0.9) between credible intervals and true data-generating function

Scales to large MLPs with 5.6 million parameters

Abstract

Bayesian neural networks (BNNs) offer the potential for reliable uncertainty quantification and interpretability, which are critical for trustworthy AI in high-stakes domains. However, existing methods often struggle with issues such as overconfidence, hyperparameter sensitivity, and posterior collapse, leaving room for alternative approaches. In this work, we advance message passing (MP) for BNNs and present a novel framework that models the predictive posterior as a factor graph. To the best of our knowledge, our framework is the first MP method that handles convolutional neural networks and avoids double-counting training data, a limitation of previous MP methods that causes overconfidence. We evaluate our approach on CIFAR-10 with a convolutional neural network of roughly 890k parameters and find that it can compete with the SOTA baselines AdamW and IVON, even having an edge in terms of calibration. On synthetic data, we validate the uncertainty estimates and observe a strong correlation (0.9) between posterior credible intervals and its probability of covering the true data-generating function outside the training range. While our method scales to an MLP with 5.6 million parameters, further improvements are necessary to match the scale and performance of state-of-the-art variational inference methods.