技術的サマリー：大規模ベイズ深層ニューラルネットワークのためのスコアベース変分事後分布推論

問題定義

ベイズ深層ニューラルネットワーク（BNN）は、頑健な不確実性の定量化、ノイズ耐性、過学習の軽減など、点推定に基づく深層学習に比べて顕著な利点を提供します。しかし、BNN における正確な事後分布推論は一般的に非現実的です。変分推論（VI）は、対象となる事後分布 $\pi(\theta)$ を扱いやすい分布 $q_\lambda(\theta)$ で近似することで、推論を最適化問題として定式化する支配的な近似手法です。

現在の支配的な手法である ELBO ベースの VI（逆 KL 発散の最小化）は、近似が対象密度のすべてのモードを捉えられないというよく知られた「モードの崩壊（mode collapsing）」問題に悩まされています。一方、対象スコア $\nabla_\theta \log \pi(\theta)$ とモデルスコア $\nabla_\theta \log q_\lambda(\theta)$ の差異を最小化するスコアベースの VI 手法は、理論的にはモードの崩壊に対処できますが、既存のスコアベース手法は大規模 BNN にはスケーラブルではありません。具体的には以下の通りです：

1. 計算的非現実性: ガウス・スコアマッチング（GSM）やバッチ・アンド・マッチ（BaM）などの手法は、状態次元に対して立方または二次的にスケールする行列の逆行列計算や乗算に依存しており、高次元のニューラルネットワークパラメータに対しては実行不可能です。
2. 確率的勾配との非互換性: 既存のスコアベース手法は通常、正確なスコアを必要とするか、大規模データセットでの学習に不可欠な、確率的勾配から導出されるノイズを含む不偏ミニバッチ・スコアを処理できません。
3. 再パラメータ化の制約: 多くの手法は再パラメータ化サンプリングに依存しており、高次元では計算コストが高く不安定になる可能性があります。

手法

著者は、Vision Transformers（ViT）や ResNets を含む大規模 BNN にスケーラブルな、新しい近接確率的勾配スコアベース VIアルゴリズムを提案します。

中核となる最適化目的関数

本手法は、各反復 $t$ において以下の最適化問題を解くことで、変分分布 $q_t(\theta)$ を $q_{t+1}(\theta)$ に反復的に更新します：

$$q_{t+1} = \arg \min_{q \in \mathcal{Q}} \mathbb{E}_{\theta \sim q_t(\theta)} \left[ \alpha_t \|\nabla_\theta \log q(\theta) - \nabla_\theta \log q_t(\theta)\|_2^2 + \|\nabla_\theta \log q(\theta) - s(\theta)\|_2^2 \right]$$

ここで：

• $s(\theta) = \nabla_\theta \log \pi(\theta)$ は対象スコアです。
• $\alpha_t$ は適応的な近接影響パラメータ（0 から定数 $\le 1$ まで線形にスケジューリング）です。
• 第一項は、スコアの観点から新しい反復値が現在の反復値に近くなることを強制する近接ペナルティです。
• 第二項はスコアマッチング損失です。

主要な技術的革新

1. ノイズを含むミニバッチ・スコアの処理: 正確なスコアを必要とする従来のスコアベース手法とは異なり、本アプローチは確率的勾配を通じて得られるノイズを含む不偏ミニバッチ・スコア $\hat{s}(\theta)$ を利用します。損失関数におけるスコア項の線形性により、ノイズを含むスコアであっても、変分パラメータ $\lambda$ に対する損失の勾配は不偏推定量のままとなります。
2. 再パラメータ化の回避: 近接項により、期待値を可変の $q$ ではなく、固定された現在の分布 $q_t$ に関して計算することが可能になります。これにより、ADVI で使用される再パラメータ化トリックの必要性を回避し、従来のフィッシャー/スコア手法で必要とされる $\theta$ に関する二階微分を避けることができ、計算オーバーヘッドとメモリ使用量を大幅に削減します。
3. スケーラビリティ: 行列逆行列を必要とする共分散誘導ノルムではなくユークリッドノルムを使用し、確率的勾配をサポートすることで、本手法は数億のパラメータを持つネットワーク（例：約 3 億 500 万パラメータの ViT-L-32）にスケーラブルです。

