面向大语言模型的基于最近邻的测试时训练

许多近期工作通过向输入上下文添加检索数据来增强语言模型。为使该方法成功,检索数据必须在训练与测试时均被加入。此外,由于输入长度随检索数据大小线性增长,对现代 Transformer 而言计算与内存成本呈二次增长。为避免这些复杂问题,我们仅在测试时使用其标准训练设置对检索数据微调模型。我们基于 Pile 数据集的文本嵌入构建了大规模分布式索引。对每个测试输入,我们的系统检索其邻居并以其文本微调模型。令人惊讶的是,检索并训练少至 20 个邻居,每个仅进行一次梯度迭代,便大幅改善了 Pile 中超过 20 个语言建模任务的性能。例如,采用最近邻的测试时训练将小型 GPT-2 与规模超其 10 倍以上的 GPT-Neo 模型间的性能差距显著缩小。然而,足够的索引质量与规模是必要的。我们的工作建立了语言建模测试时训练的首个基线。