[論文レビュー] Manifold Interpolating Optimal-Transport Flows for Trajectory Inference

Guillaume Huguet, D S Magruder+5 more|arXiv (Cornell University)|Jun 29, 2022

Traffic Prediction and Management TechniquesEngineering被引用数 19

ひとこと要約

MIOFlowは、多様体学習と最適輸送を組み合わせることで、時間的に疎な静的スナップショットから確率的で連続的な集団ダイナミクスを学習するニューラルODEベースのフレームワークを提案する。地図的自己符号化器を用いて、距離が新たなマルチスケール拡散地図距離を反映する潜在空間にデータを埋め込む。これにより、複雑な多様体上での正確な補間と軌道推定が可能となり、正規化フローおよびシュレーディンガー・ブリッジよりも分岐や生物学的ダイナミクスのモデル化で優れる。

ABSTRACT

We present a method called Manifold Interpolating Optimal-Transport Flow (MIOFlow) that learns stochastic, continuous population dynamics from static snapshot samples taken at sporadic timepoints. MIOFlow combines dynamic models, manifold learning, and optimal transport by training neural ordinary differential equations (Neural ODE) to interpolate between static population snapshots as penalized by optimal transport with manifold ground distance. Further, we ensure that the flow follows the geometry by operating in the latent space of an autoencoder that we call a geodesic autoencoder (GAE). In GAE the latent space distance between points is regularized to match a novel multiscale geodesic distance on the data manifold that we define. We show that this method is superior to normalizing flows, Schrödinger bridges and other generative models that are designed to flow from noise to data in terms of interpolating between populations. Theoretically, we link these trajectories with dynamic optimal transport. We evaluate our method on simulated data with bifurcations and merges, as well as scRNA-seq data from embryoid body differentiation, and acute myeloid leukemia treatment.

研究の動機と目的

• 高次元データの静的で時系列順に整列したスナップショットから、連続的かつ確率的な集団ダイナミクスを推定する課題に対処すること。
• 単一細胞RNAシークエンシングなど、縦断的追跡が不可能なシステムにおける軌道推定をモデル化すること。
• 推定された軌道が、埋め込まれた高次元空間ではなく、元のデータ多様体の内在的幾何構造を尊重すること。
• ガウス事前分布や決定的フローを必要としない、計算的に効率的な手法を開発すること。これは、既存の生成モデルにおける一般的な制限要因である。
• 複雑な生物学的システムにおける、未観測の時系列ポイントおよび個々の細胞軌道の正確な補間と再構築を可能にすること。

提案手法

• MIOFlowは、異なる時刻における静的集団スナップショット間の連続的フローをモデル化するために、ニューラル常微分方程式（Neural ODEs）を用いる。
• 地図的自己符号化器（GAE）を用い、潜在空間を、データ多様体上での新しいマルチスケール拡散地図距離を保存するように正則化する。
• 拡散地図距離は、カーネルベースのグラフラプラシアンを介して定義され、無限サンプルの極限において、多様体上の真の地図距離に収束する。
• フローは、多様体上の基本距離を用いた最適輸送によってペナルティが課され、時刻間の輸送コストがデータの内在的幾何構造と一致するように保証される。
• この手法はGAEの潜在空間で動作し、低次元多様体上でダイナミクスを学習することで、計算コストを低減し、一般化性能を向上させる。
• 潜在空間における初期点のサンプリングによって、決定的フローと同様に明示的なノイズ注入を避ける自然な確率性が導入される。

実験結果

リサーチクエスチョン

• RQ1 ガウス事前分布を仮定せずに、時間的に疎なスナップショットから連続的かつ確率的な集団ダイナミクスを学習できるか？
• RQ2 推定された軌道が、元のデータ多様体の内在的幾何構造を尊重するように保証できるか？
• RQ3 潜在空間に、データ駆動型でマルチスケールの距離計測を構築でき、多様体上の真の地図距離に収束するか？
• RQ4 分岐や合流といった複雑なダイナミクスをモデル化する際、MIOFlowは正規化フロー、シュレーディンガー・ブリッジ、TrajectoryNetといった既存手法と比較してどのように性能を発揮するか？
• RQ5 MIOFlowは、分化や治療反応といった生物学的プロセス中の遺伝子発現の動的変化を捉え、個々の細胞軌道を正確に再構築できるか？

主な発見

• MIOFlowは、分岐や合流を伴う合成データにおいて、正規化フローおよびシュレーディンガー・ブリッジを上回り、滑らかでより地図的距離に正確な軌道を生成する。
• シミュレーションデータにおいて、MIOFlowは発散的軌道を正確にモデル化でき、フローが多様体の分岐構造を的確に捉えている。
• 胚体様体の分化に関するscRNA-seqデータでは、MIOFlowが既知の発生マーカーと整合する生物学的に妥当な系統軌道を再構築した。
• 急性骨髄性白血病（AML）治療データでは、生存細胞が白血病幹細胞（LSC）のシグネチャーへと移行しており、後期の時系列でSLPI、MYB、B2Mの発現が上昇していることがMIOFlowによって明らかになった。
• 拡散地図距離を用いた地図的自己符号化器は、正確な多様体再構築を可能にし、潜在空間におけるフローは生物学的に意味のある距離を保持しており、それが生物学的に解釈可能な遺伝子発現パターンに復号可能である。
• MIOFlowは、ヤコビアントレースの計算を回避するため、正規化フローに比べて関数評価1回あたりの推論が高速であり、環境次元kに対してO(k)のスケーリングを示し、O(k^2)ではなくなる。