A Framework for Variational Inference of Lightweight Bayesian Neural Networks with Heteroscedastic Uncertainties

Schodt らによる論文「A Framework for Variational Inference of Lightweight Bayesian Neural Networks with Heteroscedastic Uncertainties」の詳細な技術的要約を以下に示す。

1. 問題定義

ベイズニューラルネットワーク（BNN）は、予測不確実性を定量化し、エピステミック不確実性（限られたデータに起因するモデルの不確実性）とアレトリアック不確実性（データノイズ）を区別するために不可欠である。

• 課題: 多くの応用、特にヘテロスケダスティックな不確実性（ノイズレベルが入力に応じて変化するもの）を必要とする応用においては、標準的なアプローチでは、BNN が予測平均に加えて、追加のネットワークヘッドとしてアレトリアック分散を出力することを要求する。
• 限界: 分散を出力するノードを追加すると、学習可能なパラメータ数とアーキテクチャの複雑さが増加する。これは、ハードウェア制約のあるシステムへの展開や計算時間の制限がある環境を想定した軽量 BNNにとって有害である。
• ギャップ: 既存の手法は、モデルサイズと複雑で入力依存性の高いノイズ構造を捉える能力との間でトレードオフを強いる傾向がある。

2. 手法

著者らは、アレトリアック不確実性とエピステミック不確実性の両方を BNN の学習可能パラメータの分散に直接埋め込むことで、アレトリアック不確実性用の独立した出力ノードを不要にする、新たな推論フレームワークを提案する。

A. パラメータ分布の再解釈

ベイズニューロンのパラメータの分散をエピステミック不確実性のみに帰属させるのではなく、著者らはこれらの分布を総予測不確実性（アレトリアック＋エピステミック）を符号化するものとして再解釈する。

• 標準アプローチ: $Var(\hat{y}) = \sigma^2_{epistemic}$。アレトリアック分散（$\sigma^2_{aleatoric}$）は独立した出力として予測される。
• 提案アプローチ: 入力 $k$ に対する総分散 $\sigma^2_k$ は、アレトリアック分散とエピステミック分散の和として定義される：
  $$\sigma^2_k = \sigma^2_{a,k} + \sigma^2_{e,k}$$
  訓練中、ネットワークは再構成損失が分散 $\sigma^2_k$ を持つガウス分布を仮定するエビデンス下限（ELBO）を最小化する。分散 $\sigma^2_k$ は、パラメータ分布から導出される BNN 出力の分散として計算され、実質的にパラメータ分散にアレトリアックノイズを「吸収」させる。

B. モーメント伝播によるサンプリングフリー変分推論

計算効率を維持し、モンテカルロサンプリングのコストを回避するため、著者らはモーメント伝播方式を採用する。

• メカニズム: 各層への入力の平均と分散を、ネットワークを通じて解析的に伝播させ、出力の平均と分散を計算する。
• 層の処理:
  • 全結合層と畳み込み層: ニューロン間の独立性を仮定し、平均と分散の厳密な解析的伝播を行う。
  • Leaky-ReLU: 活性化関数への入力が正規分布に従うと仮定した近似伝播を行う。著者らは、切断正規分布を含む全期待値と分散の法則を用いて、出力の平均と分散の閉形式式を導出する。
• 利点: これにより、反復的なサンプリングを必要としない決定論的推論が可能となり、軽量かつリアルタイムなアプリケーションに適したフレームワークとなる。

3. 主な貢献

1. パラメータ分散の埋め込み: BNN の重みとバイアスの分散を、アレトリアック不確実性とエピステミック不確実性の和を表すように訓練するという概念的な転換。これにより、専用の分散出力ヘッドを不要にする。
2. 軽量アーキテクチャ: 標準的な「学習された分散」アプローチと比較して、より少ない学習可能パラメータでヘテロスケダスティックな不確実性を捉えることを可能にする BNN の実現。
3. サンプリングフリー推論: モンテカルロドロップアウトやサンプリングの計算オーバーヘッドを回避し、解析的な不確実性推定を提供する簡略化されたモーメント伝播技術の統合。
4. 統合フレームワーク: 不確実性源の再解釈と、効率的なモーメント伝播を組み合わせた、訓練および推論のための完全なパイプライン。

4. 実験結果

著者らは、ヘテロスケダスティックなガウスノイズを用いた多項式回帰タスクにおいて、分散を独立した出力とする標準的な「学習された分散」アプローチに対して、「埋め込み分散」アプローチを評価した。

• 軽量領域での性能:
  • 隠れニューロン数が少ない場合（例：4 個）、埋め込み分散モデルは分布内データに対してトレンドと不確実性の両方を正確に予測し、分布外データに対しては正しく高い不確実性を示した。
  • 学習された分散モデルは、同じ小規模アーキテクチャでは分布内データの真の不確実性を回復できず、分布外データに対して不確実性を著しく過小評価した。
• パラメータ効率:
  • 固定されたニューロン数において、「学習された分散」モデルは「埋め込み分散」モデルと同等の性能を達成するために、約 34% 多くの学習可能パラメータを必要とした。
  • 「埋め込み分散」アプローチは、低パラメータ領域においてより低い再構成損失（負の対数尤度）を達成した。
• 汎化性能:
  • 「埋め込み分散」モデルは分布外データに対する汎化性能が優れており、データが訓練分布から離れるにつれて再構成損失がより緩やかに増加することを示した。これに対し、「学習された分散」モデルではそのような傾向が見られなかった。
• スケーラビリティ: ネットワークサイズが増大する（例：1024 個のニューロン）につれて性能差は縮小し、両手法は同様の性能を示した。これは、提案手法がリソース制約のある軽量モデルにおいて最も重要であることを示唆している。

5. 意義

この研究は、エッジコンピューティングや安全クリティカルなシステムへのベイズニューラルネットワークの展開における重要なボトルネックに対処するものである。

• 効率性: ノイズを予測するための追加出力ノードを不要にすることで、モデルサイズと計算コストを大幅に削減し、厳格な制約を持つハードウェアにおける BNN の実用性を可能にする。
• 堅牢性: モデル複雑性を増大させることなくヘテロスケダスティックな不確実性を捉える能力により、軽量モデルが信頼性や安全性評価を犠牲しないことを保証する。
• 実用性: 推論のサンプリングフリーな性質により予測時間が短縮され、遅延が懸念されるリアルタイムアプリケーションに不可欠である。

要約すると、本論文は、総不確実性を BNN のパラメータ分布に効果的に「焼き付ける」ことが可能であることを実証しており、軽量かつヘテロスケダスティックな不確実性定量化のための従来のマルチ出力ヘッドに対する優れた代替手段を提供する。