Atlas: 检索增强语言模型的few-shot学习

Atlas: Few-shot Learning with Retrieval Augmented Language Models

Aug 2022

Gautier Izacard, Patrick Lewis等

[Meta AI Research等]

https://arxiv.org/abs/2208.03299

https://github.com/facebookresearch/atlas

大型语言模型在广泛的任务中显示出令人印象深刻的few-shot结果。然而，当知识是这些结果的关键时，就像问答和事实核查等任务一样，似乎需要大量的参数来存储知识。众所周知，检索增强模型在不需要那么多参数的情况下就可以胜任知识密集型任务，但尚不清楚它们是否适用于few-shot设置。在这项工作中，我们介绍了Atlas，这是一个精心设计和预先训练的检索增强语言模型，能够在很少的训练示例中学习知识密集型任务。我们对各种任务进行了评估，包括MMLU、KILT和NaturalQuestions，并研究了文档索引内容的影响，表明它可以很容易地更新。值得注意的是，Atlas仅使用64个例子就在自然问题上达到了42%以上的准确率，尽管参数减少了50倍，但仍比540B参数模型高出3%。

开源的概述：

该存储库包含预训练的模型、语料库、索引和代码，用于论文Atlas:带检索增强语言模型的few-shot学习的预训练、微调、检索和评估

我们联合预训练了一个检索增强的seq2seq语言模型，该模型由基于段落的密集检索器和编码器-解码器语言模型组成。我们对各种任务进行了评估，包括MMLU、KILT和NaturalQuestions，并研究了文档索引内容的影响，表明它可以很容易地更新。值得注意的是，从2018年起，Atlas在提供维基百科索引时，仅使用64个例子，在自然问题上的准确率就达到了45%以上，尽管参数减少了50倍，但仍比540B参数模型高出6%。Atlas在更大的数据集上进行微调时也能很好地工作——在完整的自然问题数据上进行微调后，Atlas创下了64%的最新水平，比目前的技术水平高出8个百分点。

此存储库支持对大型和小型数据集进行预训练和微调。此存储库可以支持以下功能：

在解码器seq2seq模型中训练大型融合，测试高达11B的参数

使用各种不同的提取方法将解码器模型中的融合中的相关性信号提取到密集检索模型中。

在用户提供的文章语料库上执行端到端检索增强训练（测试多达400M篇文章，约40B个单词），在训练循环中进行检索

支持屏蔽语言建模、前缀语言建模、维基百科部分生成、开放域问答、多选问答、事实核查和KILT（也可以支持任意seq2seq任务）方面的培训

一种用于密集向量检索的快速、并行分布式GPU精确和近似最大内积搜索

支持快速就地刷新索引

各种内存优化和方法，用于在循环中训练检索器的同时保持快速准确的检索。

此外，有关其他功能，请参阅命令行参数或自述文件

大型语言模型（LLM）是令人印象深刻的few-shot学习者（Brown等人，2020年；Rae等人，2021；Hoffmann等人，2022年；Chowdhery等人，2022.）。他们能够在很少的例子中学习新任务，甚至只从指令中学习。为了实现这种泛化能力，关键因素是缩放模型的参数数量和训练数据的大小。大型语言模型的这种改进归功于更大的计算预算，实现了更复杂的推理，以及从更大的训练数据中记忆更多与下游任务相关的信息的能力。虽然直觉上认为，推理能力的提高会导致更好的泛化，因此很少有镜头学习，但参数内记忆却并非如此。具体而言，目前尚不清楚有效的few-shot学习在多大程度上需要大量的模型参数知识。

在本文中，我们研究了few-shot学习是否需要模型在其参数中存储大量信息，以及记忆是否可以与泛化脱钩。为此，我们利用了这样一个事实，即内存可以外包，并通过使用检索增强架构由外部非参数知识源取代。这些模型使用非参数记忆，例如，在大型、外部、潜在的非静态知识源上的神经检索器，以增强参数语言模型。除了记忆能力外，这种结构还具有吸引力，因为在适应性、可解释性和效率方面有许多其他既定优势（Guu等人，2020年；Lewis等人，2020；Yogatama等人，2021；Borgeaud等人，2021）。然而，检索增强模型尚未展示出令人信服的few-shot学习能力。在这项工作中，我们解决了这一差距，并提出了Atlas，这是一种检索增强语言模型，能够进行强大的few-shot学习，尽管它的参数计数比最近其他强大的多镜头学习者低。

Atlas通过使用基于Contriever的双编码器架构的通用密集检索器，基于当前上下文检索相关文档（Izacard et al.，2022）。检索到的文档以及当前上下文由序列到序列模型处理，该模型使用解码器中的融合架构（Izacard&Grave，2020）生成相应的输出。我们研究了训练Atlas的不同技术对其在一系列下游任务中的few-shot性能的影响，包括问答和事实核查。我们发现，对组件进行联合预训练对于few-shot性能至关重要，为此，我们仔细评估了许多现有的和新颖的预训练任务和方案。Atlas在few-shot和资源丰富的环境中都实现了强大的下游性能。例如，在只有11B参数的情况下，Atlas使用64个训练示例在NaturalQuestions上实现了42.4%的准确率（仅使用维基百科索引时为45.1%），比540B参数模型的PaLM（Chowdhery et al.，2022）高出近3个点，在使用维基百科索引的完整数据集设置中高出64.0%，以8个点的优势确立了新的技术水平。

总之，我们做出了以下贡献：

•深入研究如何设计和训练检索增强语言模型，重点关注下游的few-shot学习和样本效率。

•这项研究的发现导致了一个名为Atlas的检索增强语言模型，该模型很少表现出比标准LLM更低规模的few-shot能力。

•我们提供了微调策略的探索，以有效地使检索器和语言模型适应手头的任务。

•在few-shot设置中进行彻底的下游实验，在few-shotNaturalQuestions（+2.8%）、TriviaQA（+3.3%）、FEVER（+5.1%）上展示了最先进的结果，结果与MMLU上参数增加15倍的模型不相上下或更强。

•研究完整数据集微调的实验，在NaturalQuestions（+8.1%）、TriviaQA（+9.3%）和5个KILT任务中创造了最先进的新结果。

•实验证明了Atlas的可更新性和可解释性特征。

•实验表明，使用产品量化的压缩索引实现了与未压缩索引相当的性能，同时减少了5倍的内存。

