[论文解读] Improving language models by retrieving from trillions of tokens

Sebastian Borgeaud, Arthur Mensch+26 more|arXiv (Cornell University)|Dec 8, 2021

Topic Modeling被引用 296

一句话总结

Retro 在自回归语言模型中加入从万亿令牌数据库的检索机制，达到 GPT-3/Jurassic-1 水平的性能，但参数更少，并使知识密集型下游任务成为可能。

ABSTRACT

We enhance auto-regressive language models by conditioning on document chunks retrieved from a large corpus, based on local similarity with preceding tokens. With a $2$ trillion token database, our Retrieval-Enhanced Transformer (RETRO) obtains comparable performance to GPT-3 and Jurassic-1 on the Pile, despite using 25$ imes$ fewer parameters. After fine-tuning, RETRO performance translates to downstream knowledge-intensive tasks such as question answering. RETRO combines a frozen Bert retriever, a differentiable encoder and a chunked cross-attention mechanism to predict tokens based on an order of magnitude more data than what is typically consumed during training. We typically train RETRO from scratch, yet can also rapidly RETROfit pre-trained transformers with retrieval and still achieve good performance. Our work opens up new avenues for improving language models through explicit memory at unprecedented scale.

研究动机与目标

• 通过利用通过检索获得的巨量外部记忆来减少对日益增大的模型的依赖以提升性能的动机。
• 开发可扩展到万亿令牌的检索增强自回归架构。
• 证明检索在不同模型规模和下游任务中均能提供一致的提升。
• 展示 Retro 可以以最少额外计算量对预训练模型进行 Retro-fitted。
• 在检索系统能够访问训练数据时解决评估泄漏问题。

提出的方法

• 使用冻结的 BERT 嵌入作为键、原始文本块作为值，构建文本块的键值数据库。
• 将输入序列分成块，并用数据库中最近邻检索到的块来增强每个块。
• 使用带有分块跨注意力机制(Cca)的编码器-解码器 Transformer 来整合检索数据。
• 使用冻结的基于 BERT 的检索器实现大规模检索，无需对检索器进行训练。
• 使用以先前块的检索邻居为条件的自回归目标进行训练，从而保留因果性。
• 在大规模多语言数据（MassiveText）和测试数据集（C4、Wikitext103、Pile）上进行评估，并通过评估块重叠分析泄漏效应。

实验结果

研究问题

• RQ1 从万亿令牌数据库检索如何影响不同模型规模下的语言模型性能？
• RQ2 带检索增强的模型是否能够在标准基准上以更少的参数达到或超过 GPT-3 与 Jurassic-1？
• RQ3 增加检索数据库的大小或检索邻居的数量是否带来稳定的提升，极限在哪里？
• RQ4 Retro 是否能够以最小额外计算量有效地对已有的预训练变换器进行 Retro-fitted？
• RQ5 评估数据泄漏如何影响检索增强模型的性能？

主要发现

• 使用 2 万亿令牌数据库的 Retro 在 Pile 上的表现与 GPT-3 和 Jurassic-1 相当，同时参数少 25 倍。
• 从 150M 到 7B 参数的模型规模上，检索带来的性能提升保持不变。
• 增加检索数据库规模和检索邻居数量都会提升性能，直至出现非常大的邻居集可能降低质量的点。
• 经过训练后，Retro 在 Wikitext103 与 Pile 上取得了最先进的结果，并且可以微调以实现有竞争力的下游任务，如问答。
• 评估泄漏意识方法表明，改进来自于显式邻居拷贝和对广泛知识的提取；更大的数据库和更多的邻居会放大对更大模型的收益。