[論文レビュー] 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

ひとこと要約

本稿では、勾配やヘッセ行列の明示的計算を回避するため、ニューラルタングェントカーネル（NTK）のニストローム近似を用いた、ベイジアンディープラーニングのための高速化された線形化ラプラス近似手法であるELLAを提案する。前向きモードの自動微分を活用することで、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}.

研究の動機と目的

• Kronecker分解または対角Gauss-Newton行列（GGN）近似に依存する従来の線形化ラプラス近似（LLA）手法における計算非効率性と精度損失を解消すること。
• 大規模なヘッセ行列やヤコビ行列を明示的に計算・保存することなく、ベイジアンディープラーニングにおけるスケーラブルかつ高精度な不確実性評価を可能にすること。
• LLAとニューラルタングェントカーネル（NTK）の関係を活用し、低ランクカーネル近似によって推論を高速化すること。
• ELLAと本物のLLAとの間の近似誤差に関する理論的保証を提供し、予測不確実性の信頼性を確保すること。
• 最近代のアーキテクチャ、特に視覚変換器や標準的なImageNetベンチマークにおいて、ELLAのスケーラビリティと性能を実証すること。

提案手法

• 線形化ラプラス近似の共分散行列の効率的計算を可能にするために、ニューラルタングェントカーネル（NTK）のニストローム近似を提案する。これにより、ヘッセ行列やヤコビ行列の明示的保存を回避できる。
• NTK近似に必要なヤコビ行列・ベクトル積（JVP）を効率的に計算するために、前向きモードの自動微分（fwAD）を活用する。
• データポイントのサブセットを用いてNTKに低ランク近似を適用することで、計算コストを削減しながら予測の忠実性を維持する。
• 本物のLLAのコア構造を維持するが、全般的なGGN計算をカーネルベースの近似に置き換えることで、大規模モデルへのスケーラビリティを実現する。
• 理論的分析により、NTK近似の精度が向上するにつれて、ELLAと本物のLLAとの間の予測誤差が減少することが示された。
• 標準的なディープラーニングライブラリと互換性のある形で実装されており、事前学習済みモデルへの即時統合が可能である。

実験結果

リサーチクエスチョン

• RQ1 NTKのニストローム近似が、予測精度を損なわせることなく、線形化ラプラス近似を効率化できるか？
• RQ2 標準的なベンチマークにおいて、ELLAの不確実性推定性能は、KFAC、対角、または最終層近似を用いた標準的なLLAバージョンと比較してどうなるか？
• RQ3 全ヘッセ行列の計算が非現実的である視覚変換器のような最近代のアーキテクチャにも、ELLAはスケーラブルに適用可能か？
• RQ4 NTK近似の精度と、ELLAにおける最終的な予測不確実性との間には、どのような理論的関係があるか？
• RQ5 分布シフト、例えばImageNet-Cのノイズ汚染に対して、ELLAは信頼性のある不確実性キャリブレーションを維持できるか？

主な発見

• CIFAR-10では、ELLAはECEとNLLの観点で、KFAC、対角、最終層近似を用いたLLAと同等またはそれを上回る、最先端の不確実性キャリブレーションを達成した。
• ResNet-18、ResNet-34、ResNet-50を用いたImageNetでは、ELLAはECEが0.015〜0.018であり、MAP（0.026〜0.037）よりも低い結果を示し、より優れた不確実性キャリブレーションを実現した。
• ImageNetにおけるViT-Bでは、ELLAはECEを0.022に低下させたのに対し、MAPは0.039であった。これは、最近代のアーキテクチャでも不確実性キャリブレーションが向上したことを示している。
• ViT-Bでは、ELLAは81.6%の高い精度と0.695の低いNLLを達成したのに対し、MAPは81.5%の精度と0.700のNLLであった。精度の損失なしにキャリブレーションが向上した。
• ImageNet-Cのノイズ汚染に対しては、ELLAはNLLおよびECEの両面でMAPを常に上回り、分布シフトへのロバストネスを示した。
• 大規模モデルでは、標準的なLLA実装が直面するメモリ不足や長時間の学習を回避し、優れたスケーラビリティを示した。