Scaling Up Bayesian Neural Networks with Neural Networks

TL;DR

This paper introduces a novel Calibration-Emulation-Sampling strategy to significantly improve the computational efficiency of Bayesian Neural Networks while maintaining their ability to quantify uncertainty effectively.

Contribution

The paper presents a new CES framework that accelerates BNN sampling by using an emulator trained on initial samples, reducing computational costs compared to traditional methods.

Findings

CES improves BNN sampling speed substantially.

Maintains prediction accuracy and uncertainty quantification.

Effective on both simulated and real data.

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.