理論的収束

著者は、変分族が滑らかなベクトル場を表現するのに十分豊富であるという仮定の下で、反復が増加するにつれて変分分布のスコアが対象スコアに収束することを示す収束定理（定理 3.1）を提供します。表現の豊富さに関する仮定は理想的ですが、著者は、真の事後分布と変分族の両方がガウス分布である場合、この仮定が厳密に成り立つと指摘しています。

主要な貢献

1. 新しいアルゴリズム: スコアマッチング損失と近接ペナルティを組み合わせた、近接確率的勾配スコアマッチング VI アルゴリズムの導入。
2. 大規模 BNN へのスケーラビリティ: 本手法は、大規模な深層ネットワーク（ResNets、ViTs）および大規模データセットにおいて効果的に機能することが実証された、初のスコアベース VI 手法であり、GSM や BaM の計算ボトルネックを克服しました。
3. ノイズに対する頑健性: アルゴリズムはノイズを含む不偏ミニバッチ・スコアを効果的に処理し、フルバッチ・スコア計算を必要とせずに確率的勾配の使用を可能にします。
4. 再パラメータ化の排除: 本手法は再パラメータ化トリックを回避し、ELBO ベースの ADVI に比べて最適化の地形を単純化し、効率を向上させます。

実験結果

著者は、合成データ、PosteriorDB ベンチマーク、および実世界のタスクにおいて本手法を評価しました：

• 合成データ: ガウス分布およびガウス混合分布を対象とした場合、提案手法は、ADVI、GSM、BaM と比較して、より速い収束と初期パラメータの選択（特に小さな共分散固有値に関するもの）に対する優れた頑健性を示しました。GSM/BaM が数値的問題により失敗したモデルの不一致や高次元状態空間においても、正常に処理しました。
• MNIST（小規模 BNN）: 低データ領域において、本手法はテスト誤り率において ADVI と同等の結果を達成しましたが、期待較正誤差（ECE）および負対数尤度（NLL）が低いことから、不確実性の定量化において優れていることが示されました。
• 大規模視覚認識（ResNet-101 および ViT-L-32）: Pets、Flowers、Aircraft、DTD などのデータセットにおいて、本手法はテスト予測精度において ADVI と同等でありながら、不確実性指標（ECE および NLL）において一貫して ADVI を上回りました。競合手法（GSM、BaM、従来のスコアマッチング）は、メモリ制約または数値的爆発により失敗しました。
• 時系列予測: Koopa ネットワークアーキテクチャを使用し、ECL、ETTh2、為替レート、ILI、Traffic、Weather などの複数のデータセットにおいて、本手法は MSE、MAE、NLL において ELBO ベースの VI および非ベイズベースラインを上回りました。
• 正規化フロー: 本手法は正規化フロー変分密度に拡張され、時系列タスクにおいて ADVI に対して一貫した改善を示しました。
• 効率性: 前回の反復値と現在の反復値の両方を保持しているにもかかわらず、ウォールクロック時間および GPU メモリ使用量は ADVI と同等であり、大規模トレーニングにおける実用的な実現可能性を実証しました。

意義と主張

本論文は、この研究が、大規模ベイズ深層学習における支配的な ELBO ベースの VI に対するスケーラブルで柔軟な代替手段を提供すると主張しています。逆 KL 目的関数に固有のモード崩壊問題に対処し、従来のスコアベース手法の計算的限界を克服することで、提案アルゴリズムは Vision Transformers のような現代の高容量アーキテクチャにおける厳密なベイズ推論を可能にします。

著者は、彼らのアプローチが多様なベンチマークにおいて経験的に効果的であり、予測性能を犠牲にすることなくより優れた不確実性の定量化を提供することを強調しています。また、理論的収束は変分族の表現力に関する強い仮定に依存していますが、経験的結果は実用上の手法としての有用性を支持していると控えめに述べています。この研究は、スコアベース VI を、以前の手法のガウス分布のみの制約やフルバッチ要件を超えた、ベイズ深層学習のための実用的でスケーラブルな道筋として位置づけています。