我们的代码、经过预训练的Atlas检查点和各种支持数据开源

2方法

我们的方法遵循文本到文本的框架（Raffel et al.，2019）。这意味着所有任务的框架如下：系统获取一个文本查询作为输入，并生成一个文本输出。例如，在问答的情况下，查询对应于问题，模型需要生成答案。在分类任务的情况下，查询与文本输入相对应，并且模型生成词汇化的类标签，即与标签相对应的单词。我们从图2中的KILT基准中给出了更多下游任务的例子。由于许多自然语言处理任务都需要知识，我们的目标是通过检索来增强标准的文本到文本模型，正如我们在引言中假设的那样，这对于赋予模型few-shot的能力可能至关重要。

2.1架构

我们的模型基于两个子模型：检索器和语言模型。在执行任务时，从回答问题到生成维基百科文章，我们的模型首先用检索器从一个大型文本语料库中检索前k个相关文档。然后，这些文档与查询一起被提供给语言模型，查询反过来生成输出。检索器和语言模型都基于预先训练的Transformer网络，我们将在下面更详细地描述。

检索器。

我们的检索器模块基于Contriever（Izacard et al.，2022），这是一种基于连续密集嵌入的信息检索技术。Contriever使用双编码器架构，其中查询和文档由Transformer编码器独立嵌入（Huang et al.，2013；Karpukhin et al.，2020）。对最后一层的输出应用平均池化，以获得每个查询或文档的一个向量表示。然后，通过计算查询和每个文档对应嵌入之间的点积来获得查询和它们之间的相似性得分。Contriever模型是使用MoCo对比损失进行预训练的（He et al.，2020），并且仅使用无监督数据。如下一节所示，密集检索器的一个优点是，使用梯度下降和蒸馏等标准技术，可以在没有文档注释的情况下训练查询和文档编码器。

语言模型。

对于语言模型，我们依赖T5序列到序列架构（Raffel et al.，2019）。我们依靠Fusion in Decoder对序列到序列模型的修改，并在编码器中独立处理每个文档（Izacard&Grave，2020）。然后，我们将对应于不同文档的编码器的输出连接起来，并在解码器中对这一单个序列进行交叉关注。根据Izacard&Grave（2020），我们将查询连接到编码器中的每个文档。在语言模型中处理检索到的文档的另一种方法是将查询和所有文档连接起来，并使用这个长序列作为模型的输入。不幸的是，这种方法不会随着文档数量的增加而增加，因为编码器中的自注意导致了相对于文档数量的二次复杂度。

2.2检索器的训练目标

在本节中，我们将讨论四种不同的损失函数，以与语言模型一起训练检索器。我们考虑利用语言模型提供监督信号来训练检索器的损失函数。换句话说，如果语言模型在生成输出时发现文档有用，那么检索器的目标应该鼓励检索器将所述文档排名更高。这使我们能够仅使用感兴趣任务的查询和输出对来训练模型，而不依赖于文档注释。例如，在事实核查的情况下，模型只需要成对的事实和相应的核查，而不需要包含支持核查的证据的文件。在实践中，我们可以将这种方法应用于任何任务，包括自监督的预训练。如实验部分所示，预训练对于获得表现出很few-shot门学习能力的模型至关重要。

注意力蒸馏（ADist）。

我们考虑的第一个损失是基于语言模型的注意力得分，并受到Izacard&Grave（2021）的启发。其主要思想是，在生成输出时，输入文档和输出之间的交叉注意力得分可以用作每个输入文档重要性的代理。特别是，Izacard&Grave（2021）表明，这些分数可以在给定文档的注意头、层次和标记之间进行聚合，以获得每个文档的单个分数。然后，可以通过在前K个文档{dk}1，…上最小化具有概率分布pret的KL发散来将这些分数提取到检索器中，。。。，从检索器获得的K：

其中s是查询和文档向量之间的点积，θ是温度超参数。在最初的论文中，提出使用解码器交叉注意力的预softmax分数，并在头部、层和令牌之间进行平均。在这里，我们提出了一种替代方案，它给出了稍微更强的结果，这取决于以下观察结果。在注意力机制中，如

不能仅从注意力得分αn来评估特定令牌n对输出y的贡献，但也应该考虑值vn的范数。因此，我们使用数量αnkvnk2作为标记n的相关性度量。继Izacard&Grave（2021）之后，我们对所有注意力、层次和标记进行平均，以获得每个文档的得分。我们在所得分数上应用Softmax算子，以获得前K个检索到的文档上的分布模式（dk）。然后，我们最小化pattn（dk）和方程1中定义的来自检索器的分布pret之间的KL发散：

在这里，这种损失只用于优化检索器的参数，而不是语言模型。当使用最近的深度学习框架时，这是通过在pattn上应用StopGradient算子来实现的。

多文档读取器和检索器（EMDR2）的端到端训练。

接下来，我们考虑Sachan等人（2021）引入的方法，该方法受期望最大化算法的启发，将检索到的文档视为潜在变量。给定查询q、对应的输出a和当前检索器的前K个检索到的文档的集合DK，训练检索器的EMDR2损失为

其中pret再次是利用检索器获得的前K个文档的概率，如等式1所定义。同样，通过在plm周围应用StopGradient运算符，仅更新检索器的参数。应当注意，使该损失函数最大化的文档上的概率分布是对应于根据语言模型的输出的最高概率的文档的指示符。最后，在实践中，EMDR2损失函数应用于令牌级别，而不是序列级别。

困惑蒸馏（PDist）。

第三，我们讨论了一个更简单的损失函数，该函数受到了注意力蒸馏和EMDR2方法目标的启发（Izacard&Grave，2021；Sachan等人，2021）。更确切地说，我们希望训练检索器来预测在给定查询的情况下，每个文档将在多大程度上改善输出的语言模型困惑。为此，我们使用统一先验，最小化检索器的文档分布（等式1）和根据语言模型的文档后验分布之间的KL差异：

使用Softmax运算符，我们可以

省略一个困惑蒸馏（LOOP）。

最后，我们提出了一个目标，该目标基于在删除前k个检索到的文档时，语言模型的预测会变得多么糟糕。为此，我们计算k-1文档的每个子集的输出的对数概率，并使用负值作为每个文档的相关性得分。根据前面的损失函数，我们使用softmax算子来获得文档上的概率分布：

