[논문 리뷰] Diffusion Earth Mover's Distance and Distribution Embeddings

Alexander Tong, Guillaume Huguet외 6명|PubMed|2021. 02. 25.

Cell Image Analysis Techniques참고 문헌 8인용 수 10

한 줄 요약

이 논문은 공통 데이터 유사도 그래프 위에서 분포로 모델링함으로써 고차원 데이터셋 간의 거리를 신속하고 미분 가능한 방식으로 계산하는 Diffusion Earth Mover's Distance (Diffusion EMD)를 소개한다. 이는 $\tilde{O}(n)$ 계산 시간을 달성하고, 리만 다양체 위에서 표준 EMD와 위상적으로 동치임을 증명하며, 대규모 단일세포 환자 데이터셋의 효율적 임bedding을 가능하게 하여, 이전 방법보다 두 계급 정도 빠른 속도로 210개 샘플의 코로나19 코hort에서 임상적 이질성을 드러낸다.

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 Diffusion EMD는 기존의 트리 기반 방법과 같은 빠른 EMD 근사치보다 정확도와 속도에서 얼마나 뛰어나게 성능을 발휘하는가?
• RQ4 Diffusion EMD는 대규모 단일세포 연구에서 의미 있는 환자 수준의 만능 임베딩을 가능하게 하는가?
• RQ5 프레임워크는 완전히 미분 가능하여 기계학습 파이프라인에 통합 가능한가?

주요 결과

• Diffusion EMD는 $\tilde{O}(n)$ 계산 시간을 달성하여, $m$-by-$m$ 워셔스타인 이웃 행렬을 계산하는 복잡도를 합리적인 가정 하에 $O(m^2n^3)$에서 $\tilde{O}(mn)$으로 감소시켰다.
• 구면 MNIST 데이터셋에서, Diffusion EMD는 Spearman-$\rho$ 및 P@10 지표에서 ClusterTree 및 QuadTree를 따라하거나 능가했으며, 훨씬 더 빠른 런타임을 보였다.
• 210명의 환자 샘플로 구성된 단일세포 코로나19 데이터셋에서, Diffusion EMD는 동등한 정확도를 갖는 방법보다 적어도 두 계급 정도 더 빠른 속도로 환자 거리 행렬을 계산했다.
• k-NN 그래프에서 라플라시안 스무스니스가 뛰어나, 환자 만능 임베딩에서 기저 구조의 보존이 더 잘 이루어짐을 보였다.
• Diffusion EMD는 완전히 미분 가능했으며, 향후 엔드 투 엔드 기계학습 프레임워크에서의 활용이 가능했다.
• 이론적 분석을 통해, 데이터 그래프가 리만 다양체를 이산화할 경우, Diffusion EMD는 지오데식 기준 거리와 함께 표준 EMD와 위상적으로 동치임을 확인했다.