[論文レビュー] Automatic Differentiation of Programs with Discrete Randomness

Gaurav Arya, Moritz Schauer+2 more|arXiv (Cornell University)|Oct 16, 2022

Data Stream Mining Techniques被引用数 12

ひとこと要約

本稿では、離散的な確率的動作を伴うプログラムにおける再パrameterizationに基づく自動微分手法を提案する。この手法により、確率的経路を介して無限小の摂動を伝播させることで、不偏かつ低分散の勾配推定が可能になる。本手法は、マルコフ連鎖、エージェントベースモデル、パーティクルフィルタなどの確率的シミュレーションにおいて、前向きモードおよび逆向きモードのADをサポートし、剪定およびスムージング技術を用いてO(1)の計算オーバーヘッドを達成するオープンソースのJulia実装を提供する。

ABSTRACT

Automatic differentiation (AD), a technique for constructing new programs which compute the derivative of an original program, has become ubiquitous throughout scientific computing and deep learning due to the improved performance afforded by gradient-based optimization. However, AD systems have been restricted to the subset of programs that have a continuous dependence on parameters. Programs that have discrete stochastic behaviors governed by distribution parameters, such as flipping a coin with probability $p$ of being heads, pose a challenge to these systems because the connection between the result (heads vs tails) and the parameters ($p$) is fundamentally discrete. In this paper we develop a new reparameterization-based methodology that allows for generating programs whose expectation is the derivative of the expectation of the original program. We showcase how this method gives an unbiased and low-variance estimator which is as automated as traditional AD mechanisms. We demonstrate unbiased forward-mode AD of discrete-time Markov chains, agent-based models such as Conway's Game of Life, and unbiased reverse-mode AD of a particle filter. Our code package is available at https://github.com/gaurav-arya/StochasticAD.jl.

研究の動機と目的

• ベルヌーイ試行などの離散的確率的動作を伴うプログラムの微分を扱う挑戦に応えること。古典的自動微分は、非微分可能で不連続な出力のため、これに失敗する。
• 確率的プログラムの離散的構造を保ちつつ、統計的期待値のための勾配計算を可能にする手法を開発すること。
• 連続的パスワイズADに匹敵する不偏かつ低分散の勾配推定器を、連続的緩和化や調整可能なパラメータに依存せずに達成すること。
• マルコフ連鎖やパーティクルフィルタを含む複雑な確率的システムにおいて、前向きおよび逆向きモードADを最小限の計算オーバーヘッドでサポートすること。
• 既存のADツールとシームレスに統合され、確率的モデリングにおける勾配ベース最適化を可能にする、合成的でエンドツーエンドの微分可能なフレームワークを提供すること。

提案手法

• 無限小のパrameter摂動による結果確率の変化割合を追跡する『確率的微分』という概念を導入。これはパスワイズ推定器である。
• 再パラメータ化を用いて、離散的確率的変数を一様乱数変数とパrameterの決定的関数として表現し、確率的経路を介した勾配伝播を可能にする。
• 離散的イベントに対応するための『スムージングされた確率的微分』技術を適用。これにより、相関のあるサンプリング経路を用いて不偏な勾配推定が可能になる。
• オンラインでの剪定戦略を採用し、アクティブでないか低影響の経路を動的に削除することで、前向きモードADにおけるO(1)の計算オーバーヘッドを維持する。
• スムージングとパスワイズ推定を組み合わせて逆向きモードADメカニズムを導出し、ストレートスラッグ勾配推定器を特別なケースとして回復する。
• ForwardDiff.jl、Zygote.jl、ChainRulesCore.jlと統合可能なJuliaパッケージStochasticAD.jlとして本手法を実装。これにより、前向きおよび逆向きモードの微分がシームレスに可能になる。

実験結果

リサーチクエスチョン

• RQ1 古典的微分が定義されないベルヌーイ分布やポisson分布の抽出を伴う離散的確率的コンponentsを有するプログラムに、自動微分を拡張可能か？
• RQ2 離散的確率的要素を有する確率的プログラムの期待値に対して、不偏かつ低分散の勾配推定器を構築する方法は何か？
• RQ3 エージェントベースモデルやマルコフ連鎖のような複雑な確率的シミュレーションに、高い計算コストを伴わずに前向きモードADを効率的に適用可能か？
• RQ4 連続的パスワイズ手法と同等の性能を達成できる離散的確率的プログラム向けの逆向きモードADメカニズムを導出可能か？
• RQ5 離散的確率的プログラムの勾配推定器における分散を最小限に抑えるために、相関のあるサンプリング経路を体系的に生成する方法は何か？

主な発見

• 提案手法は、摂動サイズがゼロに近づくにつれて分散が有界のまま維持される、離散的確率的プログラムの期待値に対する不偏勾配推定を達成する。
• 本フレームワークは、剪定を用いてO(1)の計算オーバーヘッドで、離散時間マルコフ連鎖やコンウェイの人生ゲームのようなエージェントベースモデルの前向きモードADをサポートする。
• 本手法により、パーティクルフィルタのエンドツーエンドの逆向きモードADが可能になり、既存の研究で知られている手法が特別なケースとして回復される。
• スムージングされた確率的微分技術により、前向きおよび逆向きモードの両方の微分が可能になり、従来のスコア関数推定器の制限を克服する。
• 実験的評価では、有限差分法よりも分散が小さく、チューニング可能なパラメータや学習可能なパラメータを必要としない点で、連続的緩和化手法と同等の性能を示す。
• オープンソースのStochasticAD.jlパッケージは、確率的シミュレーションや確率的推論を含む多様な応用分野において、本手法の実用性を示している。