[論文レビュー] Tractable Function-Space Variational Inference in Bayesian Neural Networks

Tim G. J. Rudner, Zonghao Chen+2 more|arXiv (Cornell University)|Dec 28, 2023

Adversarial Robustness in Machine Learning被引用数 6

ひとこと要約

この論文は、パラメータではなく関数空間に直接事後分布をモデル化することで、スケーラブルな不確実性推定を可能にする、扱いやすい関数空間変分推論（FSVI）を提案する。関数空間Kullback-Leibler発散の単純な推定器を導入することで、意味のある事前分布の組み込みが可能となり、多様なタスク、特に安全性が求められる医療診断ベンチマークにおいて、最先端の予測不確実性と精度を達成する。

ABSTRACT

Reliable predictive uncertainty estimation plays an important role in enabling the deployment of neural networks to safety-critical settings. A popular approach for estimating the predictive uncertainty of neural networks is to define a prior distribution over the network parameters, infer an approximate posterior distribution, and use it to make stochastic predictions. However, explicit inference over neural network parameters makes it difficult to incorporate meaningful prior information about the data-generating process into the model. In this paper, we pursue an alternative approach. Recognizing that the primary object of interest in most settings is the distribution over functions induced by the posterior distribution over neural network parameters, we frame Bayesian inference in neural networks explicitly as inferring a posterior distribution over functions and propose a scalable function-space variational inference method that allows incorporating prior information and results in reliable predictive uncertainty estimates. We show that the proposed method leads to state-of-the-art uncertainty estimation and predictive performance on a range of prediction tasks and demonstrate that it performs well on a challenging safety-critical medical diagnosis task in which reliable uncertainty estimation is essential.

研究の動機と目的

• ベイジアンニューラルネットワークにおけるパラメータ空間変分推論の限界を解決すること。これは、しばしば不確実性の定量的評価が不十分であり、非ベイジアン手法に劣ることがある。
• データ生成プロセスに関するタスク固有の事前知識を、ベイジアンニューラルネットワーク推論に組み込むことを可能にすること。
• 高次元および過パラメータ化された設定において、信頼性のある予測不確実性推定を可能にする、スケーラブルで扱いやすい関数空間変分推論の手法を開発すること。
• 標準ベンチマークおよび実世界の安全性が求められる医療診断タスクにおける評価を通じて、分布シフト下での耐性を示すこと。

提案手法

• ベイジアン推論を関数空間に明示的に定式化し、ネットワークパラメータが誘導する関数空間上の事後分布を標的とする。
• 関数空間分布間のKullback-Leibler発散のための新規で単純な推定器を導入し、確率的変分推論を可能にする。
• 各勾配ステップで入力点をサンプリングするためのコンテキスト分布 $p_{\mathbf{X}_{\mathcal{C}}}$ を使用し、変分目的関数における上界を近似する。
• 画像入力の場合、コンテキスト分布は、訓練データの画素から構築されたモノクロ画像上での一様分布として定義される。表形式データの場合は、実測の入力範囲を用いる。
• 柔軟な事前分布設計をサポートしており、分布外領域での高い不確実性を促進するような事前分布も可能である。
• Adam最適化を用い、標準的なディープラーニングフレームワークで実装されており、ResNet や MLP などの標準アーキテクチャとも互換性がある。

実験結果

リサーチクエスチョン

• RQ1 高次元および過パラメータ化されたニューラルネットワークに対して、関数空間変分推論を扱いやすくスケーラブルにできるか？
• RQ2 関数空間事前分布にタスク固有の事前知識を組み込むことで、パラメータ空間手法と比較して予測不確実性と精度が向上するか？
• RQ3 特に安全性が求められる応用において、分布シフト下での本手法の性能はいかがなものか？
• RQ4 関数空間変分推論が、ベイジアンおよび非ベイジアンのベースラインを上回る不確実性推定および予測性能を達成できるか？

主な発見

• FSVIは、1次元回帰やTwo Moonsデータセットを含む、回帰および分類の複数のベンチマークで、最先端の予測性能と不確実性推定を達成する。
• Snelson 1次元回帰タスクでは、FSVIが訓練データにうまくフィットする一方で、訓練点から離れた領域では望ましく高い不確実性を示す予測分布を生成する。
• 糖尿病網膜症診断タスクでは、専門家へのリファラルを伴うFSVIベースの選択的予測が、ベースラインを上回り、低リファラルレートでも高い精度を維持する。
• 本手法は分布シフトに対して耐性を示し、分布外（OOD）サンプルに対して高いエピステミック的不確実性を割り当て、効果的な選択的予測を可能にする。
• タスク固有の事前分布を用いたFSVIは、標準的なパラメータ空間変分推論やディープアンサンブルと比較して、不確実性のキャリブレーションと予測精度を顕著に向上させる。
• 関数空間KL発散のための提案された推定器により、安定的かつスケーラブルな学習が可能となり、関数空間推論を実世界の応用に実用可能にする。