[论文解读] Scalable Optimal Transport in High Dimensions for Graph Distances, Embedding Alignment, and More

Johannes Klicpera, Marten Lienen+1 more|arXiv (Cornell University)|Jul 14, 2021

Caching and Content Delivery被引用 3

一句话总结

本文提出两种对数线性时间近似方法——基于局部敏感哈希（LSH）的稀疏代价矩阵与局部校正Nystrom（LCN）——以加速高维空间中的熵正则化最优传输（OT）。这些方法使深度学习应用中的可扩展OT成为可能，在词嵌入对齐任务中实现3倍加速并提升3.1%准确率，在图距离回归任务中通过图传输网络（GTN）实现48%的性能提升。

ABSTRACT

The current best practice for computing optimal transport (OT) is via entropy regularization and Sinkhorn iterations. This algorithm runs in quadratic time as it requires the full pairwise cost matrix, which is prohibitively expensive for large sets of objects. In this work we propose two effective log-linear time approximations of the cost matrix: First, a sparse approximation based on locality-sensitive hashing (LSH) and, second, a Nystrom approximation with LSH-based sparse corrections, which we call locally corrected Nystrom (LCN). These approximations enable general log-linear time algorithms for entropy-regularized OT that perform well even for the complex, high-dimensional spaces common in deep learning. We analyse these approximations theoretically and evaluate them experimentally both directly and end-to-end as a component for real-world applications. Using our approximations for unsupervised word embedding alignment enables us to speed up a state-of-the-art method by a factor of 3 while also improving the accuracy by 3.1 percentage points without any additional model changes. For graph distance regression we propose the graph transport network (GTN), which combines graph neural networks (GNNs) with enhanced Sinkhorn. GTN outcompetes previous models by 48% and still scales log-linearly in the number of nodes.

研究动机与目标

• 解决高维深度学习应用中最优传输（OT）的计算瓶颈问题，其中完整代价矩阵的计算成本过高。
• 开发高效、可扩展的标准Sinkhorn算法替代方案，依赖熵正则化但避免二次方代价矩阵计算。
• 实现在大规模场景（如词嵌入对齐与图距离回归）中OT的实际应用。
• 提出一种新颖的图传输网络（GTN），结合图神经网络（GNN）与增强版Sinkhorn算法，提升性能的同时保持对数线性可扩展性。
• 对近似质量与真实任务中端到端性能提升进行理论与实证验证。

提出的方法

• 提出基于局部敏感哈希（LSH）的稀疏代价矩阵近似方法，以减少OT中成对距离计算的数量。
• 引入局部校正Nystrom（LCN）方法，利用LSH识别低秩Nystrom近似中代价矩阵的稀疏校正项。
• 将基于LSH的近似方法集成到熵正则化OT的Sinkhorn迭代中，保持收敛性保证的同时将时间复杂度降低至对数线性。
• 设计图传输网络（GTN），利用GNN生成节点特征，并应用结合LCN近似的增强版Sinkhorn算法进行图级别距离回归。
• 理论分析表明，在底层数据分布满足弱假设条件下，LCN可保持收敛性与准确性。
• 在下游任务（如词嵌入对齐与图距离预测）中对所提方法进行端到端训练与评估。

实验结果

研究问题

• RQ1 基于LSH的代价矩阵稀疏化是否能在保持熵正则化OT精度的同时实现对数线性时间复杂度？
• RQ2 局部校正Nystrom（LCN）方法是否在近似质量与计算效率之间提供了优于标准Nystrom或单独使用LSH的更好权衡？
• RQ3 所提近似方法在真实世界深度学习应用（如无监督词嵌入对齐）中的表现如何？
• RQ4 将LCN近似OT集成到图神经网络框架（GTN）中，是否能在图距离回归任务中实现最先进性能？
• RQ5 在现有流水线中用所提近似方法替代标准Sinkhorn算法，实际性能在速度与准确率上的提升如何？

主要发现

• 基于LSH的稀疏近似与LCN方法将熵正则化OT的时间复杂度降低至对数线性，实现了大规模数据集的可扩展性。
• 在无监督词嵌入对齐任务中，所提方法使最先进方法提速3倍，同时准确率提升3.1个百分点。
• 图传输网络（GTN）结合GNN与LCN近似Sinkhorn算法，在图距离回归任务中相比先前模型性能提升48%。
• LCN近似保持了高近似质量，与完整矩阵Sinkhorn相比，OT解的精度损失可忽略不计。
• 所提方法在图任务中与节点数量呈对数线性扩展，适用于大规模图结构。
• 实证评估证实，两种近似方法在多样化高维场景（包括深度学习表征）中均具有效性。