Spatial-Temporal-Fusion BNN: Variational Bayesian Feature Layer

TL;DR

This paper introduces STF-BNN, a scalable Bayesian neural network approach that efficiently scales to large models by Bayesianizing only the first layer, leading to improved performance, robustness, and reduced training costs.

Contribution

The paper proposes a novel spatial-temporal-fusion BNN method that selectively Bayesianizes the first layer, enabling scalable and efficient Bayesian neural networks with theoretical guarantees.

Findings

Achieves state-of-the-art prediction and uncertainty quantification.

Enhances adversarial robustness and privacy preservation.

Reduces training time and memory costs significantly.

Abstract

Bayesian neural networks (BNNs) have become a principal approach to alleviate overconfident predictions in deep learning, but they often suffer from scaling issues due to a large number of distribution parameters. In this paper, we discover that the first layer of a deep network possesses multiple disparate optima when solely retrained. This indicates a large posterior variance when the first layer is altered by a Bayesian layer, which motivates us to design a spatial-temporal-fusion BNN (STF-BNN) for efficiently scaling BNNs to large models: (1) first normally train a neural network from scratch to realize fast training; and (2) the first layer is converted to Bayesian and inferred by employing stochastic variational inference, while other layers are fixed. Compared to vanilla BNNs, our approach can greatly reduce the training time and the number of parameters, which contributes to scale BNNs efficiently. We further provide theoretical guarantees on the generalizability and the capability of mitigating overconfidence of STF-BNN. Comprehensive experiments demonstrate that STF-BNN (1) achieves the state-of-the-art performance on prediction and uncertainty quantification; (2) significantly improves adversarial robustness and privacy preservation; and (3) considerably reduces training time and memory costs.