[論文レビュー] 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-BNN）を用いた無限に深いベイジアンニューラルネットワークを提案する。各無限小の層における重みの不確実性が、隠れユニットにおけるランダムな軌道を引き起こす。勾配に基づく変分推論を、新しいゼロ分散勾配推定器を用いて定式化することで、任意に表現力の高い近似事後分布を可能にし、スケーラブルでメモリ効率の良い学習が実現可能となる。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のドリフト関数と拡散関数を深層ニューラルネットワークでパrameter化し、時間相関のある重みの不確実性をモデル化する。
• 確率的変分推論を適用して、SDEのパラメータに関する変分事後分布を最適化し、修正された変分下界（ELBO）を最大化する。
• 「スターティング・ザ・ランドイング」のテクニックにインspiredされた分散低減型勾配推定器を導入し、事後分布の精度が向上する極限で分散がゼロに近づく。
• ブラックボックスの適応的SDEソルバーを用いて出力層の状態を計算し、学習中に一定のメモリコストを維持する。
• 拡張された状態次元を用いたパrameterizationにより表現力の向上を図り、より良い事後分布近似を可能にする。

実験結果

リサーチクエスチョン

• RQ1 連続的深さのベイジアンニューラルネットワークは、その事後分布近似が任意に表現力を持つように構築可能か？
• RQ2 無限パラメータのベイジアンモデルに対して、事後分布近似が改善するにつれて分散がゼロに近づくような勾配推定器を設計可能か？
• RQ3 SDEに基づくベイジアンネットワークは、ニューラルODEのメモリ効率を維持しつつ、標準ベンチマークで競争力のある性能を達成可能か？
• RQ4 本手法は、離散的深さの代替手法と比較して、分布シフトにさらされたOut-of-Distributionデータに対するキャリブレーションとロバスト性において優れているか？

主な発見

• SDE-BNNはMNISTで99.30% ± 0.09の精度、CIFAR-10で88.08% ± 1.25の精度を達成し、ベースラインのResNetや他のベイジアンモデルよりもキャリブレーションとロバスト性において優れた性能を示した。
• 「スターティング・ザ・ランドイング」（STL）勾配推定器を搭載したSDE-BNNは、CIFAR-10でテスト精度96.89%を達成し、標準的なSVIと比較して負の対数尤度を75%削減した。
• CIFAR-10における期待キャリブレーション誤差（ECE）は、SDE-BNN + STLバージョンで6.44 × 10⁻²にまで低下し、よりキャリブレーションの良い予測が得られた。
• CIFAR-10Cベンチマークにおいて、全腐敗深刻度レベルで一貫した性能を維持し、分布シフトに対して高いロバスト性を示した。
• 提案された勾配推定器は分散を顕著に低減し、非常に柔軟な事後分布パラメータ化でも安定した学習を可能にした。
• 拡張された状態次元により、余分でない次元で意味のある表現を学習可能であり、不要または冗長な次元を無視する能力も向上した。