如前所述，我们将该分布与使用检索器获得的分布之间的KL偏差最小化。这种损失的计算成本比PDist和EMDR更高，但与ADist一样，它更接近于语言模型的训练方式，即LM被训练为以一组K个文档为条件。对于LOOP，语言模型以（K−1）个文档为条件，而不是像EMDR2和PDist那样以单个文档为条件。

对于所有损失，我们也可以在计算目标或检索器分布时使用温度超参数来控制分布的峰值，这对一些任务或损失可能很重要。事实上，对于PDist和LOOP，当条件是不同的文档时，输出的困惑可能不会有太大变化，尤其是在长输出的情况下。

2.3文本前任务

在本节中，我们描述了可以用于仅使用无监督数据联合预训练检索器和语言模型的借口任务。

前缀语言建模。

首先，我们将标准语言建模任务视为潜在的预训练目标。为了在文本到文本框架中进行语言建模，我们考虑一个由N个单词组成的块，并将该块拆分为两个子序列，长度相等N/2。然后，第一个子序列被用作查询，第二个子序列对应于输出。因此，我们通过使用N/2个令牌的第一个子序列来检索相关文档，以生成输出。

掩码语言建模。

其次，我们考虑了由Raffel等人（2019）制定的掩码语言建模。同样，从一组N个单词开始，我们对平均长度为3个标记的k个跨度进行采样，从而得到15%的掩码率。然后，我们用不同的特殊标记替换每个跨度。然后对模型进行训练以生成掩码跨度，每个跨度从插入输入序列中的特殊哨兵掩码令牌开始。我们使用掩码查询检索文档，但将特殊的掩码令牌替换为检索器词汇表支持的掩码令牌

章节生成的标题。

最后，我们考虑一个更抽象的生成任务，在给定文章和章节标题的情况下，从维基百科文章中生成章节。这里，查询与文章的标题以及小节的标题相对应，输出与小节的文本相对应。我们排除了“另请参阅”、“参考文献”、“进一步阅读”和“外部链接”部分。

2.4高效检索器微调

通过使用文档索引来促进检索，文档索引是检索语料库中所有文档的文档嵌入的预先计算的集合。当联合训练检索器和语言模型时，索引需要定期更新，否则，存储在索引中的文档的嵌入相对于更新的检索器会变得过时。这意味着我们需要在训练期间定期重新计算文档的完整集合的嵌入，以保持索引的新鲜性，这对于大型索引来说可能在计算上很昂贵。在微调时尤其如此，因为相对于索引中的文档数量，训练示例的数量可能很小。因此，与标准语言模型微调相比，训练检索器可能会增加重要的计算开销。在本节中，我们分析了可能使此过程更高效的策略，从而减少了过于频繁地重新计算所有文档嵌入的需要。

完整索引更新。

让我们从分析由于更新索引而产生的开销开始，与使用固定的检索器相比。为了比较不同模型的计算时间，我们将做出以下假设：对具有P参数模型的文档执行正向传递所需的时间为O（P）。虽然这种计算模型可能看起来很天真，但主要假设是文档大小是恒定的。1由于我们将长文档拆分为字数相似的段落，并在处理不同大小的文档时使用填充，因此这种假设在实践中是合理的。设K为语言模型检索和处理的文档数，Plm为语言模型的参数数，B为批大小。每个训练步骤的复杂度为4×B×K×Plm。2.

接下来，让N是索引中的文档数，Pretr是检索器的参数数。然后，重新计算完整索引的复杂度为N×Pretr。如果我们每R个训练步骤刷新一次索引，我们将获得以下开销：

如果我们将BERT基础架构用于我们的检索器，将T5-XL用于我们的语言模型，我们得到Pretr Plm≈1 25，从而增加开销：

如果我们使用一个包含3700万个文档的索引（相当于维基百科索引的大小），用64个批量大小的20个检索到的文档进行训练，并每1000步刷新一次索引，这将导致约30%的开销。

重新排名。

第二种策略是使用检索器检索更多的文档L，并使用最新的检索器重新嵌入和重新排列这些文档，并将生成的top-K传递给语言模型。在这种情况下，对前L个文档进行重新排序的开销等于B×L×Pretr。由于我们在每个时间步长执行此操作，因此开销等于

使用与之前相同的假设，我们最终得出开销为L 100×K的数量级。如果我们对比语言模型处理的文档多10倍的文档进行重新排序（即L=10×K），我们将获得10%的开销。但是，请注意，如果在检索器上执行了多次更新，则可能仍需要完全更新索引，因为可能无法在过时索引的top-L结果中检索到真正的top-k文档。在实践中，可以跟踪前K个重新排序的文档在前L个文档中的位置，并估计何时需要更新索引。

查询端微调。

最后，最后一种策略是将查询和文档的编码解耦。在这种情况下，我们固定与文档编码器相对应的参数，而只训练与查询编码器对应的参数。因此，文档的嵌入是固定的，我们不需要刷新索引，因此没有计算开销。正如我们将在实践中看到的那样，当有大型训练数据集可用时，修复文档编码器对不同任务的影响差异很大。对于我们考虑的大多数few-shot镜头设置，查询端微调不会对性能产生很大影响，有时甚至会略微提高性能。

3相关工作

3.1自然语言处理中的检索

知识密集型任务的检索。先前的工作表明，检索可以提高各种任务的性能，如问答（Voorhees et al.，1999；Chen et al.，2017；Kwiatkowski et al.，2019）、事实核查（Thorne et al.，2018）、对话（Dinan等人，2019）或引文推荐（Petroni et al.，2022）。从历史上看，这一信息检索步骤是使用术语匹配方法实现的，如TF-IDF或BM25（Jones，1972；Robertson等人，1995年）。对于开放域问答（Voorhees et al.，1999），文档通常从维基百科中检索（Chen et al.，2017）。最近，基于神经网络的密集检索器变得很流行。这些通常遵循双编码器架构（Yih et al.，2011；Huang et al.，2013；Shen et al.，2014），其中查询和段落被独立地编码为向量，并且相关性是使用内积或欧几里得距离计算的。流行的监督检索器包括DPR（Karpukhin et al.，2020），它被训练来区分负面段落中的相关段落，以及改进硬负面挖掘过程的扩展，如ANCE（Xiong et al.，2017）。我们建议读者参考Yates等人（2021）对密集检索技术的调查。检索后，对相关文档进行处理以产生最终输出。在开放域QA中，模型可以从检索到的文档中提取一段文本作为答案（Chen et al.，2017；Clark&Gardner，2018；Wang et al.，2019；Karpukhin et al.，2020），这是一种受阅读理解启发的方法（Richardson，2013；Rajpurkar et al.，2016）。最近，使用以检索到的文档为条件的seq2seq模型将答案生成为自由格式文本已经变得普遍（Lewis等人，2020；Izacard&Grave，2020；Min等人，2020）。这些架构也被证明可以减少对话代理的幻觉（Shuster等人，2021）。

