论文题目：REPLUG: Retrieval-Augmented Black-Box Language Models
  论文日期：2023/05/24
  论文地址：https://arxiv.org/abs/2301.12652

Abstract

  这篇论文提出REPLUG，一种检索增强的语言建模框架，将语言模型(LM)视为黑盒，并用可调的检索模型对其进行增强。与之前使用特殊交叉注意力机制训练语言模型以编码检索文本的检索增强语言模型不同，REPLUG只是将检索文档放在冻结(frozen)的黑盒语言模型的输入前面。这种简单的设计可以很容易地应用于任何现有的检索和语言模型。此外，语言模型可以用来监督检索模型，然后可以找到帮助语言模型做出更好预测的文档。
  作者的实验表明带有微调检索器的REPLUG将GPT-3(175B)在语言建模上的性能提高了6.3%，将Codex在five-shot MMLU上的性能提高了5.1%。

1. Introduction

  大型语言模型(LLM)，如GPT-3和Codex，在广泛的语言任务中表现出令人印象深刻的表现。这些模型通常在非常大的数据集上训练，并在其参数中隐式地存储了大量的世界或领域知识。然而，它们也容易产生幻觉，不能代表训练语料库中知识的完整长尾。相比之下，检索增强语言模型可以在需要时从外部数据存储中检索知识，可能会减少幻觉并增加覆盖率。以前的检索增强语言模型方法需要访问内部语言模型表示(例如，训练模型或索引数据存储)，因此很难应用于非常大的语言模型。此外，许多同类中最好的LLM只能通过API访问。这些模型的内部表示没有暴露出来，也不支持微调。
  这篇论文提出REPLUG(Retrieve and Plug)，一种新的检索增强的语言模型框架，其中语言模型被视为黑盒，检索组件被添加为可调的即插即用模块。给定一个输入上下文，REPLUG首先使用现成的检索模型从外部语料库检索相关文档。检索到的文档被添加到输入上下文并输入到黑盒LLM中以做出最终预测。由于语言模型上下文长度限制了可以前缀的文档数量，这里还引入了一个新的集成方案，该方案与相同的黑盒语言模型并行编码检索的文档，允许我们轻松地以计算换取准确性。REPLUG非常灵活，可以与任何现有的黑盒语言模型和检索模型一起使用，如下图所示：

  这篇论文还提出了REPLUG LSR(REPLUG with LM-Supervised Retrieval)，一种训练方案，可以利用黑盒语言模型的监督信号进一步改进REPLUG中的初始检索模型，关键思想是使检索器适应LM，这与之前使语言模型适应检索器的工作形成了鲜明对比。作者使用了一个训练目标，该目标倾向于检索文档，以改善语言模型的困惑，同时将LM视为一个冻结的黑盒评分函数。
  实验表明，REPLUG可以提高不同黑盒语言模型在两种语言建模上的性能下游任务，包括MMLU和开放域QA。例如，REPLUG可以将Codex(175B)在MMLU上的性能提高4.5%，实现了与指令微调的Flan-PaLM 540B相当的结果。此外，用论文中的训练方案(即REPLUG LSR)调整检索器会带来额外的提高，比如GPT-3 175B语言建模提高了6.3%。本文工作首次展示了检索大型语言模型(>100B模型参数)的好处，既减少了语言模型的困惑度，又提高了上下文学习性能。
  这篇论文的主要贡献如下：
  (1) 提出REPLUG，第一个检索增强语言建模框架，用于通过检索增强大型黑盒语言模型；
  (2) 提出一种训练方案REPLUG LSR，以进一步将现成的检索模型适应于语言模型，使用语言建模分数作为监督信号，从而提高检索质量；
  (3) 在语言建模、开放域QA和MMLU上的评估表明，REPLUG可以提高各种语言模型(如GPT、OPT和BLOOM)的性能，包括高达175B参数规模的大型语言模型。

2. Background and Related Work

2.1 Black-box Language Models

  大型语言模型(>100B)，如GPT-3、 Codex和Yuan1.0，由于商业考虑而不开源，仅作为黑盒API提供，用户可以通过它发送查询和接收响应。另一方面，即使是开源语言模型，如OPT-175B和BLOOM-176B也需要大量的计算资源在本地运行和微调。例如微调BLOOM-176B需要72块 A100 GPU(80GB内存，每个1.5万美元)，使得资源有限的研究人员和开发人员无法访问它们。传统上，检索增强模型框架侧重于白盒环境，其中语言模型被微调以合并检索文档。然而，大型语言模型的规模不断扩大和黑箱特性使得这种方法不可行。为解决大型语言模型带来的挑战，这篇论文研究了黑盒环境下的检索增强，其中用户只能访问模型预测，而不能访问或修改其参数。

