Scalable Approximate Message Passing for Bayesian Neural Networks

Authors: Romeo Sommerfeld, Christian Helms, Ralf Herbrich
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: message passing, bayesian neural networks, uncertainty estimation, factor graphs

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.

Decision: reject
Reviewer scores: 3, 6, 6

Reviews and discussion:
<thread>
creation date:
1730009682583

thread_type:
official_review

<thread_object>
creator: reviewer_yeqc

<content>
summary:
This paper advances message passing (MP) methods for Bayesian neural networks by introducing a novel framework that models the predictive posterior as a factor graph. The key technical contribution is being the first MP method to handle convolutional neural networks while avoiding double-counting of training data, which was a key limitation of previous approaches that led to overconfidence. Their method shows particularly strong performance in data-constrained settings, achieving 94.62% accuracy on MNIST with LeNet-5 using only 640 samples (compared to 22.15% for SGD), while also demonstrating better calibration and out-of-distribution detection than standard approaches. While the method scales to networks with 5.6 million parameters and requires minimal hyperparameter tuning, it currently has higher computational overhead compared to standard training methods and has not yet matched the scale of state-of-the-art variational inference approaches. Nevertheless, the work represents an important step forward in providing more balanced uncertainty estimates, especially in scenarios with limited training data.

strengths:
- First message passing method to handle convolutional neural networks while solving the data double-counting problem
- Exceptional performance with limited data (94.62% MNIST accuracy with 640 samples vs 22.15% for SGD)
- Better calibration and out-of-distribution detection than standard approaches
- Scales to practical network sizes (5.6M parameters) with minimal hyperparameter tuning
- Clear theoretical framework with thorough derivations and implementation details, making it reproducible

weaknesses:
- Significantly slower than standard training methods (96.4s vs 2.3s on GPU for LeNet-5 training)
- Does not yet match the scale of state-of-the-art variational inference methods
- Limited empirical evaluation - only tested on MNIST and synthetic data, lacking results on more complex datasets
- Memory intensive during training, and requiring approximately twice the memory of standard approaches during inference.
- No comparison against other Bayesian methods like variational inference or MCMC in terms of uncertainty quality
- Limited discussion of how the approach handles different neural network architectures beyond MLPs and basic CNNs
- Does not scale to large neural networks beyond 5.6M parameters.

questions:
- Can you provide empirical results on more complex datasets beyond MNIST (e.g., CIFAR-10, ImageNet) to better demonstrate practical applicability and scalability?
- How does your method compare to modern variational inference approaches (like VOGN or IVON) in terms of uncertainty quality, calibration, and computational costs?
- What are the key bottlenecks causing the 40x slower training time compared to SGD, and are there potential optimizations to reduce this gap?
- How does your method handle modern neural network architectures with skip connections, batch normalization, or attention mechanisms - is the factor graph framework adaptable to these components?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730281469617

thread_type:
official_review

<thread_object>
creator: reviewer_mv2h

<content>
summary:
This paper introduces a scalable message-passing (MP) framework for Bayesian neural networks (BNNs), aiming to improve uncertainty quantification in deep learning models. BNNs are known for their potential in high-stakes domains due to their ability to capture predictive uncertainty. Traditional methods, like variational inference (VI), struggle with overconfidence and hyperparameter sensitivity, prompting the development of this MP framework. The authors’ approach utilizes factor graphs to model the predictive posterior, demonstrating that this avoids common issues like double-counting training data, which leads to overconfidence in other MP methods. Their implementation shows better calibration and out-of-distribution detection than standard SGD, particularly in data-constrained settings.

strengths:
* The MP approach achieves superior calibration and out-of-distribution detection, crucial for applications in high-stakes domains where understanding model uncertainty is vital.
* The method shows competitive accuracy, especially in scenarios with limited data, outperforming standard approaches like SGD on tasks with restricted data availability.
* This framework is the first to apply message passing effectively to convolutional neural networks, marking a significant advancement over previous MP implementations. It opens up a relatively unexplored branch which is very important for the advancement of the field.

weaknesses:
* The MP method is computationally demanding, resulting in slower training times compared to standard approaches such as SGD and VI. The method needs substantial memory, especially during training, making it challenging to apply to memory-intensive tasks or very large networks without further optimization.
* Although it scales better than prior MP methods, the approach still lags behind VI in handling large, complex models and datasets.
* The method is based on some assumptions and oversimplification, such as the latent variables have a Gaussian distribution. This might not hold in complex scenarios with multimodal patterns.
* There is no proper time measurement of the method to have a crisp understanding of the computation complexity.

questions:
* Would it be feasible to apply this approach to a larger network or a more challenging dataset?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730609649872

thread_type:
official_review

<thread_object>
creator: reviewer_irhr

