[论文解读] Scalable Nearest Neighbor Search for Optimal Transport

Artūrs Bačkurs, Yihe Dong+3 more|arXiv (Cornell University)|Oct 9, 2019

Advanced Image and Video Retrieval Techniques参考文献 34被引用 10

一句话总结

本文提出 Flowtree，一种新颖的算法，用于在 Wasserstein-1（W1）距离下实现可扩展的近似最近邻搜索，通过树形嵌入与基于流的估计方法加速相似性搜索。在真实世界数据集上，其运行时间相比最先进方法最高可提升 7.4 倍，同时保持高精度。

ABSTRACT

The Optimal Transport (a.k.a. Wasserstein) distance is an increasingly popular similarity measure for rich data domains, such as images or text documents. This raises the necessity for fast nearest neighbor search algorithms according to this distance, which poses a substantial computational bottleneck on massive datasets. In this work we introduce Flowtree, a fast and accurate approximation algorithm for the Wasserstein-$1$ distance. We formally analyze its approximation factor and running time. We perform extensive experimental evaluation of nearest neighbor search algorithms in the $W_1$ distance on real-world dataset. Our results show that compared to previous state of the art, Flowtree achieves up to $7.4$ times faster running time.

研究动机与目标

• 解决大规模数据集中基于 Wasserstein-1 距离的精确最近邻搜索所面临的计算瓶颈问题。
• 设计一种快速、准确且可扩展的近似算法，适用于 W1 距离计算，适合 k-NN 分类流水线。
• 克服 W1 与标准欧几里得几何不兼容的问题，以及高维空间中精确流计算效率低下的问题。
• 在真实应用场景（如文本和图像相似性）中，评估并比较 Flowtree 与现有近似方法的性能。
• 优化整合多个估计器（如 Quadtree、R-WMD、Sinkhorn）与 Flowtree 的端到端流水线，以实现更优的速度-精度权衡。

提出的方法

• Flowtree 构建底层度量空间（R^d）的分层树分解，以实现高效的流估计。
• 通过自底向上的树遍历计算近似最优传输计划，降低 W1 计算的复杂度。
• 通过在树嵌入度量上求解流问题来估计 W1 距离，可在线性时间内完成计算。
• Flowtree 可集成到多阶段流水线中，既可用作中间过滤器，也可作为现有估计器（如 R-WMD 或 Sinkhorn）的替代方案。
• 对方法进行形式化分析，给出近似因子和运行时间的理论保证，确保精度与效率。
• 利用验证集对流水线进行调优，以优化各阶段的候选剪枝阈值（c_i），在保证目标召回率的前提下最小化运行时间。

实验结果

研究问题

• RQ1 基于树的流近似方法是否能在 W1 最近邻搜索中同时实现高精度与显著加速？
• RQ2 与 R-WMD、Mean 和 Sinkhorn 等现有估计器相比，Flowtree 在运行时间与召回率表现上如何？
• RQ3 在多阶段 W1 搜索中，将 Flowtree 作为过滤器或替代方案时，最优流水线配置为何？
• RQ4 Flowtree 在 k-NN 流水线中在多大程度上缓解了精确 W1 计算的计算瓶颈？
• RQ5 Flowtree 的树嵌入方法是否能保持与高维数据几何结构（如词嵌入）的兼容性？

主要发现

• 在 20news 数据集上，当召回率@5 ≥ 0.9 时，Flowtree 相比最佳基线流水线将中位运行时间减少了高达 7.4 倍。
• 当召回率@1 ≥ 0.9 时，Flowtree 相比最佳基线流水线将运行时间提升了 3.7 倍。
• 当用作流水线第二阶段的替代方案时，Flowtree 在召回率@5 实验中将运行时间减少了 7.4 倍至 8 倍。
• Flowtree 集成流水线中的主要性能瓶颈转移至精确 W1 计算，表明 Flowtree 有效分担了大部分计算成本。
• 使用 Flowtree 作为中间过滤器的流水线，根据下游估计器的不同，可实现 1.64 倍至 6.75 倍的速度提升。
• 尽管 Sinkhorn-3 比 Sinkhorn-1 更慢，但其更高的精度使得从 Quadtree 传递的候选数量更少，因此与 Flowtree 结合后整体运行时间更短。