[論文レビュー] Scalable Inference in SDEs by Direct Matching of the Fokker-Planck-Kolmogorov Equation

Arno Solin, Ella Tamir+1 more|arXiv (Cornell University)|Oct 29, 2021

Gaussian Processes and Bayesian Inference被引用数 4

ひとこと要約

本稿では、Fokker-Planck-Kolmogorov (FPK) 方程式の一次モーメントと二次モーメントを直接一致させることで、確率的微分方程式 (SDEs) のスケーラブルでサンプリングを不要とする推論手法を提案する。SDE の解をモーメント上での決定的常微分方程式 (ODE) として定式化することで、Euler–Maruyama などのサンプリングベース手法と比較して、数個のオーダーの高速化を達成しながら、時系列および生成モデルタスクにおいて競争力のある性能を維持する。

ABSTRACT

Simulation-based techniques such as variants of stochastic Runge-Kutta are the de facto approach for inference with stochastic differential equations (SDEs) in machine learning. These methods are general-purpose and used with parametric and non-parametric models, and neural SDEs. Stochastic Runge-Kutta relies on the use of sampling schemes that can be inefficient in high dimensions. We address this issue by revisiting the classical SDE literature and derive direct approximations to the (typically intractable) Fokker-Planck-Kolmogorov equation by matching moments. We show how this workflow is fast, scales to high-dimensional latent spaces, and is applicable to scarce-data applications, where a non-parametric SDE with a driving Gaussian process velocity field specifies the model.

研究の動機と目的

• SDE の推論における計算非効率性、特に高次元およびデータが少ない状況での課題を解決すること。
• SDE の解の解が得られない Fokker-Planck-Kolmogorov (FPK) 方程式を直接近似する、非サンプリング的手法を開発すること。
• ガウス分布の仮定を用いることで、潜在空間 SDE モデルにおける高速でスケーラブルな推論を可能とすること。
• ドリフト項にガウス過程 (GP) プライアを組み合わせ、モーメント一致推論を用いることで、SDE の機械学習分野への応用範囲を拡大すること。
• 経路ごとのサンプリングを回避しながら分布の正確さを保ちつつ、確率的 Runge–Kutta スキームの決定的代替手法を提供すること。

提案手法

• ドリフトベクトル場にガウス過程プライアを適用することで、潜在ダイナミクスの柔軟でプライア情報に基づいたモデリングを可能にする。
• 解分布の一次モーメント（平均）と二次モーメント（分散共分散）を一致させることで、Fokker-Planck-Kolmogorov (FPK) 方程式の閉形式近似を導出する。
• 得られたモーメントダイナミクスを決定的常微分方程式 (ODE) として解き、軌道サンプリングの必要を回避する。
• ガウス仮定密度近似を用いて SDE の時刻周辺分布を表現することで、高次元潜在空間における効率的な計算を実現する。
• 変分推論フレームワークを用いて実装し、エンコーダが初期モーメントを出力し、デコーダがモーメント一致 SDE ダイナミクスを用いる。
• 勾配を自動微分を用いてモーメント ODE を通して計算することで、尤度最大化または変分目的関数に基づいてモデルを最適化する。

実験結果

リサーチクエスチョン

• RQ1 Fokker-Planck-Kolmogorov 方程式の直接的モーメント一致は、SDE のサンプリングベース推論に対するスケーラブルな代替手段となり得るか？
• RQ2 高次元および低データ環境下で、モーメント一致推論は Euler–Maruyama サンプリングと比較して、精度と速度の面でどの程度優れるか？
• RQ3 非線形ドリフトおよび拡散を持つ SDE の真の時刻周辺分布を、ガウス仮定密度近似がどの程度正確に捉えることができるか？
• RQ4 GP プライアを用いた SDE とモーメント一致を組み合わせることで、データが少ない応用分野においても有効なモデリングが可能か？
• RQ5 提案手法は、計算コストを削減しながらも、時系列および生成モデルベンチマークで競争力のある性能を維持できるか？

主な発見

• MOCAP データセットにおいて、同じアーキテクチャを用いた場合、モーメント一致手法はテスト MSE 7.55±0.05 を達成し、ベースラインの潜在 SDE (4.03±0.20) を上回った。これは、仮定密度制約下でも堅牢であることを示している。
• ウォルクロックタイミングでは、1回の評価に GPU で 6.0±0.1 秒、CPU で 1.8±0.1 秒を要したのに対し、200 パスの Euler–Maruyama は 37.1±0.1 秒を要した。
• 回転 MNIST タスクでは、モーメント一致手法は MSE 0.051±0.007 および NLPD 52.7±9.5 を達成し、軌道サンプリングを回避しても競争力のある性能を示した。
• 高次元潜在空間へのスケーリングにおいても効率的であり、計算コストは次元数に比例して線形に増加（O(d)）する一方、サンプリング手法は次元数の増加に伴い性能が著しく低下する。
• サンプリングベースの SDE ソルバーと比較して、推論時間を数個のオーダーで短縮した。これにより、リアルタイムおよび大規模応用に実用的であることが示された。
• GP プライアがドリフト場に強いインダクティブバイアスを提供するため、低データ環境でも本手法が有効であることが示された。