[論文レビュー] Path Integral Sampler: a stochastic control approach for sampling

Qinsheng Zhang, Yongxin Chen|arXiv (Cornell University)|Nov 30, 2021

Model Reduction and Neural Networks参考文献 55被引用数 6

ひとこと要約

パス・インテグラル・サンプラー（PIS）は、ニューラルネットワークで制御されるSDEを用いてシュレーディンガー・ブリッジ問題を解くことで、正規化されていないターゲット密度からサンプルを生成する、新しい確率的制御ベースのサンプリングアルゴリズムである。エンド・ツー・エンドの学習が可能で自由なアーキテクチャを採用でき、制御が最適でない場合でも理論的なワーシングシュタイン距離の保証を提供する。また、経路積分に基づく重要度付き重みを用いてバイアスを補正し、従来の手法に比べてサンプルの質と効率性において優れている。

ABSTRACT

We present Path Integral Sampler~(PIS), a novel algorithm to draw samples from unnormalized probability density functions. The PIS is built on the Schrödinger bridge problem which aims to recover the most likely evolution of a diffusion process given its initial distribution and terminal distribution. The PIS draws samples from the initial distribution and then propagates the samples through the Schrödinger bridge to reach the terminal distribution. Applying the Girsanov theorem, with a simple prior diffusion, we formulate the PIS as a stochastic optimal control problem whose running cost is the control energy and terminal cost is chosen according to the target distribution. By modeling the control as a neural network, we establish a sampling algorithm that can be trained end-to-end. We provide theoretical justification of the sampling quality of PIS in terms of Wasserstein distance when sub-optimal control is used. Moreover, the path integrals theory is used to compute importance weights of the samples to compensate for the bias induced by the sub-optimality of the controller and time-discretization. We experimentally demonstrate the advantages of PIS compared with other start-of-the-art sampling methods on a variety of tasks.

研究の動機と目的

• 正規化されていない確率密度関数に対して、明示的な密度モデルを必要としない柔軟でエンド・ツー・エンドで学習可能なサンプリング手法の開発。
• シュレーディンガー・ブリッジフレームワークを活用し、ターゲット分布に結びつく終端コストを持つ確率的最適制御問題としてサンプリングを定式化すること。
• 従来のMCMCに比べてステップ数を減らしつつ、高品質なサンプリングを実現し、理論的な性能バインディングを提供すること。
• 経路積分理論を用いて、制御ポリシーの非最適性と時間離散化のバイアスを補正するための重要度重みを計算すること。
• 明示的な密度制約を課える変分インファレンスの限界や、混合が遅く、トレーニングが難しい提案分布を必要とするMCMCの課題を、暗黙的でスケーラブルで微分可能なサンプリングフレームワークで克服すること。

提案手法

• 走行コストを制御エネルギーとし、終端コストをターゲット密度から導出する確率的最適制御問題としてサンプリング問題を定式化する。
• 自由なアーキテクチャ設計とエンド・ツー・エンドの学習を可能にするために、制御ポリシーをニューラルネットワークでモデル化する。
• ギルサノフの定理を用いてSDEのダイナミクスを再重み付けし、初期のウィーナー過程を制御された拡散過程を通じてターゲット分布に結びつける。
• 経路積分理論を適用して、制御ポリシーの非最適性と離散化誤差に起因するバイアスを補正するための重要度重みを計算する。
• 制御されたSDEの終端分布とターゲット測度の間のKLダイバージェンスを最適化することで、コントローラーを学習する。
• 特に低次元の場合に安定性を確保するため、アドジョイントを用いないEuler-MaruyamaによるSDE統合を採用し、メモリ制約がある場合にはReversible Heunを用いる。

実験結果

リサーチクエスチョン

• RQ1 シュレーディンガー・ブリッジ問題に基づく確率的制御フレームワークを用いて、正規化されていない密度関数のための柔軟でエンド・ツー・エンドで学習可能なサンプリングアルゴリズムを設計できるか？
• RQ2 非最適な制御ポリシーはどのように補正され、漸近的にバイアスのないサンプリングが保証されるか？
• RQ3 非最適なコントローラーを用いた場合に、サンプリング品質について理論的な性能保証をどのように提供できるか？
• RQ4 時間ホライズンTと時間刻み幅Δtの選択が、学習の安定性とサンプル品質にどのように影響するか？
• RQ5 PISは、MCMC やノーマライジング・フローなどの最先端手法に比べて、多様なタスクにおいてサンプルの質と効率性で優れているか？

主な発見

• PISは、非最適な制御ポリシーを用いても、ワーシングシュタイン距離の観点からサンプリング品質に関する理論的保証を提供する。
• 経路積分に基づく重要度重みの使用は、非最適な制御と時間離散化によって生じるバイアスを効果的に補正する。
• 大きな時間ホライズンTは学習の安定性とモードカバレッジを向上させるが、小さなTは高リプシッツ性を持つ制御ポリシーをもたらし、トレーニングが困難になる。
• 小さな時間刻み幅Δtは離散化誤差を低減し、特に統合ステップ数を増やすとサンプル品質が向上する。
• ガウス混合分布、ファンネル分布、MNIST や CIFAR-10 といった実世界のデータセットを含むタスクにおいて、PISはベースライン手法に比べてサンプルの質と効率性で優れている。
• バッチ正規化と勾配クリッピングを用いた学習は安定性を向上させ、実際の応用ではアドジョイントを用いないEuler-Maruyama統合がアドジョイントベースの手法よりも信頼性が高い。