<content>
summary:
They propose a scalable message-passing framework for Bayesian neural networks and derive message equations for various factors, which can benefit factor graph modeling across domains. The method is applied for both CNNs and FCNs on MNIST.

strengths:
- Good contribution for the message passing community and especially the experiments on CNNs.

weaknesses:
- The writeup in overall is difficult to follow and the motivation of the work is not clear.
- Poor experimental evaluation, especially since the method is tested only against synthetic data and MNIST. It should be tested on bigger models than LeNet (which is quite outdated), and it should also be tested against CIFAR-10 which is a common benchmark on BNNs.
- All of the proofs are in the Appendix. I would suggest to squeeze something in the main text since there is still space in page 10. Probably the main proof of global minimization objective?

questions:
My main question is why have you not evaluated the proposed approach on CIFAR10 and you evaluated only against MNIST? MNIST is considered as a toy dataset so we do not know if your method generalizes in bigger models even though you scaled the MLP model on MNIST. In the current state the paper is not strongly either theoretically or experimental and therefore it needs more revisions.


Misc: line 178 "predictionsfor" -> "predictions for"
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper develops approximate message passing as an approximate inference technique for Bayesian neural networks. The procedure is fairly involved and requires representing the neural network as a factor graph. Then, the authors perform loopy belief propagation with approximations to the message passing steps. The final results is a diagonal Gaussian approximation to the posterior. The authors show that the method is competitive with variational inference-based methods on a synthetic dataset with MLP and on CIFAR-10 with a small convolutional model.

Strengths:
- The main strength of the paper is in developing novel approximate message passing methodology applicable to (small) neural network models.
- The authors develop and describe a number of tricks and approximations to get the method to run and address numerical instabilities.
- The method is competitive with variational inference in small-scale experiments.

Weaknesses:
- The proposed method is involved, and it is not clear what are the advantages over existing Bayesian methods for neural networks, even variational inference.
- The authors do not do a careful literature review of existing Bayesian methods and do not compare to relevant baselines.
  + Stochastic gradient Monte Carlo methods are popular, scalable, and avoid many limitations of VI [1, 2]. 
  + For the scale of experiments in the paper, full HMC is easily applicable [3]
  + Deep ensembles should be treated as a relevant baseline for uncertainty calibration [4]
  + In the rebuttal, the authors say:
  > We would also like to note that our submission is one of the first - maybe the first - that demonstrates how to scale the training of calibrated Bayesian Neural networks with millions of weights and trained on tens of thousands of training examples.
    * SWAG [5] proposed in 2019 was run on ImageNet (1M+ datapoints, 1k classes) with ResNet-152 models. 
  +  Other scalable Bayesian methods include Laplace approximation-based methods, MC-dropout, etc.
- The empirical evaluation is not sufficient for confirming that the proposed method provides any advantages over baselines.
  + There are only two experiments: synthetic dataset and CIFAR-10. For CIFAR-10, the accuracy is 78% for the best method. The state-of-the-art accuracy on CIFAR-10 is 99+% for pretrained models and 97+% for models trained from scratch. While the authors use a small model, 78% is so far from the state-of-the-art that this result cannot be used to argue for the method providing a practical improvement.
  + Moreover, in the experiments, the proposed MP method loses to the AdamW baseline on accuracy. 
  + While the method provides improved calibration, I don't believe that is the correct target for Bayesian methods. Calibration is a highly-sensitive metric that can be improved with simple interventions [6]. I believe Bayesian methods should aim to improve accuracy, or provide some other benefit over the baselines.
- The method is computationally expensive compared to baselines, at least in the current implementation. The current implementation also does not support normalization layers and residual connections.

Decision recommendation: The paper makes a meaningful contribution in developing an approximate message passing algorithm for Bayesian neural network. In its current form, it is a proof of concept, as the method is not applicable to modern models and does not provide significant improvements in the experiments. The empirical evaluation is very limited and missing relevant baselines, but even compared to the presented baselines the method does not make a significant improvement. Thus, I recommend rejecting the paper in its current form, and encourage the authors to improve the empirical evaluation.

[1] Bayesian Learning via Stochastic Gradient Langevin Dynamics; Max Welling, Yee Whye Teh

[2] Stochastic Gradient Hamiltonian Monte Carlo; Tianqi Chen, Emily B. Fox, Carlos Guestrin

[3] What Are Bayesian Neural Network Posteriors Really Like?; Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, Andrew Gordon Gordon Wilson

[4] Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles; Balaji Lakshminarayanan, Alexander Pritzel, Charles Blundell

[5] A Simple Baseline for Bayesian Uncertainty in Deep Learning; Wesley J. Maddox, Pavel Izmailov, Timur Garipov, Dmitry P. Vetrov, Andrew Gordon Wilson

[6] On Calibration of Modern Neural Networks; Chuan Guo, Geoff Pleiss, Yu Sun, Kilian Q. Weinberger