[論文レビュー] Storchastic: A Framework for General Stochastic Automatic Differentiation

Emile van Krieken, Jakub M. Tomczak+1 more|arXiv (Cornell University)|Apr 1, 2021

Domain Adaptation and Few-Shot Learning参考文献 45被引用数 8

ひとこと要約

Storchastic は、一般化された確率的自動微分のための新しいフレームワークであり、確率的計算グラフ内の各サンプリングステップで多様で低分散の勾配推定手法を適用可能にする。推定器を提案分布、重み関数、勾配関数、制御変数の4つの構成要素に分解することで、任意の順序の導関数についても不偏な勾配推定を保証する。また、測度値導関数や SPSA といった高度な技術もサポートする。

ABSTRACT

Modelers use automatic differentiation (AD) of computation graphs to implement complex Deep Learning models without defining gradient computations. Stochastic AD extends AD to stochastic computation graphs with sampling steps, which arise when modelers handle the intractable expectations common in Reinforcement Learning and Variational Inference. However, current methods for stochastic AD are limited: They are either only applicable to continuous random variables and differentiable functions, or can only use simple but high variance score-function estimators. To overcome these limitations, we introduce Storchastic, a new framework for AD of stochastic computation graphs. Storchastic allows the modeler to choose from a wide variety of gradient estimation methods at each sampling step, to optimally reduce the variance of the gradient estimates. Furthermore, Storchastic is provably unbiased for estimation of any-order gradients, and generalizes variance reduction techniques to higher-order gradient estimates. Finally, we implement Storchastic as a PyTorch library at https://github.com/HEmile/storchastic.

研究の動機と目的

• 連続変数に限定される既存の確率的 AD フレームワークの制限や、スコア関数推定器による高分散の問題を解消すること。
• モデル開発者が各サンプリングステップで最適な勾配推定手法を選択可能となるようにし、分散を最小限に抑えつつバイアスを導入しないこと。
• 制御変数を含む分散低減技術を、確率的計算グラフにおける任意の順序の導関数推定に一般化すること。
• 推定器の構成要素に明確な条件を設定することで、n 階微分の不偏性を保証する形式的数学的基盤を提供すること。
• 実世界での展開と多様な勾配推定器の実験を可能にする、実用的でオープンソースの PyTorch ライブラリを実装すること。

提案手法

• 確率的勾配推定器を、提案分布、重み関数、勾配関数、制御変数の4つのコアな構成要素に分解する。
• 勾配伝搬の厳密な解析を可能にする数学的フレームワークを用いて、前向きモード AD の評価を形式化する。
• 4つの構成要素が特定の条件を満たす場合、結果として得られる代理損失が、不偏な n 階微分推定をもたらすことを証明する。
• 成分分解を介して、測度値導関数（MVD）や SPSA といった高度な勾配推定技術をフレームワークに統合する。
• 制御変数技術を n 階微分推定に拡張することで、高階導関数の分散低減を実現する。
• フレームワークを PyTorch ライブラリとして実装し、ディープラーニングワークフローおよび動的計算グラフとのシームレスな統合を可能にする。

実験結果

リサーチクエスチョン

• RQ1 任意の順序の導関数に対して不偏性を保証しつつ、多様な勾配推定手法をサポートできる一般化された確率的 AD フレームワークを設計可能か？
• RQ2 制御変数のような分散低減技術を、確率的計算グラフにおける高階導関数推定に体系的に拡張できるか？
• RQ3 勾配推定器の4つの構成要素にどのような形式的条件を課すと、n 階微分の不偏性が保証されるか？
• RQ4 測度値導関数や SPSA といった高度な手法を、1つの合成可能なフレームワークに統合できるか？
• RQ5 モデル開発者が各サンプリングステップで異なる勾配推定器を自由に選択しても、勾配推定の正しさに影響を及えないか？

主な発見

• Storchastic は、4つの推定器構成要素が特定の数学的条件を満たす場合、任意の順序の導関数推定が不偏のまま保たれることを形式的に証明した最初のフレームワークである。
• フレームワークは、再パラメータライゼーション、スコア関数推定器、RELAX、MAPO、順序なし集合推定器、測度値導関数を含む多様な勾配推定手法をサポートする。
• Storchastic は、制御変数に基づく分散低減を任意の順序の導関数推定に拡張可能であり、本分野における画期的な貢献である。
• PyTorch での実装により、モデル開発者がコード変更を最小限に抑えつつ、異なる推定器に切り替えられる。VAE 実験では、1行の変更で異なる手法を比較可能であることを示している。
• MNIST VAE における実験結果から、フレームワークが多様な推定器を正しくサポートしており、期待される性能順序が先行研究と一致していることが確認された。これは、限られた訓練エポック数とハイパーパramータチューニングなしでも成立する。
• フレームワークは、SPSA と測度値導関数を確率的 AD パイプラインに成功裏に統合し、その汎用性と拡張性を実証した。