Scaling Up Bayesian Neural Networks with Neural Networks

Zahra Moslemi, Yang Meng, Shiwei Lan, Babak Shahbaba

Abstract: Bayesian Neural Networks (BNNs) offer a principled and natural framework for proper uncertainty quantification in the context of deep learning. They address the typical challenges associated with conventional deep learning methods, such as data insatiability, ad-hoc nature, and susceptibility to overfitting. However, their implementation typically either relies on Markov chain Monte Carlo (MCMC) methods, which are characterized by their computational intensity and inefficiency in a high-dimensional space, or variational inference methods, which tend to underestimate uncertainty. To address this issue, we propose a novel Calibration-Emulation-Sampling (CES) strategy to significantly enhance the computational efficiency of BNN. In this framework, during the initial calibration stage, we collect a small set of samples from the parameter space. These samples serve as training data for the emulator, which approximates the map between parameters and posterior probability. The trained emulator is then used for sampling from the posterior distribution at substantially higher speed compared to the standard BNN. Using simulated and real data, we demonstrate that our proposed method improves computational efficiency of BNN, while maintaining similar performance in terms of prediction accuracy and uncertainty quantification.

Submission Length: Long submission (more than 12 pages of main content)

Previous TMLR Submission Url: https://openreview.net/forum?id=LjeF7DDm8Z&noteId=LjeF7DDm8Z

Changes Since Last Submission: Camera ready version

Video: https://drive.google.com/file/d/10cIAASi5ZN1O5MCm8ojIuTGZhEmygwE4/view

Code: https://github.com/zahraMoslemi/FBNN-codes

Assigned Action Editor: ~Shinichi_Nakajima2

Submission Number: 2529