[论文解读] Accelerated Linearized Laplace Approximation for Bayesian Deep Learning

Zhijie Deng, Feng Zhou+1 more|arXiv (Cornell University)|Oct 23, 2022

Gaussian Processes and Bayesian Inference被引用 5

一句话总结

本文提出 ELLA，一种用于贝叶斯深度学习的加速线性化拉普拉斯近似方法，通过使用神经正切核（NTK）的Nystrom近似来规避昂贵的海森矩阵或雅可比矩阵的显式计算。通过利用前向模式自动微分，ELLA 在保持理论保证的同时，实现了最先进的不确定性量化，且在视觉变换器等模型上展现出改进的可扩展性，同时在 ImageNet 和 CIFAR-10 基准测试中优于标准 LLA 近似方法。

ABSTRACT

Laplace approximation (LA) and its linearized variant (LLA) enable effortless adaptation of pretrained deep neural networks to Bayesian neural networks. The generalized Gauss-Newton (GGN) approximation is typically introduced to improve their tractability. However, LA and LLA are still confronted with non-trivial inefficiency issues and should rely on Kronecker-factored, diagonal, or even last-layer approximate GGN matrices in practical use. These approximations are likely to harm the fidelity of learning outcomes. To tackle this issue, inspired by the connections between LLA and neural tangent kernels (NTKs), we develop a Nystrom approximation to NTKs to accelerate LLA. Our method benefits from the capability of popular deep learning libraries for forward mode automatic differentiation, and enjoys reassuring theoretical guarantees. Extensive studies reflect the merits of the proposed method in aspects of both scalability and performance. Our method can even scale up to architectures like vision transformers. We also offer valuable ablation studies to diagnose our method. Code is available at \url{https://github.com/thudzj/ELLA}.

研究动机与目标

• 解决现有线性化拉普拉斯近似（LLA）方法中依赖 Kronecker 分解或对角 GGN 近似所导致的效率低下和精度损失问题。
• 在不显式计算或存储大型海森矩阵或雅可比矩阵的前提下，实现在贝叶斯深度学习中可扩展且精确的不确定性量化。
• 利用 LLA 与神经正切核（NTK）之间的联系，通过低秩核近似加速推理。
• 提供 ELLA 与原始 LLA 之间近似误差的理论保证，确保预测不确定性的可靠性。
• 在如视觉变换器等现代架构以及标准 ImageNet 基准上，展示 ELLA 的可扩展性与性能。

提出的方法

• 提出对神经正切核（NTK）的 Nystrom 近似，以实现无需显式海森或雅可比矩阵存储的 LLA 协方差矩阵的高效计算。
• 利用前向模式自动微分（fwAD）高效计算 NTK 近似所需的雅可比-向量积（JVPs）。
• 通过数据点子集对 NTK 实施低秩近似，降低计算成本，同时保持预测保真度。
• 保持原始 LLA 的核心结构，但将完整的 GGN 计算替换为基于核的近似，从而实现对大型模型的可扩展性。
• 理论分析表明，ELLA 与原始 LLA 之间的预测误差随着 NTK 近似精度的提高而减小。
• 以与标准深度学习库兼容的方式实现该方法，支持与预训练模型的即插即用式集成。

实验结果

研究问题

• RQ1 NTK 的 Nystrom 近似是否能有效加速线性化拉普拉斯近似，同时不损失预测准确性？
• RQ2 ELLA 的不确定性估计性能与标准 LLA 变体（如 KFAC、对角或最后一层近似）在标准基准上的表现相比如何？
• RQ3 ELLA 是否能够扩展到如视觉变换器这类全海森矩阵计算不可行的现代架构？
• RQ4 NTK 近似的精度与 ELLA 中最终预测不确定性之间的理论关系是什么？
• RQ5 ELLA 在分布偏移（如 ImageNet-C 损伤）情况下是否仍能保持可靠的不确定性校准？

主要发现

• 在 CIFAR-10 上，ELLA 实现了最先进的不确定性校准，其 ECE 和 NLL 指标与使用 KFAC、对角或最后一层近似的 LLA 方法相比达到匹配或更优水平。
• 在 ImageNet 上使用 ResNet-18、ResNet-34 和 ResNet-50 时，ELLA 的 ECE（0.015–0.018）低于 MAP（0.026–0.037），表明其不确定性校准更优。
• 对于 ViT-B 在 ImageNet 上，ELLA 将 ECE 降低至 0.022，而 MAP 为 0.039，表明其在现代架构上实现了更优的不确定性校准。
• 在 ViT-B 上，ELLA 保持了具有竞争力的准确率（81.6%）和更低的 NLL（0.695），优于 MAP（81.5% 准确率，0.700 NLL），表明其在不损失准确率的前提下实现了更优的校准。
• 在 ImageNet-C 损伤数据上，ELLA 在 NLL 和 ECE 上均持续优于 MAP，表明其对分布偏移具有鲁棒性。
• 该方法避免了标准 LLA 实现中因大模型导致的内存溢出和长时间训练的问题，展现出卓越的可扩展性。