一种用于 Sinkhorn 词移距离的新并行算法及其在 PIUMA 与至强 CPU 上的性能

词移距离（WMD）通过计算如下代价来衡量两篇文本文档之间的语义不相似度：将源/查询文档的所有词最优地移动到目标文档中最相似的词所需的代价。由于计算两篇文档之间的 WMD 需要求解一个 $O(V^3log(V))$ 的优化问题（其中 $V$ 为文档中不同词的数量），因此代价高昂。幸运的是，WMD 可被视为地球移动距离（EMD）的一种形式，对此可通过在优化问题中加入熵惩罚并使用 Sinkhorn-Knopp 算法求解，将算法复杂度降至 $O(V^2)$。此外，采用批处理方式（即一次性计算单个查询文档相对于多个目标文档的 WMD）可使计算高度并行化。Sinkhorn WMD 是许多 ML/NLP 应用中使用的关键核函数，通常用以 Python 实现。然而，即便内部调用了优化的 C++ BLAS 例程，朴素的 Python 实现仍可能留下显著的性能余量。我们提出了一种用于 Sinkhorn-Knopp 词移距离的新稀疏并行算法（{P}arallel {A}lgorithm for {S}inkhorn-Knopp {W}ord-movers {D}istance），通过采用 $O(V^2)$ 的 EMD 算法来计算一个文档到多个其他文档的语义距离。我们利用新的融合 SDDMM-SpMM（稠密-稠密矩阵乘的稀疏选择、稀疏-稠密矩阵乘）核函数，将 $O(V^2)$ 的稠密重计算型 EMD 版本在算法上转化为等价的稀疏版本。我们针对两种差异极大的架构——新的英特尔可编程集成统一内存架构（PIUMA）与英特尔至强（Xeon）CPU——实现并优化了该算法。我们展示了在两种平台上均能达到接近峰值性能的表现。