[論文レビュー] TensorFlow Distributions

Joshua V. Dillon, Ian Langmore+8 more|arXiv (Cornell University)|Nov 28, 2017

Gaussian Processes and Bayesian Inference参考文献 41被引用数 240

ひとこと要約

TensorFlow Distributions は、TensorFlow で高速かつ微分可能な確率的プログラミングのための2つのコア抽象概念—DistributionsとBijectors—を提供します。これにより、複雑で高次元の分布と変換をモジュール的に構築でき、Edward のようなバックエンドを支えます。

ABSTRACT

The TensorFlow Distributions library implements a vision of probability theory adapted to the modern deep-learning paradigm of end-to-end differentiable computation. Building on two basic abstractions, it offers flexible building blocks for probabilistic computation. Distributions provide fast, numerically stable methods for generating samples and computing statistics, e.g., log density. Bijectors provide composable volume-tracking transformations with automatic caching. Together these enable modular construction of high dimensional distributions and transformations not possible with previous libraries (e.g., pixelCNNs, autoregressive flows, and reversible residual networks). They are the workhorse behind deep probabilistic programming systems like Edward and empower fast black-box inference in probabilistic models built on deep-network components. TensorFlow Distributions has proven an important part of the TensorFlow toolkit within Google and in the broader deep learning community.

研究の動機と目的

• TensorFlow エコシステム内でエンドツーエンドの微分可能な確率的プログラミングを可能にする。
• 多数の分布に対し高速で数値的に安定したサンプリング、対数密度、統計量を提供する。
• バッチ処理、自動微分、アクセラレータ（GPU/TPU）対応をサポート。
• 複雑な分布を効率的に構築するための合成可能な変換（Bijectors）を提供。
• Edward や TensorFlow Estimator などの高位レベルツールと統合し、スケーラブルな研究および実運用を可能にする。

提案手法

• 2つの抽象概念を導入する: Distribution（高速サンプリングと log_prob を持つ 60+ の分布）と Bijector（22 の合成可能で微分可能な変換）。
• ベクトル化演算とブロードキャストを可能にする形状意味論（sample、batch、event）を定義。
• 可能な場合に再参数化を介してエンドツーエンドの微分を提供する、デバイス固有の C++ カーネルによるサンプリングを実装。
• 分布の関数（高階分布）および分布関数量（例: エントロピー、KL）を提供。
• TransformedDistribution と Chain/Inver t Bijectors を活用し、モジュール式の組合せと効率的な密度計算を実現。
• 重要サンプリングと変分推論を高速化するために、変換後のサンプルと log-determinants のキャッシュを組み込む。

実験結果

リサーチクエスチョン

• RQ1 深層確率的プログラミングのために、速く、微分可能で、かつスケーラブルな確率ライブラリをどのように設計できるか？
• RQ2 分布と変換をどのように組み合わせて、豊かで高次元の確率モデルを表現できるか？
• RQ3 数値的安定性を維持しつつ、複雑なモデルをモジュール的に再利用可能に構築するための抽象化（Distributions と Bijectors）は何か？
• RQ4 このようなライブラリを広範な TensorFlow エコシステムやアクセラレータハードウェアとどのように統合できるか？
• RQ5 深層学習の文脈における高階分布と分布関数量の実用的な利点と制限は何か？

主な発見

• このライブラリは、高速サンプリングと log-density の計算を備えた約60の分布と、22の合成可能な Bijectors を提供します。
• Distributions と Bijectors は、VAE、自己回帰フロー、PixelCNN ベースのアーキテクチャなど、複雑なモデルをモジュール式でエンドツーエンド微分可能に構築することを可能にします。
• 形状意味論（sample、batch、event）は、大規模テンソル全体での直感的なベクトル化とブロードキャストを可能にします。
• サンプリングはデバイス固有のカーネルで実装され、確率的ノードを介した有効な逆伝播を可能にする再パラメータ化をサポートします。
• Bijectors は変換と log-determinants を自動的にキャッシュし、サンプリングベースの推論や変分法の効率を改善します。
• TensorFlow Distributions は TensorFlow のコンポーネント（layers、data pipelines、serving、visualization）と統合され、Edward のバックエンドとして機能します。