[论文解读] Subspace Inference for Bayesian Deep Learning

Pavel Izmailov, Wesley J. Maddox+4 more|arXiv (Cornell University)|Jul 17, 2019

Gaussian Processes and Bayesian Inference参考文献 48被引用 19

一句话总结

本文提出子空间推理（Subspace Inference, SI），一种通过将近似后验推断限制在高维权重空间的低维子空间中来实现贝叶斯深度学习的方法。通过利用主成分分析（PCA）的SGD轨迹或连接解之间的低损失曲线等子空间，SI 实现了高效且精确的贝叶斯模型平均，相较于完整参数空间方法，在回归和图像分类任务中显著降低了计算成本，同时实现了最先进的不确定性校准和预测性能。

ABSTRACT

Bayesian inference was once a gold standard for learning with neural networks, providing accurate full predictive distributions and well calibrated uncertainty. However, scaling Bayesian inference techniques to deep neural networks is challenging due to the high dimensionality of the parameter space. In this paper, we construct low-dimensional subspaces of parameter space, such as the first principal components of the stochastic gradient descent (SGD) trajectory, which contain diverse sets of high performing models. In these subspaces, we are able to apply elliptical slice sampling and variational inference, which struggle in the full parameter space. We show that Bayesian model averaging over the induced posterior in these subspaces produces accurate predictions and well calibrated predictive uncertainty for both regression and image classification.

研究动机与目标

• 为解决贝叶斯深度学习中高维参数空间带来的可扩展性和精确推断难题。
• 开发一种模块化框架，实现在神经网络权重低维子空间中的有效贝叶斯模型平均。
• 改进预测不确定性估计，特别是在远离训练数据的区域，标准方法常在此失效。
• 证明出人意料的低维子空间（如5维）即可捕捉足够的功能多样性，以实现高质量的贝叶斯平均。
• 提供一种可扩展、模块化且可解释的替代方案，用于深度网络中的完整参数空间贝叶斯推断。

提出的方法

• 使用主成分分析（PCA）的SGD轨迹或独立训练模型之间沿低损失曲线的插值等技术，构建权重空间的低维子空间。
• 在子空间内使用椭圆切片采样（ESS）或变分推断（VI）等方法进行后验推断，这些方法在完整参数空间中难以实现。
• 通过从子空间参数的近似后验中采样，形成贝叶斯模型平均，结合来自多样化高性能模型的预测。
• 利用子空间结构，实现对复杂后验分布的高效探索，这些在完整权重空间中计算上是不可行的。
• 借鉴现有的可扩展贝叶斯基线方法（如SWAG）作为参考，但通过聚焦于信息丰富的子空间，实现对其的改进。
• 将该方法应用于回归和分类任务，在VGG-16到WideResNet-100等模型上，于UCI数据集和CIFAR-10/100上验证性能。

实验结果

研究问题

• RQ1 权重空间的低维子空间是否足以捕捉功能多样性，从而在深度神经网络中实现有效的贝叶斯模型平均？
• RQ2 子空间的选择（随机、SGD轨迹的PCA、低损失曲线）如何影响预测准确性和不确定性校准？
• RQ3 尽管维度显著降低，子空间推理是否能在准确性和不确定性估计方面优于完整参数空间的贝叶斯方法？
• RQ4 与标准贝叶斯深度学习方法相比，子空间推理是否能更好地表征分布外区域的不确定性增长？
• RQ5 子空间推理在多大程度上可实现模块化和跨不同架构与推断技术的适应性？

主要发现

• 在3600万参数的WideResNet上，使用2维曲线子空间（连接两个独立训练的模型）在CIFAR-100上实现了81.55%的测试准确率和0.6493的负对数似然（NLL）。
• PCA子空间（10维）实现了80.54%的准确率和0.6652的NLL，优于随机子空间（80.17%准确率，0.6858 NLL），并持续优于基线方法。
• 在PCA子空间中使用椭圆切片采样（ESS）比SWAG更准确地捕捉了后验形状，后者在第一主成分方向上高估了方差。
• 在子空间中使用变分推断（VI）相比SWAG提供了更好的后验拟合，因为它直接近似后验分布，而非SGD轨迹。
• 子空间推理在CIFAR-10和CIFAR-100基准上持续优于MC-Dropout、温度缩放和KFAC-Laplace方法。
• 即使在3600万参数网络的CIFAR-100上，5维子空间也能通过贝叶斯模型平均实现显著的性能提升，证明了低维子空间的有效性。