[论文解读] Efficient Nearest Neighbor Language Models

Junxian He, Graham Neubig+1 more|arXiv (Cornell University)|Sep 9, 2021

Topic Modeling参考文献 33被引用 5

一句话总结

本文提出了一套技术组合，通过实现自适应检索、数据存储裁剪和基于PCA的降维，加速k近邻语言模型（kNN-LM）。通过结合这些方法，该方法在WikiText-103上实现了最高6.6倍的加速，在Law-MT上实现了5.4倍的加速，同时保持或略微提升与原始kNN-LM相当的困惑度，使非参数化语言建模在实际部署中更具可行性。

ABSTRACT

Non-parametric neural language models (NLMs) learn predictive distributions of text utilizing an external datastore, which allows them to learn through explicitly memorizing the training datapoints. While effective, these models often require retrieval from a large datastore at test time, significantly increasing the inference overhead and thus limiting the deployment of non-parametric NLMs in practical applications. In this paper, we take the recently proposed $k$-nearest neighbors language model (Khandelwal et al., 2020) as an example, exploring methods to improve its efficiency along various dimensions. Experiments on the standard WikiText-103 benchmark and domain-adaptation datasets show that our methods are able to achieve up to a 6x speed-up in inference speed while retaining comparable performance. The empirical analysis we present may provide guidelines for future research seeking to develop or deploy more efficient non-parametric NLMs.

研究动机与目标

• 解决由于从大型外部数据存储中进行昂贵检索而导致的非参数化语言模型推理成本过高的问题。
• 在不牺牲性能的前提下提升kNN-LM的效率，使其更适用于实际NLP应用中的广泛部署。
• 探究是否可以有选择性地跳过某些标记的检索、移除冗余的数据存储条目，或降低向量维度以提升速度。
• 为非参数化语言建模中速度与性能的平衡提供实证指导。

提出的方法

• 引入一种基于轻量级神经网络的自适应检索机制，用于预测每个标记是否需要检索，从而将检索调用减少高达50%。
• 采用贪心合并与基于重要性的过滤方法，对数据存储条目进行裁剪，最多可减少40%的条目，同时保持模型性能。
• 使用基于PCA的降维方法压缩数据存储中的向量表示，从而同时提升速度与困惑度。
• 将自适应检索、数据存储裁剪与降维方法整合为统一的优化流水线，以实现最大效率。
• 在预留的验证集上训练检索适配器，以避免过拟合并确保泛化能力。
• 利用FAISS的PCA实现配合随机旋转，以提升向量搜索中的量化效率与性能。

实验结果

研究问题

• RQ1 我们能否通过智能地跳过某些标记的检索，在不降低性能的前提下减少推理开销？
• RQ2 在不影响模型准确率的前提下，我们能在多大程度上通过移除冗余或不重要的条目来压缩数据存储？
• RQ3 降低数据存储向量的维度是否能同时提升速度与性能？如果是，原因是什么？
• RQ4 哪种效率技术的最优组合能够在保持困惑度的同时最大化加速效果？

主要发现

• 自适应检索在WikiText-103上实现了最高1.9倍的加速，困惑度仅增加0.02，检索调用减少50%。
• 贪心合并裁剪方法使数据存储减少了40%，困惑度仅增加0.2，证明了冗余条目有效去除。
• 将向量维度通过PCA降至512维后，困惑度降低0.25点（从16.65降至16.40），并在WikiText-103上实现3.6倍加速。
• 三种技术联合使用在WikiText-103上实现了6.6倍加速，困惑度与基线kNN-LM几乎相同（16.67 vs. 16.65）。
• 在Law-MT领域自适应数据集上，联合方法实现5.4倍加速，并较基线kNN-LM的困惑度降低0.35点。
• PCA不仅加速了推理，还在Law-MT上提升了性能（从12.64降至11.59），表明其构建的向量空间更适用于基于L2的kNN搜索。