[論文レビュー] Diffusion Earth Mover's Distance and Distribution Embeddings

Alexander Tong, Guillaume Huguet+6 more|PubMed|Feb 25, 2021

Cell Image Analysis Techniques参考文献 8被引用数 10

ひとこと要約

本稿では、共通のデータ類似度グラフ上での分布としてモデル化することで、高次元データセット間の距離を高速かつ微分可能に計算する方法である Diffusion Earth Mover's Distance (Diffusion EMD) を提案する。計算時間は $\tilde{O}(n)$ に抑えられ、リーマン多様体上での標準的 EMD と位相的に同等であることが証明され、大規模な単細胞患者データセットの効率的埋め込みを可能にした。210サンプルのCOVID-19コhortにおいて、臨床的非均質性を、従来手法よりも2桁以上速く同定した。

ABSTRACT

We propose a new fast method of measuring distances between large numbers of related high dimensional datasets called the Diffusion Earth Mover's Distance (EMD). We model the datasets as distributions supported on common data graph that is derived from the affinity matrix computed on the combined data. In such cases where the graph is a discretization of an underlying Riemannian closed manifold, we prove that Diffusion EMD is topologically equivalent to the standard EMD with a geodesic ground distance. Diffusion EMD can be computed in $ ilde{O}(n)$ time and is more accurate than similarly fast algorithms such as tree-based EMDs. We also show Diffusion EMD is fully differentiable, making it amenable to future uses in gradient-descent frameworks such as deep neural networks. Finally, we demonstrate an application of Diffusion EMD to single cell data collected from 210 COVID-19 patient samples at Yale New Haven Hospital. Here, Diffusion EMD can derive distances between patients on the manifold of cells at least two orders of magnitude faster than equally accurate methods. This distance matrix between patients can be embedded into a higher level patient manifold which uncovers structure and heterogeneity in patients. More generally, Diffusion EMD is applicable to all datasets that are massively collected in parallel in many medical and biological systems.

研究の動機と目的

• 大規模な高次元データセット（例：複数患者の単細胞オミクス）間の距離を効率的に計算する課題に対処すること。
• 計算が高速であり、かつ基礎となる多様体上で標準的 Earth Mover’s Distance と位相的に忠実な方法を開発すること。
• 患者または撹乱レベルの解析のため、全データセットを高次の多様体にスケーラブルに埋め込むこと。
• 勾配ベース最適化を用いたディープラーニングパイプラインに統合可能な微分可能フレームワークを提供すること。

提案手法

• 結合データセットの類似度行列から導出される共通のデータ類似度グラフ上に、確率分布としてデータセットをモデル化する。
• グラフ上の近似カーネル密度推定にマルチスケールの拡散カーネルを用い、密度推定間の修正 $L^1$ 距離を計算する。
• 最適輸送の Kantorovich-Rubinstein 双対定式を活用し、見せ役関数の最適化により計算を効率化する。
• グラフがリーマン多様体を離散化する場合、Diffusion EMD と標準的 EMD が地図距離（geodesic ground distance）において理論的に同等であることを証明する。
• 拡散ウェーブレットと階層的クラスタリングを用いて高速アルゴリズムを実装し、$\tilde{O}(n)$ の時間計算量を達成する。
• 勾配ベース学習システムへの統合を可能にするために、フレームワークの完全な微分可能性を保証する。

実験結果

リサーチクエスチョン

• RQ1 高速でグラフベースの EMD 変種は、リーマン多様体上での標準的 EMD と位相的忠実性を保たれるか？
• RQ2 高次元データセット間の最適輸送距離を、正確性を損なわず $\tilde{O}(n)$ 時間で計算できるか？
• RQ3 Diffusion EMD は、木ベースの手法などの既存の高速 EMD 近似法と比較して、精度と速度の両面でどれほど優れているか？
• RQ4 Diffusion EMD は大規模な単細胞研究における患者レベルの多様体埋め込みを意味的に可能にするか？
• RQ5 このフレームワークは完全に微分可能であり、ディープラーニングパイプラインへの統合が可能か？

主な発見

• Diffusion EMD は $\tilde{O}(n)$ の計算時間で実現し、$m$ 個の Wasserstein 近隣行列の計算複雑度を、妥当な仮定のもとで $O(m^2n^3)$ から $\tilde{O}(mn)$ に低減した。
• 球面 MNIST データセットにおいて、Spearman-$\rho$ および P@10 の両指標で ClusterTree や QuadTree を同等以上または上回り、著しく高速な実行時間を達成した。
• 210名の患者サンプルからなる単細胞 COVID-19 データセットにおいて、同等の精度を持つ手法と比較して、患者間距離行列の計算が少なくとも2桁以上高速であった。
• k-NN グラフにおけるラプラシアンの滑らかさが優れており、患者多様体埋め込みにおける基礎的構造のより良い保存を示した。
• Diffusion EMD は完全に微分可能であり、将来のエンドツーエンドのディープラーニングフレームワークでの利用が可能である。
• 理論的分析により、データグラフがリーマン多様体を離散化する場合、Diffusion EMD が地図距離を用いた標準的 EMD と位相的に同等であることが確認された。