检索器训练。

通过利用来自语言模型的信号或使用无监督学习，可以绕过对用于训练检索器的昂贵查询文档注释的需求。REALM（Guu等人，2020）和RAG（Lewis et al.，2020）通过将文档建模为潜在变量，并通过梯度下降最小化目标，联合训练检索器和语言模型。REALM使用MLM方法进行端到端预训练，但使用提取式BERT模型（Devlin等人，2019）。Guu等人（2020）还探索了在微调时进行查询端微调，以避免索引刷新，Lee等人（2021b）也在基于短语的检索中进行了探索。Izacard&Grave（2020）提出使用交叉注意力得分作为知识提炼的监督。Sachan等人（2021）利用读者输出的困惑，对读者和检索者进行联合训练。Sachan等人（2021）和Lee等人（2021a）都利用语言模型中的困惑和注意力得分，对猎犬进行了显著跨度掩码。Lee等人提出了反向完形填空任务。（2019）以无监督的方式预先训练密集检索器。Paranjape等人（2021）提出了一种方法，使用可访问输出的第二个“知情”检索器训练检索增强生成器，测试时检索器可以从中提取，Hofstätter等人（2022年）最近提出了一个训练集过滤/加权方法，以训练更强的检索增强生成器。Izacard等人（2022）探索了不同的对比学习方法来训练检索器，而Ram等人（2022年）使用文档中的重复跨度来创建伪阳性查询文档对。

检索增强语言模型。

连续缓存模型（Grave et al.，2017b）通过计算代币的先前和当前表示之间的相似性，定义了最近代币的概率分布。然后将该分布与语言模型的分布进行插值，以改进预测。后来，通过利用近似最近邻居搜索，用于计算这种分布的令牌数量扩展到了更大的内存中（Grave et al.，2017a）。相关的kNN-LM模型（Khandelwal et al.，2020）用Transformer网络取代了LSTM，并将内存扩展到数十亿个令牌，从而大大提高了性能。最近，RETRO（Borgeaud等人，2021）通过将检索内存扩展到数万亿个令牌，并更改模型体系结构以将检索到的文档作为输入，对这些进行了扩展。

利用搜索引擎增强检索功能。

最近，不同的工作建议通过生成文本查询并使用检索到的文档作为附加上下文来训练大型语言模型与搜索引擎进行交互（Nakano等人，2021；Thoppilan等人，2022年；Shuster等人，2022.）。在few-shot问答的背景下，Lazaridou等人（2022）使用该问题执行搜索查询，并将检索到的文档添加到执行上下文学习的大型语言模型的提示中。

3.2few-shot学习

few-shot学习，即从极few-shot例子中学习的任务，已经研究了几十年（Thrun&Pratt，1998；芬克，2005；Vinyals等人，2016），但最近随着大型预训练模型的出现，人们对NLP的兴趣激增，这些模型表现出了涌现的few-shot学习能力（Wei等人，2022）。

使用大型语言模型进行上下文学习。

正如Radford等人提出的那样，为语言模型提供任务的自然语言描述。（2019）导致了few-shot学习的重大发展。GPT-3（Brown et al.，2020）证明了大型语言模型执行few-shot预测的能力，其中该模型以自然语言描述了任务，几乎没有例子。缩放模型大小、数据和计算对于实现这种学习能力至关重要，从而导致大型模型的进一步发展（Lieber等人，2021；Rae等人，2021；Smith等人，2022年；Chowdhery等人，2022.Smith等，2022）。Hoffmann等人（2022）重新审视了Kaplan等人的比例定律。（2020），表明用较小的模型对更多数据进行训练可能更有效，从而产生了参数效率提高的70B参数模型Chinchilla。

很少进行微调和基于提示的学习。

上述模型在不训练语言模型的参数的情况下，利用上下文中的指令执行few-shot学习。通过结合文本模板（“提示”）和各种形式的模型微调，也可以实现few-shot学习，要么完全更新模型的参数，例如用于分类（Schick&Schütze，2021a；Schick&Schutze，2021；Gao等人，2021；Tam等人，2021），要么生成（Schik&Schutse，2021b）。提示本身可以优化，例如通过搜索（Jiang等人，2020；Shin等人，2020）或仅更新部分模型（Logan等人，2021），或学习“软盈利”（Lester等人，2021；Li和Liang，2021）。由于其简单性，在这项工作中，我们要么使用简单的提示，要么在不进行预处理的情况下直接输入，然后执行全模型微调，这是一种类似于Le Scao&Rush（2021）的方法。

4个实验

在本节中，我们报告了对我们的语言模型在few-shot学习方面的实证评估。我们首先介绍了我们的实验设置，在第4.1节中描述了我们的评估基准，并在第4.2节中给出了我们模型的训练细节。然后，我们进行消融研究，以比较导致我们的主要模型的不同技术选择。最后，我们在few-shot镜头和完整数据集设置中，在不同的自然语言理解任务上评估了这个名为Atlas的模型。

4.1基准

为了评估我们的检索增强语言模型，我们考虑了以下基准，其中包括不同的任务。

知识密集型语言任务（KILT）。

首先，我们使用KILT评估套件（Petroni et al.，2020），包含11个数据集，对应于5项任务：事实核查、问答、对话生成、实体链接和填充。这些不同的任务需要解决关于世界的知识，这些知识可以在维基百科上找到。我们在KILT中包含的以下任务和数据集上评估了我们的模型：问答：NaturalQuestions（Kwiatkowski et al.，2019）、TriviaQA（Joshi et al.，2017）和HotpotQA（Yang et al.，2018）；空位填充：Zero Shot RE（Levy et al.，2017）和T-REx（Elsahar et al.，2018）；实体链接：AIDA CoNLL YAGO（Hoffart et al.，2011）；对话：维基百科的向导（Dinan et al.，2019）；以及事实核查：FEVER（Thorne等人，2018）。这些数据集的KILT版本与其原始版本不同，因为2019年8月维基百科转储中不存在的需要知识的实例已被删除。

