[论文解读] Linear-Time Gromov Wasserstein Distances using Low Rank Couplings and Costs

Meyer Scetbon, Gabriel Peyré+1 more|arXiv (Cornell University)|Jun 2, 2021

Topological and Geometric Data Analysis参考文献 40被引用 8

一句话总结

本文通过利用低秩耦合与低秩代价矩阵，提出了一种线性时间的 Gromov-Wasserstein (GW) 距离计算方法，实现了在异质空间中对点云的高效、可扩展对齐。通过将耦合的低秩分解与输入代价矩阵的低秩近似相结合，该方法实现了 $O(n)$ 的时间复杂度——相比最先进的熵正则化 GW 方法快了数个数量级，同时在合成数据和真实世界数据集（如单细胞基因组学数据）上保持了具有竞争力的 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) 距离的高计算成本问题，其为 NP-难问题，通常通过昂贵的 $O(n^3)$ 熵正则化方法求解。
• 通过在代价矩阵和耦合矩阵中利用低秩结构，克服 GW 求解器中迭代 $C_t = A P_{t-1} B$ 计算的瓶颈。
• 开发一种可扩展的 GW 近似方法，在保持原始 GW 公式保真度的同时实现线性时间复杂度。
• 使 GW 能够应用于大规模数据集（如单细胞基因组学）——这些数据集超出了标准熵正则化求解器的处理能力。
• 证明对耦合施加低秩约束，可在实现与熵正则化相当的 GW 性能的同时，实现快数个数量级的加速。

提出的方法

• 对耦合矩阵 $P$ 应用低秩分解，将其限制为 $P = U V^T$ 形式，其中 $U, V \in \mathbb{R}^{n \times r}$，从而减少自由度并实现更快的优化。
• 使用低秩分解近似输入代价矩阵 $A$ 和 $B$（例如，平方欧氏距离），假设其内在维度 $d, d' \ll n, m$。
• 采用一种适用于低秩耦合的 Sinkhorn 算法变体，通过在低秩因子上交替最小化来求解正则化 GW 问题。
• 通过低秩近似迭代更新合成代价矩阵 $C_t = A P_{t-1} B$，避免每次迭代中出现 $O(n^3)$ 的运算量。
• 引入一种秩感知优化方案，平衡 GW 损失与耦合秩 $r$，并自适应选择 $r$ 以保持精度。
• 利用低秩代价矩阵与耦合矩阵可使算法在 $r \ll n$ 时达到线性时间，从而实现 $O(n)$ 复杂度而非 $O(n^3)$。

实验结果

研究问题

• RQ1 对耦合矩阵进行低秩分解，能否将 Gromov-Wasserstein 距离计算的复杂度从 $O(n^3)$ 降低至 $O(n)$？
• RQ2 在输入代价矩阵和耦合矩阵中同时利用低秩结构，是否能在实现线性时间计算的同时保持 GW 的精度？
• RQ3 在真实世界数据集上，低秩 GW 的性能与最先进的熵正则化 GW 相比，在 GW 损失和对齐质量方面如何？
• RQ4 所提出的方法能否扩展到大规模数据集（如单细胞基因组学）——在这些数据集中，熵正则化 GW 因内存和时间限制而不可行？
• RQ5 低秩 GW 方法对秩选择和正则化参数是否具有鲁棒性？是否能保持对原始 GW 公式的高保真度？

主要发现

• 所提出的低秩 GW 方法通过在代价矩阵和耦合矩阵中利用低秩结构，实现了 $O(n)$ 时间复杂度，从而实现线性时间计算。
• 在 $n=5000$ 的合成数据上，该方法达到的 GW 损失与最先进的熵正则化方法相当，但速度高出数个数量级，即使在 $r = n/100$ 时亦然。
• 在大规模 BRAIN 数据集（$n=34,079$, $m=27,906$）上，该方法成功计算了 GW 对齐，而熵正则化求解器因内存和时间限制而失败。
• 在 BRAIN 数据集上，该方法在对齐精度上优于 MREC，在多种秩设置下均实现了更高比例的正确匹配标签细胞。
• 该方法对秩的选择具有鲁棒性，且无需像熵正则化方法那样对正则化参数 $\varepsilon$ 进行精细调参。
• 该方法在实现显著加速的同时，保持了对原始 GW 公式的高保真度，使 GW 能够应用于此前难以处理的问题。