[論文レビュー] Trajectory Inference via Mean-field Langevin in Path Space

Lénaïc Chizat, Stephen X. Zhang+2 more|arXiv (Cornell University)|May 14, 2022

Single-cell and spatial transcriptomics被引用数 4

ひとこと要約

本稿では、ウィーナー測度に関する最小エントロピー推定量を経路空間で計算するために、スクリーンフリーで平均場ランジエビアン（MFL）ダイナミクスのアルゴリズムを提案する。この手法はシュレーディンガー橋とノイズ付き勾配降下法を活用し、スパースで時間分解能の高いデータからの一貫した推論を可能にする。単一細胞RNA-Seqや質量変動モデリングへの応用が可能である。

ABSTRACT

Trajectory inference aims at recovering the dynamics of a population from snapshots of its temporal marginals. To solve this task, a min-entropy estimator relative to the Wiener measure in path space was introduced by Lavenant et al. arXiv:2102.09204, and shown to consistently recover the dynamics of a large class of drift-diffusion processes from the solution of an infinite dimensional convex optimization problem. In this paper, we introduce a grid-free algorithm to compute this estimator. Our method consists in a family of point clouds (one per snapshot) coupled via Schrödinger bridges which evolve with noisy gradient descent. We study the mean-field limit of the dynamics and prove its global convergence to the desired estimator. Overall, this leads to an inference method with end-to-end theoretical guarantees that solves an interpretable model for trajectory inference. We also present how to adapt the method to deal with mass variations, a useful extension when dealing with single cell RNA-sequencing data where cells can branch and die.

研究の動機と目的

• 経路空間における最小エントロピー推定量を計算するスクリーンフリーでスケーラブルなアルゴリズムの開発。これは無限次元の凸最適化のため、計算的に実行不可能である。
• 提案されたMFLダイナミクスのグローバル収束を証明することで、経路推定に向けたエンドツーエンドの理論的保証を提供すること。
• 細胞の分岐や死滅といった質量変動を扱えるようにこの手法を拡張すること。これは単一細胞RNA-Seq応用で一般的な現象である。
• 非凸的な深層学習ベースの経路推定手法に代わる、理論的根拠があり解釈可能な代替手法を提供すること。これらの手法はきめ細やかな収束保証を欠いている。

提案手法

• 本手法は、各時間スナップショットごとに点群の集合を用い、エントロピー正則化付き最適輸送、すなわちシュレーディンガー・ブリッジによって結合することで、時間経過に伴うダイナミクスをモデル化する。
• 推定問題は、データ適合性とエントロピー正則化の両方を組み合わせた目的関数に対するノイズ付き勾配降下法によって進化する平均場ランジエビアンダイナミクスとして定式化される。
• ダイナミクスは、観測された周辺分布への忠実性と、ウィーナー測度との相対エントロピーによる滑らかさのバランスを取る確率微分方程式によって支配される。
• ステップサイズ $\eta$、ノイズレベル $\sigma$、温度 $\tau$ のためのアニーリングスケジュールを用いることで、安定な収束を保証する。
• 質量変動を扱えるように、修正されたダイナミクスが導入される。これには既知の増殖率 $g(t,x)$ を組み込み、細胞の増殖や死滅のモデリングが可能になる。
• 本手法は境界の明示的強制を一切行わないCPUベースのフレームワークで実装されており、観測データのサポート内での粒子の安定性に依存している。

実験結果

リサーチクエスチョン

• RQ1 最小エントロピー推定量を経路空間で計算するスクリーンフリーで確率的最適化手法を設計でき、理論的収束保証を備えているか？
• RQ2 平均場ランジエビアンダイナミクスはどのようにして真の経路推定解へグローバルに収束するのか。収束速度は何か？
• RQ3 本手法は、標準的な輸送モデルが捉えきれない細胞の分岐や死滅といった質量変動を扱えるように拡張可能か？
• RQ4 再構成精度とスパースサンプリングへのロバスト性という観点から、本手法のMFLは、Global Waddington-OTのような既存手法と比較してどのように性能を発揮するか？

主な発見

• MFLダイナミクスは、最小エントロピー推定量へ指数関数的収束を示し、経路推定に向けた強力な理論的保証を確立する。
• 本手法は合成データおよび実際の単一細胞データにおいて、最先端の性能を達成。真の値とのエネルギー距離は、MFLでは $\lambda = 0.025$、gWOTでは $\lambda = 0.01$ で最小化される。
• 本アルゴリズムは、分岐過程や時間依存ポテンシャルを含む複雑なダイナミクスを、1時間刻みあたり50粒子のわずかな粒子数でも効果的に再構成できる。
• 既知の増殖率 $g(t,x)$ の組み込みにより、細胞増殖などの質量変動の正確なモデリングが可能となり、生物学的に現実的な状況での再構成忠実度が向上する。
• 境界条件を明示的に設定する必要がなく、最適化の全過程で粒子が観測データのサポート内に留まるため、安定性が保たれる。
• ハイパーパrameterの選択に強く、複数のデータセットおよびパrameter設定において一貫した性能を発揮する。エネルギー距離と視覚的再構成品質の両面から検証済み。