Approximate Message Passing for Bayesian Neural Networks

TL;DR

This work tackles uncertainty quantification in deep networks by introducing an Approximate Message Passing (AMP) framework for Bayesian neural networks that models the predictive posterior as a factor graph. It derives Gaussian-message passing rules for sums, nonlinearities, and products, and implements a layer-wise, batched inference engine with a diagonal Gaussian posterior q(θ) for scalable learning. The approach is demonstrated on synthetic data and CIFAR-10, where it achieves competitive accuracy and, notably, improved calibration compared to strong baselines such as AdamW and IVON. While scalable to millions of parameters, the method currently lags VI in scale and speed, motivating future integration with residual connections, normalization, and transformer-like components to close the gap with state-of-the-art VI methods.

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.

Figures (4)

• Figure 1: Conceptual vector-valued factor graph for a simple MLP. Each training example has its own "branch" (a copy of the network), while predictions for an unlabeled input ${\bm{x}}$ are computed on a separate prediction branch. All branches are connected by the shared model parameters. Grayed-out variables are conditioned on (observed). Colored arrows indicate the three iteration orders: a forward / backward pass on training examples, and a forward pass for prediction.
• Figure 2: A full FactorGraph models all messages for one batch of training examples. To iterate the FactorGraph, we only need one joint message summarizing the prior and all other examples. When switching to a new batch, we aggregate messages from the previous batch and store them in the Trainer.
• Figure 3: Fitting MLPs of width 16 with increasing depth. Between any linear layers we apply LeakyReLU with a leak of $0.1$. As the depth increases, the network becomes more expressive but harder to fit.
• Figure 4: A deterministic factor corresponding to the $\text{LeakyReLU}_\alpha$ activation function.

Theorems & Definitions (3)

• proof
• proof
• proof