[论文解读] Why do Nearest Neighbor Language Models Work?

Frank F. Xu, Uri Alon+1 more|arXiv (Cornell University)|Jan 7, 2023

Topic Modeling被引用 6

一句话总结

本文研究了为何k-最近邻语言模型（kNN-LMs）在从与基础语言模型相同的训练数据中检索时仍能提升困惑度。研究识别出三个关键因素：来自不同Transformer层的表征集成（贡献了55%的性能提升）、使用近似最近邻搜索（带来2.6%的提升）以及温度超参数的调优（最高可带来3.7%的提升），这三者共同解释了kNN-LM的成功，而并非源于记忆化或过拟合。

ABSTRACT

Language models (LMs) compute the probability of a text by sequentially computing a representation of an already-seen context and using this representation to predict the next word. Currently, most LMs calculate these representations through a neural network consuming the immediate previous context. However recently, retrieval-augmented LMs have shown to improve over standard neural LMs, by accessing information retrieved from a large datastore, in addition to their standard, parametric, next-word prediction. In this paper, we set out to understand why retrieval-augmented language models, and specifically why k-nearest neighbor language models (kNN-LMs) perform better than standard parametric LMs, even when the k-nearest neighbor component retrieves examples from the same training set that the LM was originally trained on. To this end, we perform a careful analysis of the various dimensions over which kNN-LM diverges from standard LMs, and investigate these dimensions one by one. Empirically, we identify three main reasons why kNN-LM performs better than standard LMs: using a different input representation for predicting the next tokens, approximate kNN search, and the importance of softmax temperature for the kNN distribution. Further, we incorporate these insights into the model architecture or the training procedure of the standard parametric LM, improving its results without the need for an explicit retrieval component. The code is available at https://github.com/frankxu2004/knnlm-why.

研究动机与目标

• 理解为何kNN-LMs在从与基础LM相同的训练数据中检索时仍能提升性能。
• 分离并评估kNN-LMs中影响性能提升的具体设计选择。
• 确定性能提升是否源于检索、架构集成，或其他因素（如正则化或标签平滑）。
• 提出kNN-LM的通用形式化表达，以实现对设计组件的系统性消融分析。
• 证明性能提升可在无需显式检索的标准参数化LM中复现，从而提升效率。

提出的方法

• 提出kNN-LM的通用形式化表达，将非参数组件与参数化LM的softmax和嵌入层统一。
• 通过系统性地修改组件进行消融研究：用于相似度计算的输入表征、使用近似kNN搜索与精确kNN搜索的对比，以及kNN组件中的softmax温度。
• 训练并比较使用不同输入表征（如注意力层输出与前馈层输出）的基础LM变体，以隔离表征集成的影响。
• 用可微矩阵运算替代kNN检索，以在保持性能的同时提升推理效率。
• 测试替代的训练目标和标签类型（如软标签），以排除kNN-LM成功可能的其他解释。

实验结果

研究问题

• RQ1 为何kNN-LM在从与训练基础LM相同的训练数据中检索时仍能降低困惑度？
• RQ2 性能提升是源于检索、架构集成，还是近似搜索带来的正则化效应？
• RQ3 kNN组件中的softmax温度是否显著影响模型性能，且在原始kNN-LM设置中是否为最优？
• RQ4 能否通过修改输入表征或训练过程，在无需检索的标准参数化LM中复现kNN-LM的性能提升？
• RQ5 近似最近邻搜索在泛化中扮演什么角色，为何其表现可能优于精确搜索？

主要发现

• 将Transformer不同层（如注意力与前馈输出）的表征进行集成，可解释kNN-LM中55%的困惑度降低，相当于困惑度相对减少6.5%。
• 使用近似最近邻搜索替代精确搜索，带来2.6%的相对困惑度降低，表明精度降低可能具有正则化效应。
• 调优kNN组件中的softmax温度可实现最高达3.7%的相对困惑度降低，表明温度是一个关键但常被忽视的超参数。
• kNN-LM的成功与使用不同基础模型进行集成无关，表明性能提升并非源于模型平均。
• 通过受控的消融实验，已排除记忆化、过拟合、稀疏化、概率窃取或软标签等因素的影响，性能提升并非源于这些因素。
• 可微矩阵运算可替代kNN搜索，同时保持超过6.5%的相对困惑度降低，实现无需检索的高效推理。