2.2 Retrieval-augmented Models

  事实证明，用从各种知识库中检索到的相关信息来增强语言模型，可以有效地提高各种NLP任务的性能，包括语言建模和开放域问答。具体来说，使用输入作为查询，检索器首先从语料库中检索一组文档(即token序列)，然后语言模型将检索到的文档作为附加信息进行最终预测。这种检索风格可以添加到encoder-decoder和decoder-only模型中。例如，Atlas通过将文档建模为潜在变量来与检索器联合共同对encoder-decoder模型进行微调，而RETRO则改变了decoder-only的架构，以合并检索到的文本，并从头开始预训练语言模型。这两种方法都需要通过梯度下降更新模型参数，不能应用于黑盒LM。另一条线路检索增强的语言模型，如kNN-LM检索一组tokens，并在LM的下一个token分布和推理时从检索的token计算出的kNN分布之间进行插值。虽然kNN-LM不需要额外的训练，但它需要访问内部的LM表示来计算kNN分布，这对于像GPT-3这样的大型语言模型并不总是可用的。
  本文研究了用检索来改进大型黑盒语言模型的方法。虽然相关工作已经证明了使用冻结检索器可以提高GPT-3在开放域问答上的性能，但作者在更通用的环境中处理这个问题，包括语言建模和理解任务。本文还提出了一种集成方法来合并更多的文档，以及一种训练方案来进一步使检索器适应大型语言模型。

3. REPLUG

  本文提出REPLUG(Retrieve and Plug)，一种新的检索增强的语言模型范式，其中语言模型被视为黑盒，检索组件被添加为一个潜在的可调模块。
  如下图所示，给定一个输入上下文，REPLUG首先使用检索器从外部语料库中检索一小部分相关文档，然后通过LM并行传递每个检索文档与输入上下文的连接，并集成预测概率。

3.1 Document Retrieval

  给定输入上下文 $x$，检索器旨在从语料库 D = { d 1 , . . . , d m } \mathcal D = \{d_1, ..., d_m\} D={d1​,...,dm​} 中检索与 $x$ 相关的一小部分文档。根据之前的工作，作者使用基于双塔编码器(dual encoder)架构的密集检索器，其中编码器用于编码输入上下文 $x$ 和文档 $d$。具体来说，编码器通过对 $d$ 中的tokens进行最后一个隐藏表示的均值池化，将每个文档 $d\in \mathcal{D}$ 映射到嵌入 $E(d)$。在查询时，将相同的编码器应用于输入上下文 $x$ 以获得查询嵌入 $E(x)$。查询嵌入和文档嵌入之间的相似度是通过它们的余弦相似度计算的：$$s(d,x)=cos(E(d),E(x))$$在此步骤中检索与输入 $x$ 相比具有最高相似度分数的top-k个文档。为了高效检索，作者预先计算每个文档 $d\in \mathcal{D}$ 的embedding，并在这些嵌入上构建FAISS索引。

3.2 Input Reformulation

  检索到的top-k个文档提供了关于原始输入上下文 $x$ 的丰富信息，并可以潜在地帮助语言模型做出更好的预测。将检索到的文档作为LM输入的一部分的一种简单方法是在 $x$ 前面加上所有 $k$ 个文档。然而，考虑到语言模型的上下文窗口大小，这个简单的方案从根本上受到我们可以包括的文档数量(k)的限制。为了解决这个限制，作者采用了一种集成策略描述如下。根据上述相似度计算中的评分函数，假设 ${\mathcal{D}}^{\prime }\subset \mathcal{D}$ 包含 $k$ 个与 $x$ 最相关的文档，将每个文档 $d\in {\mathcal{D}}^{\prime }$ 添加到 $x$，将这个拼接分别传递给LM，然后从所有 $k$ 个传递中集成输出概率。形式上，给定输入上下文 $x$ 及其 $top-k$ 个相关文档 ${\mathcal{D}}^{\prime }$，下一个token $y$ 的输出概率可以通过加权平均计算得到：$$p(y|x,{\mathcal{D}}^{\prime })=\sum _{d\in {\mathcal{D}}^{\prime }}p(y|d\circ x)\cdot \lambda (d,x)$$其中 $\circ$ 表示两个序列的拼接，权重 $\lambda (d,x)$ 是基于文档 $d$ 和输入上下文 $x$ 之间的相似度得分：$$\lambda (d,x)=\frac{e^{s(d,x)}}{{\sum }_{d\in {\mathcal{D}}^{\prime }}{e}^{s(d,x)}}$$虽然这种集成方法需要运行LM $k$ 次，但交叉注意力在每个检索文档和输入上下文之间执行。因此，与将所有检索到的文档放在前面的方法相比，该集成方法不会产生额外的计算开销。