大规模多任务语言理解（MMLU）。

我们的第二个主要评估基准是MMLU（Hendrycks等人，2021），它包含57个多选问答数据集（称为域），这些数据源于为人类设计的真实考试。这些涵盖了非常广泛的主题，例如高中数学、专业法、逻辑谬误和临床知识，可以大致分为四个子类：人文科学、社会科学、STEM和“其他”。我们专注于few-shot学习，基准测试的作者建议每个领域使用5个训练示例。除了5点设置之外，我们还考虑了三个额外的设置。第一种是零样本设置，根本没有训练数据。第二种，我们称之为多任务few-shot，是在所有任务的6-shot数据上训练单个模型，从而产生285个例子的训练集。最后一种，我们称之为迁移学习，利用了MMLU作者提供的其他多项选择QA任务中的额外训练示例，即MCTest（Richardson，2013）、RACE（Lai et al.，2017）、ARC（Clark et al.，2018）和OBQA（Mihaylov et al.，2019），从而形成了一组95k个示例的训练集。

其他基准。

此外，我们报告了流行的NaturalQuestions（Kwiatkowski et al.，2019）和TriviaQA（Joshi et al.，2017）数据集的原始开放域版本的结果。我们还根据FEVER的原始版本（Thorne et al.，2018）评估了我们的模型，该模型将事实核查作为文本声明的三方分类问题（“支持”：文本在维基百科中得到证据支持，“反驳”：声明与维基百科中的证据不一致，或者“信息不足”，没有足够的证据做出判断）。我们还进行了实验来评估我们模型的时间敏感性。在这里，我们从TempLAMA构建了一个数据集（Dhingra et al.，2022），由一组关于一系列主题的时间敏感完形填空问题组成，其中答案从2017年到2020年发生变化。当提供2017年与2020年的指数时，我们评估模型的准确性，以评估模型在多大程度上忠实地反映了测试时提供给他们的指数的内容，以及更新指数作为一种持续学习或模型可更新性方法的有效性。

4.2技术细节

我们现在描述预训练和微调模型的过程。我们重点关注第4.3节和第4.4节中进行的消融研究所使用的设置。我们稍后会给出更多关于用于最终模型的超参数的详细信息。

训练前。

对于预训练，我们使用无监督的Contriever模型初始化检索器模块，该模型使用BERT基础架构。我们使用T5预训练的权重来初始化语言模型。由于最初的T5预训练模型在训练集中包括监督数据，我们使用仅在未标记文本上训练的1.1版本模型。具体来说，我们从T5 lm自适应变体开始初始化，因为它们的稳定性得到了提高。

对于第4.3节和第4.4节中进行的消融研究，我们使用包含3B重量的T5-XL。我们使用AdamW对所有模型进行10000次迭代的预训练，批处理大小为64，读取器的学习率为10-4，检索器的学习率是10-5，具有线性衰减和1000个预shot步骤。我们每1000步刷新一次索引。这意味着，与使用固定检索器进行训练相比，在预训练期间重新计算该指数10次，导致约30%的开销。我们将检索到的文档数设置为20。我们在第4.5节开始时详细介绍了用于最终模型训练的超参数。

微调。

当执行下游任务时，无论是在few-shot镜头设置中还是在大型训练集中，我们都会进行微调，使我们的模型适应这些任务。对于少量KILT消融实验，我们执行固定数量的微调迭代，而不是使用早期停止。更确切地说，我们决定对64次拍摄设置使用50次迭代，对1024次拍摄设置则使用200次迭代。在这两种情况下，我们都使用了32个示例的批量大小，具有线性衰减的4×10−5的学习率，以及读取器和检索器的5个预shot步骤。

未标记的数据集。

最后，我们讨论了用于训练模型的未标记文本数据集，这些数据集形成了检索索引。首先，我们考虑2021 12月20日的维基百科转储，我们保留了列表和信息框，通过在条目之间添加分号分隔符将其线性化。我们将文章按节拆分，并将长段拆分为大小相等且包含不到200个单词的段落。这导致了总共37M个段落，平均包含78个单词。我们还使用了来自2020-10常见爬网转储的文档，这些文档经过CCNet管道预处理（Wenzek et al.，2020）。我们以与Gopher类似的方式执行额外的文档过滤（Rae等人，2021）。更准确地说，我们根据文档长度、平均单词长度、字母数字字符的比例和重复标记的数量来过滤文档。这导致总共有350M个通道。相同的段落用于索引和模型预训练。在预训练期间，我们确保从检索到的文档中筛选出我们正在训练的段落，以防止模型简单地检索它正在去挖掘/生成的段落，并简单地使用它来解决预训练任务。

4.3训练前损失和任务

我们通过比较不同的预训练任务和目标函数来联合训练检索器和语言模型，从而开始我们的消融研究。我们的目标是回答以下研究问题：

（RQ 1）对整个模型进行联合预训练是否会带来更好的少杆表现？

（RQ 2）检索器的最佳目标函数和最佳借口任务是什么？

我们首先通过使用掩码语言建模任务对模型进行预训练来比较第2.2节中介绍的检索器的训练目标。我们在64次和1024次KILT基准的子集上评估了这些模型：NaturalQuestions、FEVER和维基百科的Wizard，以及两个基线：“闭卷”（即非增强T5）基线，在相同数据上预训练，并从Contriever和T5 lm adapt初始化。我们在表1中报告了结果。首先，我们注意到闭卷基线的较差性能，这表明了扩充的重要性。接下来，我们观察到，用检索对我们的模型进行预训练对于在少量拍摄任务中获得良好性能很重要。事实上，所有在预训练期间包括检索的模型都大大优于没有联合预训练的基线。接下来，我们比较了一个使用固定检索器预先训练的模型，以及使用各种检索器训练目标的模型。在与预训练目标相对应的MLM验证指标上，我们观察到联合训练检索器会带来很大的改进。在64次shot的下游任务中，这种效果往往不太明显，在1024次shot中几乎不存在。我们认为，这证明了预训练的最大影响是对语言模型的影响，该模型学习使用和聚合检索到的文档中的信息。最后，我们没有观察到不同的检索器训练目标之间存在显著的系统差异。因此，我们决定在随后的实验中使用困惑蒸馏，因为它往往比EMDR2或ADist更稳定，并且比LOOP更高效。

