[論文レビュー] Path-Gradient Estimators for Continuous Normalizing Flows

Lorenz Vaitl, Kim A. Nicoli+2 more|arXiv (Cornell University)|Jun 17, 2022

Model Reduction and Neural Networks被引用数 7

ひとこと要約

本稿では、変分推論におけるより低い分散の学習を可能にする、連続正規化フロー向けの効率的なパス勾配推定器を提案する。尤度の対数の勾配を計算するために補助ODEを解くことで、標準の再パラメータ化勾配と比較して1イテレーションあたり5–14%の実行時間増加で、収束が速く、性能が向上する。

ABSTRACT

Recent work has established a path-gradient estimator for simple variational Gaussian distributions and has argued that the path-gradient is particularly beneficial in the regime in which the variational distribution approaches the exact target distribution. In many applications, this regime can however not be reached by a simple Gaussian variational distribution. In this work, we overcome this crucial limitation by proposing a path-gradient estimator for the considerably more expressive variational family of continuous normalizing flows. We outline an efficient algorithm to calculate this estimator and establish its superior performance empirically.

研究の動機と目的

• 現代の変分推論で広く使われる表現力のある連続正規化フローに対して、効率的なパス勾配推定器が不足しているという問題に取り組む。
• 従来のパス勾配手法が正規化フローに適用可能でない、あるいは実行コストが高すぎるという制限を克服する。
• VAE や格子場理論のような高次元問題における連続正規化フローの、より速く安定した学習を可能にする。
• 最小限の計算オーバーヘッドと優れた実験的性能を備えた、標準の再パラメータ化勾配の即時交換可能な代替手法を提供する。

提案手法

• 変分対数密度の勾配をフローを逆方向に伝搬させる補助ODE（式11）を導出することで、パス勾配の計算を可能にする。
• 主フローと同じODEソルバーを用いて補助ODEを前向きに解くことで、互換性と低メモリオーバーヘッドを確保する。
• 補助ODEの解を用いてパス勾配を計算し、再パラメータ化サンプリングを介したELBOへのパラメータの隠れ依存関係のみを捉える。
• 最小限のコード変更で既存の学習パイプラインに統合可能であり、標準勾配の即時交換が可能になる。
• 連続正規化フローの構造を活用して勾配を効率的に計算し、ODE統合全体に対する逆方向微分の必要を回避する。
• 主フローのODEソルバーで得た中間状態を再利用することで、1イテレーションあたりのメモリ使用量を一定に保ち、計算効率を維持する。

実験結果

リサーチクエスチョン

• RQ1 表現力が高く、パス勾配計算に課題を伴う連続正規化フローに対して、効率的なパス勾配推定器を実装できるか？
• RQ2 提案手法のパス勾配推定器は、高次元設定下で標準の再パラメータ化勾配と比較して、収束が速く、ELBO最適化が優れているか？
• RQ3 パス勾配推定器の実行時間とメモリ使用量は、標準学習と比較してどの程度のオーバーヘッドを示すか？
• RQ4 VAE や格子場理論の両方において、異なるアーサテクチャ、データセット、ODEソルバーの下で、パス勾配推定器の性能はどのように変化するか？
• RQ5 計算物理学や生成モデル分野における大規模な変分推論タスクにおいて、パス勾配推定器を実用的に導入できるか？

主な発見

• パス勾配推定器は学習の分散を低減し、収束を加速させ、最高次元の格子場理論タスクにおいて有効サンプルサイズ（ESS）が30%以上増加した。
• A100 GPU上では1イテレーションあたり実行時間が5–14%増加にとどまり、格子サイズが大きくなるに従い、その増加率は低下する。
• VAEでは標準の再パラメータ化勾配と同等またはそれ以上の性能を達成しており、Omniglotではベースラインの111%以内、MNISTでは98%のウォールクロック時間で学習が可能だった。
• Agrawalら（2020）の先行研究の最先端手法と比較して、1イテレーションあたりの処理がはるかに速く、メモリ使用量も半分に抑えられた（メモリコストが2倍であった）。
• パス勾配推定器と標準勾配推定器との間で関数評価回数の顕著な差は観察されず、ODEソルバーの動作が安定していることが示された。
• 実験的結果から、標準勾配をパス勾配に置き換えることで、多様なタスクとアーサテクチャにおいて一貫して顕著な性能向上が得られた。