[論文レビュー] A Computational Framework for Solving Wasserstein Lagrangian Flows

Kirill Neklyudov, Rob Brekelmans+4 more|arXiv (Cornell University)|Oct 16, 2023

Model Reduction and Neural Networks被引用数 4

ひとこと要約

本論文は、ラグランジアン作用関数の双対表現を活用することで、軌道のバックプロパゲーションや最適なカップリングへのアクセスを必要とせず、効率的な最適化を可能にする、Wassersteinラグランジュ流を統一的に解くためのディープラーニングフレームワークを提案する。この手法は、カスタマイズ可能な運動エネルギーおよびポテンシャルエネルギー項を用いて多様な事前知識を柔軟に組み込むことができ、単一細胞軌道推定において先行手法を上回る性能を発揮する。

ABSTRACT

The dynamical formulation of the optimal transport can be extended through various choices of the underlying geometry (kinetic energy), and the regularization of density paths (potential energy). These combinations yield different variational problems (Lagrangians), encompassing many variations of the optimal transport problem such as the Schrödinger bridge, unbalanced optimal transport, and optimal transport with physical constraints, among others. In general, the optimal density path is unknown, and solving these variational problems can be computationally challenging. We propose a novel deep learning based framework approaching all of these problems from a unified perspective. Leveraging the dual formulation of the Lagrangians, our method does not require simulating or backpropagating through the trajectories of the learned dynamics, and does not need access to optimal couplings. We showcase the versatility of the proposed framework by outperforming previous approaches for the single-cell trajectory inference, where incorporating prior knowledge into the dynamics is crucial for correct predictions.

研究の動機と目的

• 単一細胞RNA-Seqデータにおける希釈で、周辺観測値からの軌道推定という不適切に定式化された問題に対処すること。
• シュレーディンガー・ブリッジ、アンバランスなOT、物理的制約付きOTといった多様な最適輸送の変種を、単一の計算フレームワークで統一すること。
• ラグランジアン定式化におけるカスタマイズ可能な運動エネルギーおよびポテンシャルエネルギー項を通じて、ダイナミクスに事前知識を柔軟に組み込むこと。
• 個々の粒子軌道のシミュレーションやバックプロパゲーションを回避する、微分可能でスケーラブルな手法を開発すること。
• 双対最適化を活用することで、単一細胞軌道推定および画像生成において優れた性能を示すこと。

提案手法

• 軌道推定を、運動エネルギー（Wasserstein-2またはWasserstein-Fisher-Rao幾何）とポテンシャルエネルギー（例：エントロピーまたは空間的ポテンシャル）の組み合わせからなるラグランジアン作用関数の最小化問題として定式化する。
• ラグランジアンの双対表現を用いることで、個々の粒子軌道のシミュレーションや微分を必要としない、取り扱いやすい最適化目的関数を導出する。
• 深層ニューラルネットワークを用いてベクトル場 $ s_t(x, \theta) $ と密度経路 $ \rho_t(x, \eta) $ をパラメータ化し、周辺制約に一致するように双対目的関数で学習する。
• 最適化は、明示的な軌道展開やカップリング推定を必要とせず、エンドツーエンドの学習を可能にする双対目的関数を用いる。
• 画像生成の文脈では、ハミルトニアン・ジャコビ条件により1ステップでのサンプリングが可能であり、零加速度の下で $ X_1 = X_0 + \nabla s_0(X_0) $ が成り立つ。
• エネルギー項の組み合わせにより、特定の生物学的・物理的制約に適合する誘導的バイアスを備えた、複数の構成をサポートする。

実験結果

リサーチクエスチョン

• RQ1 統一的なディープラーニングフレームワークは、アンバランスなOT、シュレーディンガー・ブリッジ、物理的制約付きOTといった多様な最適輸送問題を、同一の定式化で解けるか？
• RQ2 細胞ダイナミクスに関する事前知識を、変分フレームワークにおける軌道推定に柔軟に埋め込む方法は何か？
• RQ3 ラグランジアンの双対最適化は、軌道のバックプロパゲーションの計算負荷を回避しつつ、精度を維持できるか？
• RQ4 本フレームワークは、既存の手法と比較して、単一細胞RNA-seqデータにおける軌道推定性能をどの程度向上させるか？
• RQ5 本フレームワークは、繰り返しODEソルバーを必要とせず、最適輸送マップを用いて1ステップでの画像生成を可能にするか？

主な発見

• 提案されたフレームワークは、特化したラグランジアンを用いて生物学的知識を組み込む場合に、先行手法を上回る性能を発揮する。
• Embryoid Body、CITE-seq、Multiomeデータを含む低次元および高次元の単一細胞データセットにおいて、最先端の性能を達成する。
• 画像生成の文脈では、ハミルトニアン・ジャコビ条件により1ステップでのサンプリングが可能であり、反復的ODEソルバーと同等の品質を、1回の関数評価で達成する。
• 双対最適化アプローチにより、軌道のバックプロパゲーションやカップリング推定を回避し、計算コストを顕著に削減する。
• ラグランジアンの構成要素を調整することで、アンバランスなOTやシュレーディンガー・ブリッジといった複数の最適輸送の変種に、本フレームワークが柔軟に一般化可能である。
• MNISTおよびCIFAR-10における実験から、学習済みベクトル場を用いた1ステップ生成が、反復的拡散モデルと同等の高精細なサンプルを生成できることを示した。