接下来，我们比较表2中第2.3节中介绍的不同的自监督借口任务。在这里，我们观察到所有三个任务都有类似的结果，掩码语言建模有一个小优势。因此，在接下来的内容中，我们采用掩码语言建模进行预训练。

最后，我们在预训练期间考虑了索引和训练数据的数据源Wikipedia和普通爬网的不同组合。在所有情况下，当执行少量微调时，我们都使用2021维基百科转储作为索引。我们在表3中报告了结果。首先，我们观察到，使用基于维基百科的索引可以带来更好的下游性能。对此可能有两种解释：首先，当我们使用维基百科进行few-shot拍摄任务时，当使用相同的数据进行训练时，模型可能会更好地适应。另一种解释可能是，维基百科是一个比普通爬行更高质量、更密集的知识来源。其次，当使用通用爬网索引时，我们观察到维基百科数据上的预训练会导致比使用通用爬网数据更低的性能。我们认为，主要原因是两个域之间的分布不匹配导致检索到的文档通常不太相关。反过来，这可能意味着预训练的效率较低，因为语言模型没有充分利用文档中的信息。因此，在下文中，我们决定将来自两个域的数据组合用于索引和预训练数据。

4.4微调

在本节中，我们对如何将我们的模型应用于下游任务进行了消融研究，这依赖于微调。特别是，我们想调查以下研究问题：

（RQ 3）如何在训练数据有限的任务中有效地微调Atlas？

为了回答这个问题，我们比较了不同的策略来微调检索器模块，如第2.4节所述。我们在表4中报告了结果。首先，关于预训练，我们观察到，在微调过程中保持检索器固定会导致64次和1024次few-shot设置的性能显著下降。其次，重新排序策略（第2行）导致与完全更新索引（第1行）非常相似的结果，同时显著提高了效率。最后，仅对查询编码器进行微调也会产生强大的结果：特别是，在64镜头设置中，这比执行完全微调略强，我们将其归因于过度拟合的机会较少。另一方面，在1024镜头的设置中，进行全面的微调会产生更强的结果，尤其是在NaturalQuestions上。因此，在下文中，我们对少量示例的实验使用查询端微调，对较大数据集使用标准微调。

4.5训练和评估Atlas

在本节中，我们应用前几节消融的结果来训练一系列Atlas模型，参数范围从770M到11B。更具体地说，我们使用困惑提取目标函数，以及掩码语言建模借口任务。我们使用Wikipedia和Common Crawl数据的组合对这些模型进行预训练，用于训练数据和索引内容。我们检索20个文档，每2500步更新一次索引，并对前100个文档进行重新排序。我们使用批量大小为128的AdamW预训练10000次迭代的模型。

4.5.1 MMLU结果

如第4.1节所述，我们考虑MMLU的四种设置：1）零样本设置，我们直接应用预处理模型，无需进行少量的shot微调；2）五shot设置，我们针对57个域中的每个域使用5个训练示例微调模型；3）五shot多任务设置，其中，不是针对每个域独立微调模型，我们训练单个模型来执行所有任务，4）访问多个辅助数据集的设置，总共有95K个训练示例。我们训练模型生成与正确答案选项（“A”、“B”、“C”或“D”）相对应的字母，并在测试时从4个字母中选择最有可能的答案。完整的技术细节见附录A.1。

性能与参数。

首先，我们将Atlas与不同尺寸的5shot、5shot多任务和完整设置的闭卷模型进行比较，并在表5中报告结果。在这些设置中，Atlas比闭卷基线高出6.6到15.6分，证明了检索在57个领域中对few-shot语言理解的一致实用性。在770M参数的少量拍摄设置中，闭卷T5的表现很难明显优于随机（25%），而等效的Atlas实现了约40%的表现，尽管其尺寸较小，但明显优于随机。所有模型都会随着更多的数据而改进，但有趣的是，与较大的模型相比，770M模型并没有从few-shot多任务学习中受益那么多（对于闭卷模型，它实际上损失了3分），这表明较小的模型很难在few-shot设置中把握任务之间的协同作用。较大的模型很好地利用了多任务设置，Atlas的改进程度超过了闭卷。例如，Atlas-11B提高了13分（43.4→ 56.4），但等效的闭式书本只提高了7（36.1→ 43.5）。最后，在迁移学习环境中，所有模型都有所改进，但Atlas封闭式模型之间的相对差距仍然相似。

去偏置。

在微调时，我们将哪个答案选项与哪个答案字母排列在一起，以减少过度拟合，并鼓励在答案字母上使用统一的先验。然而，该模型可能仍然对一些字母表现出偏见，特别是在few-shot镜头设置中，因此除了上述标准推理之外，我们还包括第二种“去偏见”推理模式。在这里，我们运行4次正向传递，每个循环排列一次问题中的答案字母-答案选项分配，例如，分配给字母“A”的答案选项变为“B”，原来的“B”变为“C”等等。3然后，我们将4个概率求和，以获得最终预测，这减少了对其中一个答案字母的伪偏差（更多细节见附录A.1）。结果如表6所示。我们发现，在零样本和五拍设置中，去偏非常有效，分别提高了10.3和4.5个点。当有更多的训练数据可用时，对去偏的需求就会减少，导致多任务和完整数据设置仅提高0.2点。

与已出版作品的比较

接下来，我们将带有去偏的Atlas-11B结果与最近报道的具有最先进的大型语言模型（如GPT-3或Chinchilla）的结果进行比较，这些模型需要大量的计算来训练。我们在表7中报告了结果。我们发现，Atlas在零样本方面的表现明显优于随机，结合去偏倚推理，获得的零样本分数超过文献中报告的GPT3五shot结果（47.1%比43.9%）（Hendrycks等人，2021）。对于5shot点设置，Atlas比GPT-3高4%，同时使用的参数减少了15倍，预训练计算减少了10倍。4当在组合的5shot点数据上进行多任务训练时，Atlas提高到56.6%，接近Gopher的5shot点性能（60.0%）。最后，在全数据设置上，我们在MMLU作者推荐的辅助数据上进行训练，Atlas的总体准确率达到65.6%，接近最先进水平。有趣的是，在这种设置中，Atlas显著优于GPT-3，而在5shot设置中，它们的性能相似。

4.5.2开放域问答结果

