[論文レビュー] Wide Bayesian neural networks have a simple weight posterior: theory and accelerated sampling

Jiri Hron, Roman Novak+2 more|arXiv (Cornell University)|Jun 15, 2022

Gaussian Processes and Bayesian Inference被引用数 4

ひとこと要約

本稿では、幅広いベイジアンニューラルネットワーク（BNN）の事後分布を、データに依存する再パラメータライゼーションであるreparametrization（再優先化）によって変換し、そのKullback–Leibler（KL）ダイバージェンスが幅が増加するにつれて標準正規事前分布から消失することを示す。この手法により、マルコフ連鎖モンテカルロ（MCMC）サンプリングにおける混合速度が向上し、有効サンプルサイズ（ESS）が最大50倍にまで向上する。このことは、全結合ネットワークおよび残差ネットワークの両方において、特に幅の広いネットワークで顕著に顕著な事後分布推論の高速化を実現する。

ABSTRACT

We introduce repriorisation, a data-dependent reparameterisation which transforms a Bayesian neural network (BNN) posterior to a distribution whose KL divergence to the BNN prior vanishes as layer widths grow. The repriorisation map acts directly on parameters, and its analytic simplicity complements the known neural network Gaussian process (NNGP) behaviour of wide BNNs in function space. Exploiting the repriorisation, we develop a Markov chain Monte Carlo (MCMC) posterior sampling algorithm which mixes faster the wider the BNN. This contrasts with the typically poor performance of MCMC in high dimensions. We observe up to 50x higher effective sample size relative to no reparametrisation for both fully-connected and residual networks. Improvements are achieved at all widths, with the margin between reparametrised and standard BNNs growing with layer width.

研究の動機と目的

• ベイジアンニューラルネットワーク（BNN）の理論的理解の不足と、特にパrameter空間における高い計算コストを解消すること。
• 既存の関数空間の結果（例：NNGP極限）を補完する形で、重み空間における広いBNNの事後分布行動を同定することで、広いBNNに関する理解のギャップを埋めること。
• 次元が非常に高い空間においても効果的に動作する、実用的で効率的なMCMCサンプリング手法をBNNに開発すること。
• ネットワーク幅が増加するにつれて標準正規分布に近づくように事後分布を変換することで、BNNにおける事後分布推論をより高速かつ効果的に実現すること。

提案手法

• BNN重み $ \theta $ を新しいパラメータ化 $ \theta = T(\theta) $ に変換する再パラメータライゼーション写像 $ T(\theta) $ を導入し、その結果、再パラメータライズされた事後分布と標準正規事前分布とのKLダイバージェンスが層の幅が増加するにつれて消失することを保証する。
• 理論的に（定理2.1）$ \mathrm{KL}(\mathcal{N}(0,I_d) \| p(\phi|\mathcal{D})) \to 0 $ が成り立つことを証明し、パrameter空間における事後分布が事前分布に収束することを示した。
• Cholesky分解を用いた計算効率の良い実装により、同時に再パラメータライゼーションとそのヤコビアン行列式の計算が可能となり、MCMCサンプリングが現実的に行えるようになった。
• 再パラメータライズされた事後分布に対してLangevin Monte Carlo（LMC）を適用し、対数事後分布の勾配が標準正規分布の勾配の周囲に集中することを利用することで、混合速度が向上した。
• 2次までのテイラー近似とHutchinsonの推定法を用いて、ヤコビアン行列式の対数を効率的に近似し、メモリ使用量を削減し、スケーラブルな勾配計算を可能にした。
• JAXの `vjp` および `jvp` 関数を用いて勾配を効率的に計算し、バックプロパゲーション中にメモリオーバーを回避するためのカスタムフォワードパス戦略を採用した。

実験結果

リサーチクエスチョン

• RQ1 ネットワークの幅が増加するにつれて、広いベイジアンニューラルネットワークの事後分布はパrameter空間でどのように振る舞うか？
• RQ2 無限大の幅の極限において、再パラメータライズされた事後分布と標準正規事前分布とのKLダイバージェンスが消失するような再パラメータライズーションを設計できるか？
• RQ3 このような再パラメータライズーションは、特に高次元の重み空間において、MCMCサンプリングの混合速度を向上させるか？
• RQ4 この手法は、残差ネットワークや異なるデータセットサイズを含むさまざまなアーキテクチャに、効率的かつスケーラブルに適用可能か？
• RQ5 reparametrizationの利点は、厳密なNNGP領域外でも、特に有限幅の実用的状況においても持続的か？

主な発見

• reparametrizationにより、無限大の幅の極限において、再パラメータライズされたBNN事後分布と標準正規事前分布とのKLダイバージェンスがゼロに収束することが保証され、パrameter空間における事後分布の簡素化の理論的基盤が得られた。
• この手法により、標準BNNと比較してMCMCサンプリングの有効サンプルサイズ（ESS）が最大50倍まで向上し、すべての幅において改善が観察され、層の幅が増加するにつれてその効果が顕著になった。
• CIFAR-10で3層の全結合ネットワーク（1層あたり1024ユニット）を用いた実験では、NNGP領域から大きく離れた状況でも、ESSに10倍の向上が観察された。
• ResNet-20では、上位層の幅が訓練例の数と同等以上である場合にのみ、ESSに10倍の向上が得られ、これは再パラメータライズーションの効果が状況依存であることを示唆している。
• Choleskyに基づくヤコビアン計算と、Hutchinsonの推定法およびJAXの `vjp` を用いたメモリ最適化された勾配推定により、計算の効率が向上した。
• 非ガウス型の尤度に対しても有効であることが示された。小さな実験でカテゴリカル尤度を用いた場合、再パラメータライズされたバージョンは標準パラメータライズーションよりも高いステップサイズで約99％の受容率を達成した。