Bayesian (deep) neural networks (BNN) are often more attractive than the vanilla point-estimate deep learning in various aspects including uncertainty quantification, robustness to noise, resistance to overfitting, and more. The variational inference (VI) is one of the most widely adopted approximate inference methods. Whereas the ELBO-based variational free energy method is a dominant choice in the literature, in this paper we introduce a score-based alternative for BNN variational inference. Score-based VI can address the known issue of mode collapsing in ELBO-based VI. Although several score-based VI methods have been proposed in the community, most are not adequate for large-scale BNNs for various computational and technical reasons. We propose a novel scalable VI method where the learning objective combines the score matching loss and the proximal penalty term in iterations, which helps our method avoid the reparametrized sampling, and allows for noisy unbiased mini-batch scores through stochastic gradients. This in turn makes our method scalable to large-scale neural networks including Vision Transformers. On several benchmarks including visual recognition and time-series forecasting with large-scale deep networks, we empirically show the effectiveness of our approach.

翻译：贝叶斯（深度）神经网络在不确定性量化、噪声鲁棒性、抗过拟合及诸多方面通常优于传统的点估计深度学习。变分推断是最广泛采用的近似推断方法之一。尽管基于ELBO的变分自由能方法是文献中的主流选择，但本文为贝叶斯神经网络变分推断引入了一种基于评分的方法。基于评分的变分推断能够解决基于ELBO的变分推断中已知的模态坍塌问题。虽然学界已提出若干基于评分的变分推断方法，但多数方法因计算和技术原因难以适用于大规模贝叶斯神经网络。我们提出了一种新颖的可扩展变分推断方法，其学习目标在迭代过程中结合了评分匹配损失和近端惩罚项，这有助于避免重参数化采样，并通过随机梯度获得含噪声的无偏小批量评分。这进而使我们的方法可扩展至包括视觉Transformer在内的大规模神经网络。在视觉识别、时间序列预测等多项采用大规模深度网络的基准测试中，我们通过实验验证了该方法有效性。