[논문 리뷰] Linear-Time Gromov Wasserstein Distances using Low Rank Couplings and Costs

Meyer Scetbon, Gabriel Peyré외 1명|arXiv (Cornell University)|2021. 06. 02.

Topological and Geometric Data Analysis참고 문헌 40인용 수 8

한 줄 요약

이 논문은 저랭크 결합과 저랭크 비용 행렬을 활용하여 선형 시간 Gromov-Wasserstein (GW) 거리 계산을 제안한다. 이는 이질적 공간 내에서 점군의 효율적이고 확장 가능한 정렬을 가능하게 한다. 결합의 저랭크 분해와 입력 비용 행렬의 저랭크 근사화를 조합함으로써, 기존의 최첨단 엔트로픽 GW보다 수개의 주기 빠른 $O(n)$ 복잡도를 달성하면서도, 합성 및 실세계 데이터셋(예: 단세포 유전체 분석 데이터 포함)에서 경쟁 가능한 GW 손실과 정렬 정확도를 유지한다.

ABSTRACT

The ability to align points across two related yet incomparable point clouds (e.g. living in different spaces) plays an important role in machine learning. The Gromov-Wasserstein (GW) framework provides an increasingly popular answer to such problems, by seeking a low-distortion, geometry-preserving assignment between these points. As a non-convex, quadratic generalization of optimal transport (OT), GW is NP-hard. While practitioners often resort to solving GW approximately as a nested sequence of entropy-regularized OT problems, the cubic complexity (in the number $n$ of samples) of that approach is a roadblock. We show in this work how a recent variant of the OT problem that restricts the set of admissible couplings to those having a low-rank factorization is remarkably well suited to the resolution of GW: when applied to GW, we show that this approach is not only able to compute a stationary point of the GW problem in time $O(n^2)$, but also uniquely positioned to benefit from the knowledge that the initial cost matrices are low-rank, to yield a linear time $O(n)$ GW approximation. Our approach yields similar results, yet orders of magnitude faster computation than the SoTA entropic GW approaches, on both simulated and real data.

연구 동기 및 목표

• Gromov-Wasserstein (GW) 거리의 높은 계산 비용을 해결한다. GW는 NP-완전이며 일반적으로 $O(n^3)$의 비용을 요구하는 엔트로픽 정규화를 통해 해결된다.
• GW 솔버에서 반복적인 $C_t = A P_{t-1} B$ 계산이 발생하는 성능 저하 문제를 해결하기 위해, 비용 행렬과 결합 양쪽 모두에 저랭크 구조를 활용한다.
• 원래 GW 공식화에 충실하면서도 선형 시간 복잡도를 달성하는 확장 가능한 GW 근사법을 개발한다.
• 기존 엔트로픽 솔버가 처리할 수 없는 대규모 데이터셋(예: 단세포 유전체 분석)에 GW를 적용할 수 있도록 한다.
• 저랭크 제약 조건을 가진 결합이 엔트로픽 정규화와 유사한 GW 성능을 달성하면서도, 수개의 주기 빠른 속도를 제공함을 보여준다.

제안 방법

• 결합 행렬 $P$에 저랭크 분해를 적용하여, $U, V \in \mathbb{R}^{n \times r}$ 이면 $P = U V^T$ 형태로 제한함으로써 자유도를 감소시키고 최적화 속도를 향상시킨다.
• 입력 비용 행렬 $A$와 $B$(예: 제곱 유클리드 거리)를 저랭크 분해를 통해 근사화하며, 내재 차원성 $d, d' \ll n, m$을 가정한다.
• 저랭크 결합에 적합한 Sinkhorn 알고리즘의 변종을 사용하여, 저랭크 요소에 대한 교대 최소화를 통해 정규화된 GW 문제를 해결한다.
• 각 반복에서 $O(n^3)$ 연산을 피하기 위해, 저랭크 근사를 사용하여 합성 비용 행렬 $C_t = A P_{t-1} B$를 반복적으로 업데이트한다.
• GW 손실과 결합 랭크 $r$ 사이의 균형을 고려한 랭크 인식 최적화 기법을 도입하며, 정확도를 유지하기 위해 $r$을 적응적으로 선택한다.
• 저랭크 비용 행렬과 저랭크 결합이 $r \ll n$일 때 선형 시간 알고리즘을 이끌어내어, $O(n)$ 복잡도로 $O(n^3)$ 대비 빠른 계산을 가능하게 한다.

실험 결과

연구 질문

• RQ1 결합 행렬의 저랭크 분해가 GW 거리 계산의 계산 복잡도를 $O(n^3)$에서 $O(n)$으로 줄일 수 있는가?
• RQ2 입력 비용 행렬과 결합 행렬 양쪽에 저랭크 구조를 활용하면 GW 정확도를 유지하면서도 선형 시간 계산이 가능한가?
• RQ3 실세계 데이터셋에서 저랭크 GW의 성능은 최첨단 엔트로픽 GW와 GW 손실 및 정렬 품질 측면에서 어떻게 비교되는가?
• RQ4 제안된 방법은 엔트로픽 GW가 비현실적인 메모리 및 시간 제약으로 인해 처리 불가능한 대규모 데이터셋(예: 단세포 유전체 분석)에 대해 확장 가능한가?
• RQ5 저랭크 GW 방법은 랭크 선택과 정규화 파rameter에 대해 강건한가? 또한 원래 GW 공식화에 충실한가?

주요 결과

• 제안된 저랭크 GW 방법은 비용 행렬과 결합 양쪽의 저랭크 구조를 활용하여 $O(n)$ 시간 복잡도를 달성함으로써 선형 시간 계산이 가능하다.
• 합성 데이터($n=5000$)에서, 이 방법은 최첨단 엔트로픽 방법과 유사한 GW 손실을 달성하지만, $r = n/100$일 때조차도 수개의 주기 빠른 속도를 기록한다.
• 대규모 BRAIN 데이터셋($n=34,079$, $m=27,906$)에서는 엔트로픽 솔버가 메모리 및 시간 제약으로 실패하는 상황에서도 GW 정렬 계산을 성공적으로 수행한다.
• BRAIN 데이터셋에서 MREC보다 정렬 정확도가 뛰어나며, 다양한 랭크 설정에서 더 높은 수준의 정확히 매칭된 레이블이 부여된 세포 비율을 확보한다.
• 이 방법은 랭크 선택에 대해 강건하며, 엔트로픽 방법과 달리 정규화 파rameter $\varepsilon$의 미세조정이 필요하지 않다.
• 원래 GW 공식화에 높은 충실도를 유지하면서도 상당한 속도 향상을 이끌어내어, 이전에는 처리 불가능했던 문제에 GW를 적용할 수 있도록 한다.