K 近邻语言模型

这是论文《通过记忆推广：最近邻语言模型》的 PyTorch 实现。它使用 k 最近邻来改善自回归变压器模型的困惑度。

自回归语言模型估计$p(w_t|c_t)$，步骤中的标记在哪里$w_t$$t$，$c_t$是上下文，$c_t=(w_1,w_2,...,w_{t-1})$。

本文改进了$p(w_t|c_t)$使用带搜索键的键值对$(f(c_i),w_i)$使用 k 最近邻搜索的功能$f(c_t)$。$f(c_t)$这是上下文的嵌入$c_t$。本文（以及本实现）使用变压器最后一层前馈层的输入作为$f(c_t)$。

我们使用 FAISS 进行索引$f(c_i)$。

实施

因此，要运行$k$ NN-LM，我们需要：

• 训练变压器模型
• 建立索引$(f(c_i),w_i)$
• 使用 NN 搜索来评估 k$k$ nn-ML$(f(c_i),w_i)$$f(c_t)$

这个实验使用了一个小数据集，这样我们就可以在不占用几百千兆字节的索引磁盘空间的情况下运行它。

$k$NN-LM 的官方实现可以在这里找到。