接下来，我们在两个开放领域问答基准上评估Atlas：NaturalQuestions和TriviaQA。我们与之前的工作进行了比较，包括在使用64个示例的few-shot镜头设置中，以及使用完整的训练集，并在表8中报告了结果。在这些需要高度记忆的基准上，我们清楚地看到了检索增强的好处。Atlas-11B在NaturalQuestions和TriviaQA的64次问答中都获得了最先进的结果。特别是，它的性能明显优于更大的模型，如PaLM，或需要更多训练计算的模型，例如Chinchilla。在使用完整的训练集时，Atlas还获得了最先进的结果，例如，将NaturalQuestions的准确率从55.9%提高到60.4%。该结果是通过使用CCNet和2021 12月的维基百科语料库（我们的默认索引设置）组成的索引获得的。在第5.2节中，我们考虑使用由不同日期存档的维基百科语料库组成的索引，并在使用与NaturalQuestions在时间上匹配的索引时，证明NaturalQuentions的额外+3.6%。我们在附录a.2中报告了作为模型大小函数的性能以及详细的超参数。

Atlas还与最近探索检索的工作相比较，后者用非常大的模型增强了few-shot问答。Lazaridou等人（2022）使用Gopher在15个镜头的设置中探索NaturalQuestions，使用谷歌搜索检索到的50个段落来增加问题。该方法包括从每个检索到的段落中生成4个候选答案，然后使用受RAG启发的分数（Lewis et al.，2020）或更昂贵的方法重新排序。这种方法（未在我们的表格中显示）实现了32.7%（RAG）和38.4%（Ensemble）的精确匹配分数，每个测试时间问题需要50次（RAG或450次（Ensembles）Gopher-280B的前向传球。Atlas使用相同的15个训练示例和50个段落，实现了38.7 EM，尽管参数减少了25倍，并且需要相对可以忽略不计的计算。

4.5.3发shot结果

我们在表9中报告了原始3类FEVER事实核查测试集的结果。我们考虑64镜头设置，其中训练示例从整个训练集中均匀采样。与开发集和测试集不同，训练集是不平衡的，阳性标签比阴性标签多，这对few-shot学习构成了挑战。在这种情况下，我们获得了64.3%的准确率。我们还报告了一个15杆的设置，从每个类中均匀抽样5个示例，以与Gopher（Rae等人，2021）发布的结果进行比较，其中Atlas得分56.2%，比Gopher高5.1分。最后，我们在完整的训练集上对模型进行了微调，在ProoFVer的1.5%范围内取得了78%的分数，ProoFVer使用了一种特殊的体系结构，一种经过句子级注释训练的检索器，并提供了用FEVER发布的维基百科语料库，而Atlas则从CCNet和2021 12月的Wikipedia转储中检索。如果我们给Atlas一个由FEVER维基百科语料库组成的索引，我们就设定了80.1%的最新水平

4.5.4KILT结果

最后，我们对KILT上的Atlas进行了评估，KILT是一个由几个不同的知识密集型任务组成的基准，如第4.1节所述。我们在表10中报告了测试集的结果，其评估可在线获取5。数据集的KILT版本被过滤，因此我们在其他地方评估的数据集的结果在KILT上无法直接比较（即FEVER、NQ和TQA）。我们同时考虑64镜头设置和完全微调设置，在这两种情况下，我们在每个数据集上单独训练Atlas。有关超参数和开发集结果的更多详细信息，请参见附录A.3。对于64shot，我们大大超过了随机性能，甚至与排行榜上的一些完全调整过的模型相竞争，例如对于FEVER，我们的64shotAtlas仅落后于Sphere、SEAL和Re2G 2-2.5点，在零样本RE上优于Sphere和SEAL。在全数据集设置中，Atlas与3个数据集的最先进水平相差3%，并在剩下的五个数据集中设置了最先进的技术。

5分析

5.1可解释性和泄漏

Atlas这样的半参数模型的一个优点是能够检查检索到的项目，以帮助解释。为了更好地了解Atlas检索的效果，以及它如何使用检索到的段落，我们检查了多任务few-shotMMLU的检索到的章节。如图3的左侧面板所示，该模型从CCNet检索其大部分段落（平均85%）。维基百科约占检索段落的15%，这高于我们在统一先验下的预期，因为维基百科只占索引的10%左右。维基百科检索的比例在MMLU领域之间有所不同，该模型在STEM领域更大程度上使用维基百科，在社会科学领域使用最少。维基百科使用最多的领域是“抽象代数”（73%），最少的是“道德场景”（3%）。我们还注意到，MMLU微调后的Atlas并没有大量使用维基百科信息框中的段落。

我们还可以分析段落的内容，以评估它们对完成下游任务的用处。图3的中间面板显示了检索到的文档包含正确答案选项文本的频率。在前25段中，30%的试题中至少有一个提到了正确答案。6右侧面板显示，当正确答案选项文本在检索到的段落中出现得更频繁时，MMLU的准确性会提高，从没有出现答案选项时的55%上升到提及15次以上的77%。

对检索到的文档进行的人工分析表明，文档有助于以多种不同的方式回答问题。对50个正确回答的问题样本进行的手动检查显示，44%的问题至少包含部分有用的背景信息。这些文件将提高非专业人员正确回答的可能性，例如围绕问题引用的上下文线索，或基于数量的问题的有用数字，这有助于将答案选项缩小到更小的范围。在另外26%的案例中，一段话包含了回答问题的所有必要信息，以直截了当的方式陈述。如果阅读得当，这些段落会使问题变得简单，并且通常包括规范定义或问题中要求的确切数字答案等信息。28%的检索集不包含明显的信息，这会使问题变得更容易。最后，2%的人在一段话中包含了逐字逐句的问题及其答案。

鉴于MMLU是由预先存在的考试创建的，这些问题可能会出现在开放网络上。根据网络数据训练的模型（或者，在我们的情况下，从中检索）冒着不是通过概括而是通过逐字记忆来正确回答的风险，这可能会导致误导性的高分。在一些非常大的语言模型中，它们可以逐字记忆和回忆大部分的预训练数据（Carlini等人，2021），有时会努力从预训练数据中过滤出现的下游实例，但这在文献中尚未针对MMLU执行。为了评估我们索引中MMLU泄漏的普遍性，我们手动检查了问题的检索结果，其中问题（没有答案选项）和文章之间最长的n-gram重叠至少是问题长度的75%。这导致我们的CC-Net语料库中2.8%的问题被泄露。

