[论文解读] 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 队列中揭示临床异质性，速度比先前方法快两个数量级。

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 在测地线基距离下具有理论等价性。
• 利用扩散小波和分层聚类实现快速算法，达到 $\tilde{O}(n)$ 时间复杂度。
• 确保框架的完全可微性，以实现与基于梯度学习系统的集成。

实验结果

研究问题

• RQ1 一种快速、基于图的 EMD 变体能否在黎曼流形上保持与标准 EMD 的拓扑保真度？
• RQ2 如何在不牺牲准确性的前提下，以 $\tilde{O}(n)$ 时间复杂度计算高维数据集之间的最优传输距离？
• RQ3 与现有快速 EMD 近似方法（如基于树的方法）相比，Diffusion EMD 在准确性和速度方面有多大的性能提升？
• RQ4 Diffusion EMD 是否能实现在大规模单细胞研究中具有意义的患者水平流形嵌入？
• RQ5 该框架是否完全可微分，从而支持集成到深度学习流水线中？

主要发现

• Diffusion EMD 实现了 $\tilde{O}(n)$ 的计算时间，在合理假设下，将计算 $m$ 个样本的 Wasserstein 邻居矩阵的时间复杂度从 $O(m^2n^3)$ 降低至 $\tilde{O}(mn)$。
• 在球面 MNIST 数据集上，Diffusion EMD 在 Spearman-$\rho$ 和 P@10 指标上达到或优于 ClusterTree 和 QuadTree，且运行时间显著更短。
• 在包含 210 例患者样本的单细胞 COVID-19 数据集中，Diffusion EMD 计算患者距离矩阵的速度比同等准确度的先前方法快至少两个数量级。
• 该方法在 k-NN 图中表现出更优的拉普拉斯平滑性，表明其在患者流形嵌入中更好地保留了底层结构。
• Diffusion EMD 具有完全可微性，为未来在端到端深度学习框架中的应用铺平了道路。
• 理论分析证实，当数据图对黎曼流形进行离散化时，Diffusion EMD 与标准 EMD 在测地线基距离下具有拓扑等价性。