[论文解读] Tractable Function-Space Variational Inference in Bayesian Neural Networks

Tim G. J. Rudner, Zonghao Chen+2 more|arXiv (Cornell University)|Dec 28, 2023

Adversarial Robustness in Machine Learning被引用 6

一句话总结

本文提出了一种可处理的函数空间变分推断（FSVI）方法，用于贝叶斯神经网络，通过直接对函数而非参数的后验分布进行建模，实现了可扩展的不确定性估计。通过引入一种简单的函数空间KL散度估计器，FSVI能够有效融入先验信息，并在多种任务上实现了最先进的预测不确定性与准确性表现，包括一个关键安全性的医疗诊断基准任务。

ABSTRACT

Reliable predictive uncertainty estimation plays an important role in enabling the deployment of neural networks to safety-critical settings. A popular approach for estimating the predictive uncertainty of neural networks is to define a prior distribution over the network parameters, infer an approximate posterior distribution, and use it to make stochastic predictions. However, explicit inference over neural network parameters makes it difficult to incorporate meaningful prior information about the data-generating process into the model. In this paper, we pursue an alternative approach. Recognizing that the primary object of interest in most settings is the distribution over functions induced by the posterior distribution over neural network parameters, we frame Bayesian inference in neural networks explicitly as inferring a posterior distribution over functions and propose a scalable function-space variational inference method that allows incorporating prior information and results in reliable predictive uncertainty estimates. We show that the proposed method leads to state-of-the-art uncertainty estimation and predictive performance on a range of prediction tasks and demonstrate that it performs well on a challenging safety-critical medical diagnosis task in which reliable uncertainty estimation is essential.

研究动机与目标

• 为解决贝叶斯神经网络中参数空间变分推断的局限性，后者常导致不确定性量化不佳且性能低于非贝叶斯方法。
• 实现将关于数据生成过程的任务特定先验知识融入贝叶斯神经网络推断中。
• 开发一种可扩展且可处理的函数空间变分推断方法，以在高维和过参数化设置下支持可靠的预测不确定性估计。
• 在标准基准和一个真实世界的安全关键性医疗诊断任务上评估该方法，证明其在分布偏移下的鲁棒性。

提出的方法

• 该方法在函数空间中明确定义贝叶斯推断，目标是通过网络参数诱导出的函数后验分布。
• 提出一种新颖且简单的函数空间分布之间KL散度的估计器，使随机变分推断成为可能。
• 该方法在每个梯度更新步骤中使用上下文分布 $p_{\mathbf{X}_{\mathcal{C}}}$ 采样输入点，以近似变分目标中的上确界。
• 对于图像输入，上下文分布被定义为基于训练数据像素构建的灰度图像上的均匀分布；对于表格数据，则使用经验输入边界。
• 该方法支持灵活的先验设计，包括可鼓励在分布外区域产生高不确定性的一类先验。
• 使用标准深度学习框架实现，采用Adam优化，且与ResNet和MLP等标准架构兼容。

实验结果

研究问题

• RQ1 能否使函数空间变分推断在高维、过参数化的神经网络中实现可处理且可扩展？
• RQ2 与参数空间方法相比，将任务特定先验知识融入函数空间先验是否能提升预测不确定性与准确性？
• RQ3 在分布偏移下，该方法的性能如何，特别是在安全关键应用中？
• RQ4 函数空间变分推断能否在不确定性估计和预测性能方面优于贝叶斯与非贝叶斯基线方法？

主要发现

• FSVI在多个回归与分类基准任务中均实现了最先进的预测性能与不确定性估计，包括1D回归和Two Moons数据集。
• 在Snelson 1D回归任务中，FSVI生成的预测分布能良好拟合训练数据，同时在远离训练点的区域表现出高不确定性，符合预期。
• 在糖尿病视网膜病变诊断任务中，基于FSVI的可选预测（专家转诊）方法优于基线方法，在低转诊率下仍保持高准确率。
• 该方法对分布偏移表现出鲁棒性，模型能为分布外样本分配高认知不确定性，从而实现有效的可选预测。
• 与标准参数空间变分推断和深度集成相比，使用任务特定先验的FSVI显著提升了不确定性校准与预测准确性。
• 所提出的函数空间KL散度估计器实现了稳定且可扩展的训练，使函数空间推断在实际应用中变得切实可行。