[论文解读] Infinitely Deep Bayesian Neural Networks with Stochastic Differential Equations

Winnie Xu, Ricky T. Q. Chen+2 more|arXiv (Cornell University)|Feb 12, 2021

Gaussian Processes and Bayesian Inference参考文献 40被引用 15

一句话总结

本文提出了一种使用随机微分方程（SDE-BNNs）的无限深度贝叶斯神经网络，其中每个无穷小层的权重不确定性在隐藏单元中诱导出随机轨迹。通过采用一种新型的零方差梯度估计器进行基于梯度的变分推断，该方法实现了任意表达力的近似后验分布，并支持可扩展、内存高效的训练——在MNIST和CIFAR-10上实现了最先进的性能，同时提升了校准能力和鲁棒性。

ABSTRACT

We perform scalable approximate inference in continuous-depth Bayesian neural networks. In this model class, uncertainty about separate weights in each layer gives hidden units that follow a stochastic differential equation. We demonstrate gradient-based stochastic variational inference in this infinite-parameter setting, producing arbitrarily-flexible approximate posteriors. We also derive a novel gradient estimator that approaches zero variance as the approximate posterior over weights approaches the true posterior. This approach brings continuous-depth Bayesian neural nets to a competitive comparison against discrete-depth alternatives, while inheriting the memory-efficient training and tunable precision of Neural ODEs.

研究动机与目标

• 开发一种连续深度的贝叶斯神经网络模型，结合神经ODE和贝叶斯不确定性量化的优势。
• 在无限参数模型中实现可扩展的、基于梯度的变分推断，支持任意表达力的近似后验分布。
• 推导一种新型梯度估计器，其方差随着近似后验趋近真实后验而趋于零。
• 在分类任务中实现具有竞争力的性能，同时保持神经ODE的内存效率和自适应计算特性。

提出的方法

• 将无限深度贝叶斯神经网络的极限建模为描述隐藏单元动态的随机微分方程（SDE）。
• 使用深度神经网络参数化SDE的漂移和扩散函数，以建模时间相关的权重不确定性。
• 应用随机变分推断来优化SDE参数的变分后验，最大化一个修改后的变分下界（ELBO）。
• 引入一种受‘稳住着陆’技巧启发的方差减少梯度估计器，其方差在后验逼近精度提高时趋于零。
• 使用黑箱自适应SDE求解器计算输出层状态，从而在训练过程中保持恒定的内存开销。
• 采用扩展状态维度的参数化方法，以提升表达力并实现更优的后验近似。

实验结果

研究问题

• RQ1 是否可以构建一种连续深度的贝叶斯神经网络，使其后验近似具有任意表达力？
• RQ2 能否为无限参数的贝叶斯模型设计一种梯度估计器，使其方差随着后验近似质量提高而趋于零？
• RQ3 基于SDE的贝叶斯网络是否能在标准基准上实现具有竞争力的性能，同时保持神经ODE的内存效率？
• RQ4 与离散深度的替代方法相比，该方法在分布外数据上的校准性和鲁棒性如何？

主要发现

• SDE-BNN在MNIST上达到99.30% ± 0.09的准确率，在CIFAR-10上达到88.08% ± 1.25的准确率，优于基线ResNet和其他贝叶斯模型，在校准性和鲁棒性方面表现更优。
• 采用‘稳住着陆’（STL）梯度估计器的SDE-BNN在CIFAR-10上的测试准确率达到96.89%，相比标准SVI将负对数似然降低了75%。
• SDE-BNN + STL变体在CIFAR-10上的期望校准误差（ECE）降低至6.44 × 10⁻²，表明预测具有更好的校准性。
• 在CIFAR-10C基准中，该模型在所有损坏严重程度级别下均保持一致的性能，表现出对分布偏移的鲁棒性。
• 所提出的梯度估计器显著降低了方差，即使在高度灵活的后验参数化下也能实现稳定训练。
• 扩展状态维度使模型能够在非冗余维度中学习有意义的表征，同时忽略无关或冗余的维度。