[论文解读] Faster Nearest Neighbor Machine Translation

Shuhe Wang, Jiwei Li+6 more|arXiv (Cornell University)|Dec 15, 2021

Natural Language Processing Techniques被引用 9

一句话总结

本文提出 Faster $k$ NN-MT，一种基于层次聚类的方法，通过使用聚类中心近似距离，加速神经机器翻译中的最近邻搜索，将推理时间减少至仅标准模型的1.2倍，同时保持最先进性能。该方法通过仅在聚类中心和预计算的类内距离中进行搜索，显著加快检索速度，使 $k$ NN-MT 在实际系统中得以实用化部署。

ABSTRACT

$k$NN based neural machine translation ($k$NN-MT) has achieved state-of-the-art results in a variety of MT tasks. One significant shortcoming of $k$NN-MT lies in its inefficiency in identifying the $k$ nearest neighbors of the query representation from the entire datastore, which is prohibitively time-intensive when the datastore size is large. In this work, we propose extbf{Faster $k$NN-MT} to address this issue. The core idea of Faster $k$NN-MT is to use a hierarchical clustering strategy to approximate the distance between the query and a data point in the datastore, which is decomposed into two parts: the distance between the query and the center of the cluster that the data point belongs to, and the distance between the data point and the cluster center. We propose practical ways to compute these two parts in a significantly faster manner. Through extensive experiments on different MT benchmarks, we show that extbf{Faster $k$NN-MT} is faster than Fast $k$NN-MT \citep{meng2021fast} and only slightly (1.2 times) slower than its vanilla counterpart while preserving model performance as $k$NN-MT. Faster $k$NN-MT enables the deployment of $k$NN-MT models on real-world MT services.

研究动机与目标

• 解决 $k$ NN-MT 在推理过程中从大规模数据存储中检索最近邻时效率低下的问题。
• 克服 Fast $k$ NN-MT 的局限性，后者在高频词具有大规模特定标记数据存储时仍显缓慢。
• 在保持高翻译性能的同时，降低 $k$ NN-MT 的时间和内存开销。
• 通过提升推理速度和可扩展性，实现 $k$ NN-MT 在真实世界机器翻译服务中的实用化部署。

提出的方法

• 应用层次聚类将数据存储划分为多个聚类，并以聚类中心代表每个聚类。
• 将查询与数据点之间的距离近似为两部分之和：(1) 查询到聚类中心的距离，(2) 数据点到其所属聚类中心的距离。
• 预计算并缓存所有数据点与其对应聚类中心之间的类内距离，以加速检索。
• 在推理过程中仅将搜索范围限制在最近的聚类中心，从而减少距离计算次数。
• 在源端和目标端分别构建更小的、基于聚类的数据存储，以降低内存使用并提高批量推理效率。
• 使用温度缩放的 Softmax 聚合检索到的邻居的概率，以保持原始 $k$ NN-MT 的概率校准特性。

实验结果

研究问题

• RQ1 能否通过层次聚类高效近似 $k$ 近邻搜索，同时在 $k$ NN-MT 中实现性能损失可忽略？
• RQ2 所提出的基于聚类的方法在推理速度和内存使用方面与 Fast $k$ NN-MT 相比如何？
• RQ3 在保持翻译质量的前提下，通过聚类中心能将搜索空间减少到何种程度？
• RQ4 该方法能否扩展至具有长输入序列和高频标记的大规模真实世界机器翻译系统？

主要发现

• Faster $k$ NN-MT 的推理速度仅比标准神经机器翻译模型慢1.2倍，而 Fast $k$ NN-MT 则慢2倍。
• 该方法显著降低了时间和空间复杂度，支持更大的批量大小，并在推理过程中实现进一步加速。
• 在 WMT 和领域适应基准测试（医学、IT、古兰经、字幕）中，Faster $k$ NN-MT 的平均性能与 Fast $k$ NN-MT 和原始 $k$ NN-MT 相当或仅低不到1 BLEU 点。
• 该模型在多种不同领域中均保持强劲性能，SacreBLEU 分数在 WMT'14 上为52.7，WMT'16 上为44.9，古兰经数据集上为37.1，性能下降可忽略。
• 层次聚类策略显著减少了数据存储大小，使 $k$ NN-MT 在生产环境中部署成为可能。
• 该方法通过聚类限制高频词（如 'the'、'is'）的每个标记数据存储大小，有效避免了检索瓶颈。