前言

• 原文地址：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，并重复上述过程