技術要約：ベイズ深層学習における較正済みサンプリングフリー不確実性推定

1. 問題提起

現代のディープラーニングモデル、特にトランスフォーマー・アーキテクチャに基づくモデルは、過剰な自信（オーバーコンフィデンス）に陥りやすいことで知られており、これが疾患検出、法的推論、視覚障害者向けの視覚支援といった高リスクなアプリケーションにおける信頼性を制限しています。ベイズニューラルネットワーク（BNN）は、モデルパラメータの分布を学習することで不確実性定量化への原理的な経路を提供しますが、テスト時の計算コストが大きなボトルネックとなっています。

標準的なベイズ推論は、重みを事後分布からサンプリングした多数のフォワードパス（モンテカルロ（MC）サンプリング）を通じて予測を平均化する必要がありますが、これは大規模なアーキテクチャにおいては計算量的に極めて困難です。逆に、平均ネットワークのみを使用する場合（分散を無視する場合）は効率的ですが、不確実性を捉えることができません。平均と分散をレイヤーを通じて単一のパスで解析的に伝播させる既存の分散伝播手法は、歴史的に多層パーセプトロン（MLP）に限定されてきました。それらを現代的なアーキテクチャ（Transformer、ResNetを含むCNN）へ拡張することは、深層化や正規化レイヤーの存在により困難でした。先行研究である「Streamlining」[30]などは、正規化レイヤーと活性化関数を線形化していますが、その結果、伝播される平均が最終レイヤーまで平均ネットワークと同一になってしまい、追加の計算コストを支払っているにもかかわらず、選択的予測（selective prediction）の改善には至っていません。

2. 手法：較正済み分散伝播 (Calibrated Variance Propagation: CVP)

著者らは、TransformerやCNNを含む現代的なアーキテクチャに対して、単一のフォワードパスで正確な不確実性推定を実現するために設計された包括的なフレームワークである**Calibrated Variance Propagation (CVP)**を提案しています。CVPは、標準的な決定論的フォワードパスを、平均と分散のペア $(\mu, \Sigma)$ をネットワーク内を伝播させる一連のオペレータへと置き換えることで動作します。

この手法は、以下の3つの具体的な技術的ギャップに対処しています。

A. 活性化関数における正確なモーメント

入力の分散を伝播する平均を、入力の分散に依存しないものにするデルタ法（一次テイラー展開による線形化）とは異なり、CVPは非線形活性化関数に対して正確なモーメント・マッチングを採用しています。ガウス入力を仮定することで、CVPはGELU、sigmoid、tanh、およびReLUの平均と分散に関する閉形式の式を利用します。これにより、伝播される平均が各非線形層において平均ネットワークの予測から乖離することを保証し、入力依存の不確実性を捉えます。

B. 正規化レイヤーのための期待値ベースの規則

核心的な貢献は、正規化レイヤー（例：LayerNorm）のための新しい伝播規則です。従来の手法は、分散項 $s^2(x)$ を $s^2(\mu_x)$ に線形化し、操作を平均の決定論的なアフィン変換へと簡略化していました。CVPは、$s^2(x)$ を入力分布の下での期待値に置き換えます：
$$E[s^2(x)] = s^2(\mu_x) + \left(1 - \frac{1}{D}\right)\bar{\Sigma}_x$$
ここで $\bar{\Sigma}_x$ は対角入力分散の平均です。これにより、LayerNorm操作が入力モーメントに対してアフィンとなり、Transformerアーキテクチャの中心である正規化レイヤーを通じて不確実性を正確に伝播させることが可能になります。

C. レイヤーごとの較正 (Per-Layer Calibration)

深さが増すにつれて（対角共分散およびガウス入力の仮定に起因する）近似誤差が蓄積するのを防ぐため、CVPは学習可能なレイヤーごとのスケーリング係数 ($\alpha_i$) を導入しています。これらの係数は、主要なレイヤー（具体的には活性化および正規化レイヤーの前）において、伝播される分散を乗算的に再スケーリングします。$\alpha_i$ パラメータは、温度スケーリング（temperature scaling）に類似して、保持された検証セット上で負の対数尤度（NLL）を最小化することにより、軽量な較正フェーズによって適合されます。

