[论文解读] Function Space Particle Optimization for Bayesian Neural Networks

Ziyu Wang, Tongzheng Ren+2 more|arXiv (Cornell University)|Feb 26, 2019

Adversarial Robustness in Machine Learning被引用 14

一句话总结

该论文提出了一种新型变分推断方法——函数空间粒子优化（f-SVGD），用于贝叶斯神经网络，其在回归函数空间而非权重空间中直接执行粒子优化。通过避免在过参数化模型中出现的梯度混合与粒子坍缩问题，f-SVGD在预测性能、对抗鲁棒性以及强化学习任务中均表现出色，在多个基准测试中超越了强基线模型。

ABSTRACT

While Bayesian neural networks (BNNs) have drawn increasing attention, their posterior inference remains challenging, due to the high-dimensional and over-parameterized nature. To address this issue, several highly flexible and scalable variational inference procedures based on the idea of particle optimization have been proposed. These methods directly optimize a set of particles to approximate the target posterior. However, their application to BNNs often yields sub-optimal performance, as such methods have a particular failure mode on over-parameterized models. In this paper, we propose to solve this issue by performing particle optimization directly in the space of regression functions. We demonstrate through extensive experiments that our method successfully overcomes this issue, and outperforms strong baselines in a variety of tasks including prediction, defense against adversarial examples, and reinforcement learning.

研究动机与目标

• 解决过参数化贝叶斯神经网络中权重空间粒子优化的失效模式，即粒子坍缩至单一模式的问题。
• 克服高维多模态权重空间后验分布中的维度灾难与模式坍缩问题。
• 通过将优化从权重空间转移到函数空间（其中粒子表示回归函数）来提升后验近似质量。
• 通过仅在粒子之间交换函数空间误差信号，实现可扩展且可并行化的推理，适用于大规模模型。
• 在多种任务中实现一致的性能提升，包括回归、对抗防御与强化学习。

提出的方法

• 通过将每个粒子表示为由神经网络权重参数化的回归函数，在函数空间中执行粒子优化。
• 使用函数值核，基于函数相似性而非权重距离来计算粒子间的相互作用。
• 推导出一种粒子更新规则，通过利用 f^j 关于 θ^j 的雅可比矩阵，将来自粒子 j 的梯度应用于粒子 i，从而避免梯度混合。
• 实现更新规则的随机小批量近似，以支持大规模数据集与模型的扩展。
• 利用模型并行性：仅在粒子之间交换函数空间误差信号，从而降低通信开销。
• 采用重参数化技巧，计算核函数相对于粒子权重的梯度，从而实现高效优化。

实验结果

研究问题

• RQ1 在函数空间中进行粒子优化，能否缓解过参数化BNN中权重空间粒子方法所出现的模式坍缩与梯度混合问题？
• RQ2 与权重空间方法相比，在函数空间中优化是否能带来更好的泛化性能与不确定性校准？
• RQ3 在回归、对抗鲁棒性与强化学习等下游任务中，f-SVGD相较于强基线模型表现如何？
• RQ4 函数空间核的设计对BNN中收敛性与近似质量有何影响？
• RQ5 f-SVGD能否通过模型并行性实现高效扩展，同时保持极低的通信开销？

主要发现

• f-SVGD成功避免了在过参数化模型中困扰权重空间粒子优化的粒子坍缩问题。
• 在回归任务中，f-SVGD的测试误差更低，且不确定性校准优于SVG D及其他基线方法。
• 在对抗防御任务中，f-SVGD在MNIST与CIFAR-10数据集上显著提升了鲁棒准确率，优于标准BNN与SVG D。
• 在深度强化学习中，f-SVGD实现了更优的探索能力，并在MuJoCo环境中获得了更高的累积回报。
• 该方法具有高效可扩展性：数百个粒子可容纳于单张GPU中，且由于仅交换函数空间信号，通信开销保持较低。
• 消融实验表明，权重空间方法中的梯度混合是导致失败的关键原因，而f-SVGD的设计有效消除了该问题。