Model Architecture Adaption for Bayesian Neural Networks
This addresses the problem of prohibitive computation in BNNs for practitioners needing uncertainty quantification, though it is incremental as it builds on existing NAS and BNN methods.
The paper tackles the high computational cost of Bayesian Neural Networks (BNNs) by proposing a novel network architecture search (NAS) that optimizes for both accuracy and uncertainty, reducing inference runtime by up to 2.98x compared to baselines on CIFAR10.
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.