3. 主な貢献

1. 正規化レイヤーのための分散伝播規則： 著者らは、線形化と比較してMCサンプリングへの忠実度が大幅に向上する、正規化レイ層のための新しい規則を開発しました。これは合成データおよび実際のモデルの活性化の両方で検証されています。
2. レイヤーごとの較正ステージ： 学習可能なスケーリング係数の導入は、系統的な誤差の蓄積を補正し、較正、そして極めて重要なことに、選択的予測の性能を大幅に向上させます。
3. 最適化された実装： CVPの実装は、決定論的フォワードパスの約2〜3倍の計算コストで単一のフォワードパスを実行し、繰り返しのサンプリングのオーバーヘッドなしに、大幅な信頼性の向上を実現します。

4. 実験結果

CVPは、単一モード（画像分類：CIFAR-10, CIFAR-100, ImageNet）およびマルチモード（VQA：VQAv2; 視覚的推論：NLVR2）のベンチマークにおいて、IVON変分最適化器を用いて訓練されたモデルを用いて評価されました。

• 効率と信頼性のトレードオフ： CVPは、平均ネットワークおよびMCサンプリングの両方に対して一貫して**パレート優位（Pareto-dominates）**を示しました。CVPは、MCサンプリング（例：4サンプル）と同等の不確実性の質を、そのわずかな部分のコスト（単一パスの約2.3〜2.9倍 vs 4サンプルの場合の4倍）で達成しています。
• 選択的予測： CVPは、Coverage at Risk (C@R) において顕著な改善を示しました。例えば、BEiT-3を用いたNLVR2ベンチマークにおいて、CVPはリスク0.5%におけるカバレッジを、Streamliningの8.2%から**14.6%へと向上させました。ViLTを用いたVQAv2では、2.6%から10.8%**へと向上しました。
• Streamliningとの比較： Streamliningは平均ネットワークに対して較正をわずかに改善しますが、伝播される平均が平均ネットワークと同一であるため、選択的予測の改善には至りません。CVPの早期の平均の乖離とレイヤーごとの較正により、CVPはStreamliningを大幅に上回る性能を発揮します。
• アブレーション研究： CVPのいずれかの構成要素（正確な活性化モーメント、正規化規則、またはレイヤーごとの較正）を取り除くと性能が低下し、観察された利得にはこれらすべての要素が必要であることが確認されました。

5. 意義と主張

本論文は、CVPが大規模な現代的アーキテクチャに対するサンプリングフリーの不確実性推定を可能にすることで、ベイズ深層学習における重要なギャップを埋めるものであると主張しています。その意義は以下の通りです。

• 実用性： 繰り返しの推論が実行不可能なリアルタイムかつ安全性が重視されるシステムにおいて、MCサンプリングに代わる実行可能な選択肢を提供します。
• アーキテクチャの汎用性： 従来のメソッドがMLPのような構造に限定されていたのに対し、マルチモーダルTransformerやResNetへの適用に成功した最初の分散伝播フレームワークです。
• パフォーマンス： 一貫したパレート改善を提供し、低レイテンシを維持しながら、高価なMCサンプリング手法に匹敵、あるいはそれを上回る信頼性指標（較正および選択的予測）を実現します。

著者らは、CVPが決定論的パスの約2倍のメモリフットプリント（分散を保持するため）を必要とし、較正のために保持されたセットを必要とすることなど、限界についても認めています。さらに、現在の実験はエンコーダー型のアーキテクチャと分類またはVQAヘッドに限定されており、自己回帰的な生成については今後の課題としています。しかし、本手法は、効率的で信頼性の高い深層学習における不確実性定量化の新しい基準を確立しています。