[论文解读] An Efficient Shared-memory Parallel Sinkhorn-Knopp Algorithm to Compute the Word Mover's Distance

Jesmin Jahan Tithi, Fabrizio Petrini|arXiv (Cornell University)|May 14, 2020

Topic Modeling参考文献 13被引用 4

一句话总结

本文提出了一种共享内存并行Sinkhorn-Knopp算法，通过将密集矩阵运算转换为稀疏核融合，加速了词移距离（WMD）的计算，在96个核心上实现了67倍的加速，并且比原始Python实现快逾700倍。该方法结合了SDDMM_SpMM核融合与优化的欧几里得距离计算，在现代CPU上实现了高性能语义文本相似度计算。

ABSTRACT

The Word Mover's Distance (WMD) is a metric that measures the semantic dissimilarity between two text documents by computing the cost of moving all words of a source/query document to the most similar words of a target document optimally. Computing WMD between two documents is costly because it requires solving an optimization problem that costs \(O(V^3log(V))\) where \(V\) is the number of unique words in the document. Fortunately, the WMD can be framed as the Earth Mover's Distance (EMD) (also known as the Optimal Transportation Distance) for which it has been shown that the algorithmic complexity can be reduced to \(O(V^2)\) by adding an entropy penalty to the optimization problem and a similar idea can be adapted to compute WMD efficiently. Additionally, the computation can be made highly parallel by computing WMD of a single query document against multiple target documents at once (e.g., finding whether a given tweet is similar to any other tweets happened in a day). In this paper, we present a shared-memory parallel Sinkhorn-Knopp Algorithm to compute the WMD of one document against many other documents by adopting the \(O(V^2)\) EMD algorithm. We used algorithmic transformations to change the original dense compute-heavy kernel to a sparse compute kernel and obtained \(67 imes\) speedup using \(96\) cores on the state-of-the-art of Intel extregistered{} 4-sockets Cascade Lake machine w.r.t. its sequential run. Our parallel algorithm is over \(700 imes\) faster than the naive parallel python code that internally uses optimized matrix library calls.

研究动机与目标

• 为解决计算词移距离（WMD）的高计算成本问题，其时间复杂度为O(V³log V)，其中V为唯一词的数量。
• 实现在一个查询文档与多个目标文档之间同时进行高效并行的WMD计算。
• 通过算法转换，将Sinkhorn-Knopp算法中密集矩阵运算的计算负担转化为稀疏核，以减轻计算压力。
• 在使用OpenMP和C++的共享内存多插槽x86系统上，实现高性能与强可扩展性。

提出的方法

• 本文采用带熵正则化的地球移动距离（EMD）公式，将WMD计算复杂度从O(V³log V)降低至O(V²)。
• 应用Sinkhorn-Knopp矩阵平衡算法，高效求解正则化EMD问题，实现快速且稳定的收敛。
• 提出一种新型稀疏核融合技术SDDMM_SpMM，将稀疏-密集-密集矩阵乘法（SDDMM）与密集-稀疏矩阵乘法（SpMM）融合为单一优化核。
• 将欧几里得距离计算重构为分块的、类似GEMM的矩阵乘法核，以提升数据局部性与缓存效率。
• 使用OpenMP在多个CPU核心和NUMA域之间实现任务级并行，并对内存访问模式进行调优以提升可扩展性。
• 通过算法转换，将密集计算核替换为稀疏核，降低内存带宽压力并提高算术强度。

实验结果

研究问题

• RQ1 能否有效将WMD计算中Sinkhorn-Knopp算法的密集计算核转化为稀疏核，以提升性能？
• RQ2 在共享内存系统上，通过并行化多个目标文档的WMD计算，可实现多大程度的加速？
• RQ3 高度优化的C++/OpenMP实现与使用优化数学库的原始并行Python实现相比，性能如何？
• RQ4 对欧几里得距离核进行算法优化，对整体WMD计算性能的影响有多大？
• RQ5 在4插槽x86服务器架构中，该算法在多个NUMA节点间的可扩展性如何？

主要发现

• 所提出的并行Sinkhorn-WMD算法在Intel Cascade Lake系统上，4个NUMA插槽、96个核心的环境下，相比串行实现实现了67倍的加速。
• 该实现比使用优化并行数学库的原始并行Python实现快逾700倍。
• SDDMM_SpMM核融合降低了内存带宽压力，提升了算术强度，从而高效利用了现代CPU缓存。
• 优化后的类似GEMM的欧几里得距离核通过增强数据局部性，显著提升了性能，尤其对WMD工作负载中常见的高瘦矩阵效果明显。
• 该算法在NUMA域之间表现出良好的可扩展性，在4插槽系统上，源文档中v_r = 37个词时达到峰值性能。
• 性能跟踪中观察到冷启动效应，v_r = 31时速度提升最差，归因于缓存未命中，表明输入排序与数据布局的重要性。