【论文阅读】REPLUG: Retrieval-Augmented Black-Box Language Models

前言

• 原文地址：REPLUG: Retrieval-Augmented Black-Box Language Models

• 本文提出REPLUG，一个将语言模型视为黑盒检索增强的语言模型架构。在REPLUG中，仅将检索得到的文档拼接到原有输入前面即可，不需要像以前一样更新语言模型参数。该架构中可以通过更新检索器进一步提升性能。
  

REPLUG

• 给一个输入上下文
• REPLUG会首先从外部资源 D = { d 1 , … , d m } D=\{d_1,\dots,d_m\} D={d1​,…,dm​}中检索出一些相关文档
  • 使用基于双塔encoder（共享参数）的dense retrieval来检索文档，一个encoder用来编码输入$x$和文档$d$
  • 文档与输入的embedding都是对其中每个token最后一个隐藏层表达的平均值
  • 通过cos similarity计算$x$与$d$的相关性：$s(d,x)=cos(E(d),E(x))$
  • 预先计算文档的embedding，并利用FAISS来快速找到top-k文档
• 之后我们将每个检索到的文档与输入上下文进行拼接，并行输入到大模型中
  • 由于模型输入限制，无法将所有检索文档与输入$x$进行拼接
  • 采用聚合策略，拼接时，将每个top-k文档分别拼接在$x$前面，并将拼接结果分别输入到语言模型中。
• 最后聚合每个并行输入得到的预测概率
  • 对上面分别计算的结果进行聚合
    • 给定上下文输入 $x$ 和 top-k 相关文档集合 $D^{{}^{\prime }}$，下一个token $y$ 的生成概率由加权平均决定
      • $p(y|x,D^{{}^{\prime }})={\sum }_{d\in D^{{}^{\prime }}}p(y|d\circ x)\cdot \lambda (d,x)$
        • 其中 $\lambda (d,x)$是 $d$ 与 $x$ 相似度 $s(d,x)$ 进行softmax的结果

REPLUG LSR: Training the Dense Retriever

REPLUG LSR 可以看做 REPLUG的一个增强版本。在REPLUG中，我们使用的检索器可能不够适配语言模型，因此这里利用语言模型本身反馈的监督信号，来调整REPLUG中的检索器。

• 这里的监督信号可以告诉我们，什么样的文档应该被检索回来

核心思想：our approach can be seen as adjusting the probabilities of the retrieved documents to match the probabilities of the output sequence perplexities of the language model

• 其实就是匹配检索文档的概率与语言模型输出序列的概率
  • 输出序列的概率就是语言模型提供的监督信号
  • 这样做的理由
    • 如果模型输出的ground truth序列的概率更大，那么我们认为模型的效果越好
    • 我们认为，如果一个文档对模型的输出更有帮助，那么我们就认为这个文档更应该被检索回来，其检索的概率也应该更大。
    • 所以说，一个文档被检索回来的概率应该与使用这个文档得到输出序列的概率是正相关的，因此我们想要匹配检索文档的概率与语言模型输出序列的概率

---

这部分介绍如何计算检索文档概率分布与输出序列概率分布

Computing Retrieval Likelihood

给定输入 $x$，我们检索回来概率最大的top-k个文档，为$D^{{}^{\prime }}\subset D$，文档$d$的检索概率(likelihood)为

$P_R(d\mid x)=\frac{e^{s(d,x)/\gamma }}{{\sum }_{d\in {\mathcal{D}}^{\prime }}{e}^{s(d,x)/\gamma }}$

• $\gamma$是用来控制 softmax 温度的超参

• 按理应该在整个 $D$ 上进行，但是那样计算量太大，因此在 $D^{{}^{\prime }}$ 上近似计算

Computing LM likelihood

将语言模型用来评估每个文档对语言模型困惑度的提升程度，首先计算 $P_{LM}(y|d,x)$，这是给定 $x$ 和文档 $d$ 时，ground truth $y$ 的生成概率。如果这个概率越大，则说明当前文档对困惑度的提升程度越大。然后计算分布：

$Q(d\mid x,y)=\frac{e^{P_{LM}(y\mid d,x)/\beta }}{{\sum }_{d\in {\mathcal{D}}^{\prime }}{e}^{P_{LM}(y\mid d,x)/\beta }}$

• $\beta$是超参

---

有了两个分布之后，用loss function 对二者进行匹配

在给定 $x$ 和 $y$ 时，计算检索概率分布和语言模型概率分布，我们利用KL divergence来匹配两个分布，并用来优化dense retriever

$\mathcal{L}=\frac{1}{|\mathcal{B}|}{\sum }_{x\in \mathcal{B}}KL\left(P_R(d\mid x)\Vert Q_{\mathrm{LM}}(d\mid x,y)\right)$

• $B$ 是输入 $x$ 的集合
• 我们最小化损失函数来优化检索器，LM保持不动

因为检索器参数在训练过程中更新，参数更新后document embedding会变化，因此每隔$T$步就重新算一次document embedding，并重复上述过程

Training Setup

Model

• LM: GPT-3(for REPLUG LSR)
• Retriever：Contriver(2022新模型)

Training data

• 所有训练数据都来自 Pile training data(包含不同领域文本的language model benchmark)

• 800K 个 256 token长的序列作为训练queries

  • 每个query分成两部分，前128token作为 input context $x$，后一半作为需要续写的ground truth $y$

• 外部语料库 $D$, 采样36M 128 token长的文档

Results

Language Modeling

• randomly subsampled Pile training data (367M documents of 128 tokens) and use them as the retrieval corpus for all models

MMLU

• Atlas trains both the retriever and the language model, which we consider a white-box retrieval LM setting.
• 对于检索增强的版本，我们将test question作为query，从Wikipedia中检索10个文档，与question拼接成10个输入，最后的结果是10个输出的聚合

Open Domain QA

• dataset: Natural Question and TriviaQA

  • For evaluation, we consider the few-shot(use a few training data) and full data(use all training data)

• RETRO, R2-D2, Atlas are finetuned on the training data, either in a few-shot setting or with full training data

Analysis

• 性能的提升不止源于聚合不同的输出结果，聚合相关的文档是成功的关键
• 随着聚合文档数目的提升，REPLUG 和 REPLUG LSR 的性能单点提升，不过 a small number of documents(e.g., 10)就可以做的不错

• REPLUG带来的性能增益与模型大小保持一致, 且能够应用到不同模型上

• REPLUG is more helpful when texts contain rare entities