[論文レビュー] Joint Distributions for TensorFlow Probability

Dan Piponi, Dave Moore+1 more|arXiv (Cornell University)|Jan 22, 2020

Bayesian Modeling and Causal Inference参考文献 16被引用数 11

ひとこと要約

この論文では、TensorFlow ProbabilityにおけるJointDistributionsを紹介する。JointDistributionsは、複数の宣言的および命令的インターフェースを通じて、柔軟で組み合わせ可能な確率的グラフィカルモデルの指定を可能にするフレームワークである。TensorFlow Distributionsを拡張して構造的で、バッチ処理可能かつ微分可能な同時分布をサポートすることで、さまざまなモデルアーキテクチャにおいて、ハードウェア加速された効率的な推論を実現する。ベクトル化されたサンプリングと勾配計算により、推論が統合される。

ABSTRACT

A central tenet of probabilistic programming is that a model is specified exactly once in a canonical representation which is usable by inference algorithms. We describe JointDistributions, a family of declarative representations of directed graphical models in TensorFlow Probability.

研究の動機と目的

• TensorFlow Probabilityにおける複雑な確率的モデルを指定するための統一的で拡張可能なインターフェースを提供すること。
• 1つの分布契約の下で、宣言的、命令的、コルーチンベースの複数のモデル指定のスタイル—-をサポートすること。
• TensorFlowのバッチおよび形状の意味論を構造的な同時分布に拡張することで、効率的なベクトル化された推論を可能にすること。
• 標準化されたインターフェースを通じてモデル定義と推論アルゴリズムを分離し、微分可能性と勾配サポートを維持すること。
• 自動ベクトル化を通じて、マルチチェーンMCMCやマルチサンプル変分境界といったスケーラブルな推論パラダイムを促進すること。

提案手法

• JointDistributionsは、テンソルの構造的出力（タプル、辞書、リスト）をサポートするようにTensorFlow Distribution APIを拡張し、イベント、バッチ、サンプルの形状の意味論を保持する。
• 3つの主要なスタイルが実装されている：JointDistributionNamed（辞書ベース）、JointDistributionSequential（リストベース）、JointDistributionCoroutine（ジェネレータベース）、それぞれが異なるコーディングスタイルを可能にする。
• サンプル形状の伝播を通じて、モデルの自動ベクトル化が可能で、効率的なバッチ化されたサンプリングと対数確率計算が実現される。
• sample_distributionsを導入し、条件付き分布とそのサンプル値の両方を返すことで、局所的な対数密度評価などの細粒度な確率的計算が可能になる。
• TensorFlowの自動微分を活用して、log_probとサンプルの勾配をモデルパラメータに関して自動的に計算し、HMC やブラックボックス変分推論のような勾配ベースの推論が可能になる。
• モデルパラメータはtf.VariableおよびTransformedVariableによって管理され、微分可能性と制約（例：非負のスケール）が保証され、tfp.math.minimizeを通じてエンドツーエンドの最適化が可能になる。

実験結果

リサーチクエスチョン

• RQ1 確率的モデルを、人間が読みやすく、かつハードウェア加速推論とネイティブに互換性を持つ方法でどのように指定できるか？
• RQ2 複数の指定イディオム（例：命令的 vs 宣言的）が、1つの統一インターフェースの下で相互に動作可能となる設計パターンは何か？
• RQ3 複雑でネストされたモデルにおいて、構造的同時分布がイベント、バッチ、サンプルの形状意味論を一貫して保持できるか？
• RQ4 自動ベクトル化とバッチ処理を通じて、同時分布がどの程度スケーラブルな推論を可能にするか？
• RQ5 微分可能性と制約を維持しながら、モデルパラメータを効率的に最適化する方法は何か？

主な発見

• JointDistributionsは、辞書ベース、リストベース、コルーチンベースの複数の指定スタイルをサポートする統一的で拡張可能なインターフェースを提供し、推論アルゴリズムとの相互運用性を損なわない。
• フレームワークは、複雑でネストされた同時分布においても、ハードウェアと自動バッチ処理により加速された効率的なベクトル化されたサンプリングと対数確率計算を可能にする。
• sample_distributionsメソッドにより、ユーザーは中間の条件付き分布とその値にアクセスでき、個別成分の対数密度計算のような細粒度の分析が可能になる。
• log_probとサンプルの勾配は、逆方向自動微分により自動的に計算され、ハミルトニアンモンテカルロやブラックボックス変分推論とのシームレスな統合が可能になる。
• tf.VariableまたはTransformedVariableとして表現されたモデルパラメータは自動的に追跡され、微分可能であり、標準的なTensorFlowオプティマイザを通じてエンドツーエンドの最適化が可能になる。
• JointDistributionオブジェクトのネスティングはネイティブにサポートされており、例えば、その構成要素自体が同時分布であるようなモジュラで階層的なモデル構築が可能になる。