Atlas等检索增强模型的一个好处是其知识的可编辑性（其他分析见第5.2节）。为了估计纯粹的、未泄漏的性能，我们可以从检索到的结果中筛选出任何可能泄漏的段落，并重新运行语言模型。当控制这种泄漏时，MMLU得分略有下降，从56.4%降至55.8%（-5.5%）。我们注意到，与最近非常大的模型的预训练语料库相比，我们的CC网络语料库相对较小，这些模型在多达1.4万亿个代币上训练（Hoffmann et al.，2022），是我们指数的35倍，这使得在这种大小的语料库上训练的模型可能会观察到更多的MMLU泄漏示例，但在非检索增强模型中检测这种泄漏是具有挑战性的。

5.2时间敏感性和可更新性

检索增强模型的一个好处是，通过在测试时更新或交换索引，它们可以保持最新，而无需重新训练。为了评估Atlas中这一机制的有效性，我们首先构建了一个从TempLAMA导出的时间敏感问题数据集（Dhingra et al.，2022）。TempLAMA是一个模板化完形填空问题的集合，源于Wikidata和Wikidata，其中正确答案会随着时间的推移而变化。我们从这个数据集中选择了一个子集的问题，这些问题在2017年和2020年有不同的答案，例如，问题：西奥·沃尔科特效力于___答案：阿森纳足球俱乐部（2017），埃弗顿足球俱乐部（2020），并形成了一个由248个训练、112个发展和806个测试问题组成的小训练集。

使用这个数据集，我们使用问题和2017年的答案对闭卷T5-XXL和Atlas进行了微调，为Atlas提供了2017年维基百科索引，然后在2017年的测试集上测量精确匹配的准确性。结果可以在表11的第一行和前两列中找到。我们首先观察到，正如预期的那样，Atlas大大优于T5（57.7%c.f.12.1%）。我们还注意到，正如所期望的那样，T5和Atlas在使用2017年的答案进行训练时，几乎从未从2020年开始生成答案，得分分别为2.8%和1.5%（表11的第一行和第二列）。然而，如第2行所示，我们可以将Atlas索引转换为2020年维基百科索引，而无需重新训练，并发现Atlas相应地更新了其预测，2020年的准确率上升到与其2017年的性能（53.1%）相似的水平，而纯参数T5没有这样的可更新性机制。

这表明Atlas可以忠实地使用其提供的索引。此外，这种零样本可更新性机制具有保持最新的有用特性，而无需最新的注释数据，也无需持续、终身的预训练，这对于大型仅参数模型可能是必需的。表11的第3行和第4行完成了这幅图，这次我们用2020年的答案进行了训练，并证明阿特拉斯也可以有效地将零样本及时向后转移到2017年（50.1%）。有趣的是，T5无法很好地回答2020年的问题，即使是用2020年答案进行了训练（3.6%），很可能是因为它是根据2020年之前的数据进行的预训练（Dodge等人，2021）。

我们还研究了NaturalQuestions的时间效应。NaturalQuestions是一个由谷歌搜索引擎在短时间内收集的搜索查询组成的数据集。因此，数据具有很强的时间偏见，例如，2018年世界杯有很多问题。此外，在没有具体说明时间上下文的情况下，有些问题是模棱两可的。例如，对于“爱尔兰上一次在特威克纳姆击败英格兰是什么时候”的问题，NaturalQuestions的预期答案是2018年，而爱尔兰也在2022年以及之前的许多其他时间在特威肯纳姆击败了英格兰。在表12中，我们报告了使用不同的维基百科转储对Atlas进行微调所获得的结果。我们观察到，2018年12月的维基百科转储接近数据收集日期，无论是少量微调还是完全微调，都会产生最佳结果。特别是，它带来了一个新的最先进的64 EM自然问题。

5.2.1指数压缩

保持密集的检索索引可能会占用大量内存，尤其是当索引项的数量被缩放时。在本节中，我们简要分析了在a）维基百科索引和b）上述大多数实验中使用的CC和维基百科索引组合的情况下，Atlas索引的内存需求。

Atlas的检索组件有两个记忆压力来源——段落本身和文档嵌入索引。标记化的段落一旦二进制化，维基百科和组合索引分别需要11GB和130GB的存储空间。这些段落不需要存储在昂贵的GPU RAM中，甚至可以将内存映shot到磁盘，在节点之间进行分片，或者在需要时进行压缩，因此在这种情况下不代表限制性的硬件挑战。然而，嵌入索引本身必须存储在GPU RAM中以进行快速搜索，因此其大小更敏感。在上面的实验中，我们对索引执行精确搜索，这是通过在所有可用的GPU上分割索引并并行计算搜索来实现的。索引以fp16精度存储，导致维基百科和组合索引的GPU总内存需求分别为49 GB和587 GB。

对索引的这种大GPU内存需求限制了可访问性和部署的简易性。然而，许多索引压缩技术可用于最近邻搜索，这通常可以以一定的检索精度为代价显著降低内存需求。继Izacard等人（2020）之后，我们探讨了乘积量化（PQ，Jégou等人，2011），这是一种流行的有损压缩技术，在不同压缩级别下对Atlas-3B 64次NQ任务精度的影响。

结果如图4所示。我们发现，在性能显著下降之前，可以进行实质性压缩。也就是说，维基百科索引可以从49GB压缩到4GB，检索精度和精确匹配的下降可以忽略不计。同样，组合索引可以从587GB压缩到50GB，而不会严重退化，这表明组合索引可以加载到单个80GB GPU上。

6讨论

在本文中，我们介绍了一个大型检索增强语言模型Atlas。通过联合预训练检索器模块和语言模型，我们表明Atlas在各种知识密集型任务上具有强大的few-shot学习能力，包括NaturalQuestions、TriviaQA、FEVER、8个KILT任务和57个MMLU任务。例如，当对64个例子进行训练时，Atlas-11B在NaturalQuestions上的准确率达到42%以上，在TriviaQA上的准确度达到84.7%，与需要50倍训练前计算的540B参数模型PaLM相比，这提高了近3分。我们还提供了详细的消融和分析，以确定在训练这种检索增强模型时哪些因素是重要的，并展示了Atlas的可更新性、可解释性和可控制性能力。最后，我们证明了Atlas在完整的数据集设置中也很强大，在NaturalQuestions、TriviaQA、FEVER和5 KILT任务中获得了最先进的新结果。