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

这篇具有很好参考价值的文章主要介绍了【论文阅读】REPLUG: Retrieval-Augmented Black-Box Language Models。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

前言

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

REPLUG

【论文阅读】REPLUG: Retrieval-Augmented Black-Box Language Models,论文阅读,信息检索,论文阅读,语言模型,算法

给一个输入上下文
REPLUG会首先从外部资源 $D=\{d_1,\dots,d_m\}$ 中检索出一些相关文档
- 使用基于双塔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^{'}$ ，下一个token $y$ 的生成概率由加权平均决定
    - $p(y|x,D^{'}) = \sum_{d \in D^{'}}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: Retrieval-Augmented Black-Box Language Models,论文阅读,信息检索,论文阅读,语言模型,算法

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^{'} \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^{'}$ 上近似计算

Computing LM likelihood

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

$\mid x, y)=\frac{e^{P_{L M}(y \mid d, x) / \beta}}{\sum_{d \in \mathcal{D}^{\prime}} e^{P_{L M}(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}} K L\left(P_R(d \mid x) \| 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

【论文阅读】REPLUG: Retrieval-Augmented Black-Box Language Models,论文阅读,信息检索,论文阅读,语言模型,算法

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

MMLU

【论文阅读】REPLUG: Retrieval-Augmented Black-Box Language Models,论文阅读,信息检索,论文阅读,语言模型,算法

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

【论文阅读】REPLUG: Retrieval-Augmented Black-Box Language Models,论文阅读,信息检索,论文阅读,语言模型,算法

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