摘要

近年来，许多研究通过将检索到的数据添加至输入上下文，来增强语言模型的能力。然而，要使这一方法取得成功，检索数据必须在训练和推理阶段均被引入。此外，随着检索数据量的增加，输入长度呈线性增长，而现代Transformer模型在计算和内存消耗上则面临二次方增长的挑战。为避免这些复杂性，我们提出一种简化方案：在测试阶段，仅使用模型标准的训练设置，对检索到的数据进行微调。我们基于Pile数据集的文本嵌入构建了一个大规模分布式索引。对于每个测试输入，系统检索其最近邻，并在这些邻近文本上对模型进行微调。令人惊讶的是，仅需检索并针对20个最近邻进行一次梯度更新，即可在Pile数据集中的20多个语言建模任务上显著提升模型性能。例如，通过最近邻的测试时训练，能够将小型GPT-2模型与规模超过其十倍的GPT-Neo模型之间的性能差距显著缩小。然而，索引的质量与规模必须足够充分。本工作为语言建模中的测试时训练（test-time training）建立了一个首个基准。