4. REPLUG LSR: Training the Dense Retriever

  作者不再仅仅依赖于现有的神经密集检索模型，而是进一步提出了REPLUG LSR(REPLUG with LM Supervisored retrieval)，它通过使用LM本身来调整REPLUG中的检索器，以提供关于应该检索哪些文档的监督信息。
  受"Questions are all you need to train a dense passage retriever"这篇论文的启发，该方法可以被视为调整检索文档的概率以匹配语言模型的输出序列困惑度的概率。换句话说，作者希望检索器找到困惑度分数较低的文档。如下图所示，训练算法包括四个步骤：
  (1) 检索文档并计算检索似然度；
  (2) 通过语言模型对检索文档进行评分；
  (3) 通过最小化检索似然度和语言模型的分数分布之间的KL散度来更新检索模型参数；
  (4) 异步更新数据存储索引。

  似然度，即可能性，likelihood

4.1 Computing Retrieval Likelihood

  从语料库 $\mathcal{D}$ 中检索具有最高相似度分数的 $k$ 个文档 ${\mathcal{D}}^{\prime }\subset \mathcal{D}$，给定输入上下文 $x$，计算每个检索文档 $d$ 的检索似然度：$$P_R(d|x)=\frac{e^{s(d,x)/\gamma }}{{\sum }_{d\in {\mathcal{D}}^{\prime }}{e}^{s(d,x)/\gamma }}$$其中 $\gamma$ 是控制softmax温度的超参数。理想情况下，检索似然度是通过边缘化语料库 $\mathcal{D}$ 中的所有文档来计算的，这在实际中是难以实现的。因此，作者通过仅在检索文档 ${\mathcal{D}}^{\prime }$ 上边缘化来近似检索似然度。

4.2 Computing LM likelihood

  使用LM作为评分函数来衡量每个文档对LM复杂度的改善程度。具体来说，给定输入上下文 $x$ 和文档 $d$，首先计算$P_{LM}(y|d,x)$，即LM的真实值输出y的概率，概率越高，文档 $d$ 在改善LM的困感度方面越好。然后，再计算每个文档 $d$ 的LM似然度：$$Q(d|x,y)=\frac{e^{P_{LM}(y|d,x)/\beta }}{{\sum }_{d\in {\mathcal{D}}^{\prime }}{e}^{P_{LM}(y|d,x)/\beta }}$$其中 $\beta$ 是另一个超参数。

4.3 Loss Function

  给定输入上下文 $x$ 和相应的基本事实连续性 $y$，计算检索似然和语言模型似然。稠密检索器通过最小化这两个分布之间的KL散度来训练：$$\mathcal{L}=\frac{1}{|\mathcal{B}|}\sum _{x\in \mathcal{B}}KL\left(P_R\right(d|x)\left|\right|Q_{LM}(d|x,y\left)\right)$$其中 $\mathcal{B}$ 是一组输入上下文。在最小化损失时，只能更新检索模型参数。由于黑盒假设，LM的参数是固定的。

4.4 Asynchronous Update of the Datastore Index

  因为检索器中的参数在训练过程中被更新，所以之前计算的文档嵌入不再是最新的。因此，作者在每个T个训练步骤中重新计算文档嵌入并使用新的嵌入重建高效的搜索索引。然后使用新的文档嵌入和索引进行检索，并重复训练过程。

5. Training Setup

  训练时，首先在REPLUG中描述模型设置，然后在REPLUG LSR中描述训练检索器的程序。

5.1 REPLUG

  理论上，任何类型的检索器，无论是密集的还是稀疏的，都可以用于REPLUG。根据之前的工作，作者使用Contriever作为REPLUG的检索模型，因为它已经证明了强大的性能。

  Contriever是一个通过对比学习框架来预训练的信息检索模型。Unsupervised Dense Information Retrieval with Contrastive Learning

5.2 REPLUG LSR

  对于REPLUG LSR，作者使用Contriever模型初始化检索器，使用GPT-3 Curie作为监督LM来计算LM似然度。

5.2.1 Training data

  作者使用从Pile训练数据中采样800K条序列作为训练query，每条序列有256 tokens，每个query被分成两部分：前128 tokens用作输入上下文 $x$，后128 tokens用作真实值连续性 $y$。对于外部语料库 $\mathcal{D}$，从Pile训练数据中采样36M个文档，每个文档有128 tokens。为了避免繁琐的检索，确保外部语料库文档与训练查询的采样文档不重叠。

5.2.2 Training data

  为了使训练过程更高效，作者预先计算了外部语料库 $\mathcal{D}$ 的文档嵌入，并创建FAISS索引用于快速相似性搜索。给定一个查询 $x$，从FAISS索引中检索前20个文档，并计算温度为0.1时的检索似然和LM似然。使用Adam优化器训练检索器，learning rate为2e-5，batch size为64，warmup ratio为0.1。每3k steps重新计算文档嵌入，并对检索器进行微调，总共有25k steps。