Model Architecture Adaption for Bayesian Neural Networks

TL;DR

This paper introduces a novel neural architecture search method that optimizes Bayesian Neural Networks for accuracy, uncertainty quantification, and reduced inference latency, outperforming traditional BNNs in efficiency.

Contribution

It proposes a NAS scheme that optimizes BNN placement for uncertainty and accuracy, reducing inference time significantly while maintaining performance.

Findings

Models achieve comparable uncertainty and accuracy to deep ensembles.

Inference runtime reduced by approximately 3 times compared to MCDropout and deep ensemble.

Search effectively identifies optimal Bayesian layer placement.

Abstract

Bayesian Neural Networks (BNNs) offer a mathematically grounded framework to quantify the uncertainty of model predictions but come with a prohibitive computation cost for both training and inference. In this work, we show a novel network architecture search (NAS) that optimizes BNNs for both accuracy and uncertainty while having a reduced inference latency. Different from canonical NAS that optimizes solely for in-distribution likelihood, the proposed scheme searches for the uncertainty performance using both in- and out-of-distribution data. Our method is able to search for the correct placement of Bayesian layer(s) in a network. In our experiments, the searched models show comparable uncertainty quantification ability and accuracy compared to the state-of-the-art (deep ensemble). In addition, the searched models use only a fraction of the runtime compared to many popular BNN baselines, reducing the inference runtime cost by $2.98 \times$ and $2.92 \times$ respectively on the CIFAR10 dataset when compared to MCDropout and deep ensemble.