REPLUG: Retrieval-Augmented Black-Box Language Models

问题 (Problem)

大型语言模型（LLMs）如GPT-3虽然强大，但存在两个核心问题：

1. 知识局限性：模型参数中存储的知识是静态的，无法实时更新，且对于长尾知识（rare knowledge）的覆盖不全，容易产生事实性错误或“幻觉”。
2. 黑盒特性：当前最先进的LLMs（通常 >100B 参数）往往通过API提供服务，用户无法访问模型的内部参数、梯度或进行微调。这使得传统的、需要“白盒”访问权限的检索增强方法（如RETRO、Atlas）无法适用。

因此，本文的核心问题是：如何在只能“黑盒”访问（即只能输入文本、获取输出）的前提下，通过外部知识库对大型语言模型进行有效的检索增强，以提升其性能并减少幻觉？

方法 (Method)

作者提出了 REPLUG (Retrieve and Plug) 框架，它将语言模型视为一个不可更改的黑盒，并将一个可调优的检索器作为插件来增强它。

REPLUG 推理过程

REPLUG的推理过程分为两步：文档检索 和 输入重构与集成。

REPLUG 推理流程图（论文Figure 2）

1. 文档检索 (Document Retrieval)

   • 给定一个输入上下文 $x$，使用一个**密集检索器（Dense Retriever）**从一个大规模的外部语料库 $\mathcal{D}$ 中检索出与 $x$ 最相关的 $k$ 个文档。
   • 该检索器采用双编码器架构，将输入 $x$ 和文档 $d$ 分别编码为向量 $E(x)$ 和 $E(d)$。
   • 相关性得分通过余弦相似度计算：

     $$s(d,x) = \cos(E(d), E(x))$$

   • 为了效率，所有文档的向量会预先计算并用FAISS等工具建立索引，以实现快速检索。

2. 输入重构与集成 (Input Reformulation & Ensemble)

   • 直接将所有 $k$ 个文档拼接到输入 $x$ 前面会受到模型上下文窗口长度的限制。
   • REPLUG 采用了一种巧妙的集成策略：
     • 对于检索到的top-k个文档 ${d_1, d_2, ..., d_k}$，分别与原始输入 $x$ 进行拼接，构造出 $k$ 个独立的输入：${d_1 \circ x, d_2 \circ x, ..., d_k \circ x}$。
     • 将这 $k$ 个输入并行地送入黑盒语言模型中，得到 $k$ 组关于下一个词的输出概率分布。
     • 最终的输出概率是对这 $k$ 组概率分布进行加权平均的结果。权重由检索器的相关性得分决定。
   • 下一个词 $y$ 的最终概率计算公式为：

   $$p(y|x, \mathcal{D}') = \sum_{d \in \mathcal{D}'} p(y|d \circ x) \cdot \lambda(d, x)$$

   其中，$\mathcal{D}'$ 是 top-k 文档集合，$d \circ x$ 表示文档和输入的拼接。权重 $\lambda(d,x)$ 是检索得分经过Softmax归一化后的结果：

   $$\lambda(d,x) = \frac{e^{s(d,x)}}{\sum_{d' \in \mathcal{D}'} e^{s(d',x)}}$$

   • 这种方法的优势是，虽然需要调用模型 $k$ 次，但它绕过了上下文长度限制，可以集成更多的文档信息，并且计算可以完全并行，只是增加了API调用成本。

REPLUG LSR: 训练可适变的检索器

为了让检索器更能“投其所好”，即检索到能最大化帮助特定黑盒LLM的文档，作者提出了 REPLUG LSR (LM-Supervised Retrieval) 训练方案。其核心思想是用黑盒LLM本身作为监督信号来指导检索器的训练。

REPLUG LSR 训练流程图（论文Figure 3）

训练过程包含四个步骤：

1. 计算检索似然度 (Computing Retriever Likelihood, $P_R$)

   • 对于一个输入 $x$，用当前检索器检索 top-k 个文档 $\mathcal{D}'$。
   • 基于检索器的相似度得分 $s(d,x)$，计算出一个在 $\mathcal{D}'$ 上的概率分布 $P_R(d|x)$，代表了检索器认为每个文档的重要程度。
   • 公式为：

   $$P_R(d|x) = \frac{e^{s(d,x)/\tau}}{\sum_{d' \in \mathcal{D}'} e^{s(d',x)/\tau}}$$

   其中 $\tau$ 是温度超参数。

2. 计算语言模型似然度 (Computing LM Likelihood, $Q_{LM}$)

   • 这是获取监督信号的关键一步。
   • 对于 $\mathcal{D}'$ 中的每一个文档 $d$，将其与输入 $x$ 拼接后，喂给黑盒LLM，计算模型对于真实答案（ground truth） $y$ 的预测概率 $P_{LM}(y|d,x)$。
   • 这个概率值越高，说明文档 $d$ 对LLM做出正确预测的帮助越大。
   • 将这些概率值也转换成一个目标分布 $Q_{LM}(d|x,y)$，代表了LLM认为每个文档的重要程度。
   • 公式为：

   $$Q_{LM}(d|x, y) = \frac{e^{P_{LM}(y|d,x)/\beta}}{\sum_{d' \in \mathcal{D}'} e^{P_{LM}(y|d',x)/\beta}}$$

   其中 $\beta$ 是另一个温度超参数。

3. 优化损失函数 (Loss Function)

   • 通过最小化检索器分布 $P_R$ 和语言模型目标分布 $Q_{LM}$ 之间的 KL散度 (KL Divergence) 来更新检索器的参数。
   • 这会迫使检索器学习去预测那些能够最大化LLM性能的文档。
   • 损失函数为：

     $$\mathcal{L} = \frac{1}{|\mathcal{B}|}\sum_{x \in \mathcal{B}}KL(P_R(d|x) \ || \ Q_{LM}(d|x,y))$$

     其中 $\mathcal{B}$ 是一组输入上下文
   • 在训练过程中，只有检索器的参数被更新，LLM始终保持冻结和黑盒。

4. 异步更新数据索引 (Asynchronous Update of Datastore Index)

   • 由于检索器的编码器在训练中不断变化，预先计算的文档向量会变得“过时”。
   • 因此，训练过程中需要周期性地（例如每T个训练步）使用更新后的编码器重新计算整个文档库的向量，并重建FAISS索引。这是一个计算量巨大的步骤。

Baseline

论文在多个任务上与不同的基线模型进行了对比：

• 语言建模 (Language Modeling)：

  • 基线: 原始的、未增强的GPT-2和GPT-3系列模型。
  • 对比模型: REPLUG增强版和REPLUG LSR增强版。

• MMLU (多任务语言理解)：

  • 基线LLMs: Codex (175B), PaLM (540B), Flan-PaLM (540B)。
  • 白盒检索模型: Atlas (11B)。

• 开放域问答 (Open-Domain QA)：

  • 基线LLMs (few-shot): Chinchilla, PaLM, Codex。
  • 白盒检索模型 (fine-tuned): RETRO, R2-D2, Atlas。

数据集 (Datasets)

• REPLUG LSR 训练:

  • 训练查询: 从 The Pile 训练集中采样的80万个序列。
  • 外部文档库: 从 The Pile 训练集中采样的3600万个文档（与查询无重叠）。

• 评估:

  • 语言建模: The Pile 基准测试集。
  • 多任务语言理解: MMLU 数据集，采用5-shot设置。
  • 开放域问答: Natural Questions (NQ) 和 TriviaQA (TQA) 数据集，采用few-shot设置。
  • 泛化性分析: Wikitext-103 数据集，用于测试REPLUG在GPT-2, BLOOM, OPT等不同模型家族上的效果。

可复现性 (Reproducibility)

• 代码: 论文中未提供公开的代码库链接。
• 算力: 复现该方法需要巨大的计算资源。
  1. API成本: REPLUG LSR的训练需要大量调用一个中等大小的LLM（如GPT-3 Curie）作为监督，推理过程则需要将API调用成本乘以 $k$（集成文档的数量）。
  2. GPU资源: 训练一个强大的密集检索器本身就需要大量GPU算力。
  3. 索引更新瓶颈: REPLUG LSR训练中最耗费计算的步骤是异步索引更新。周期性地为数千万文档重新编码并建立索引，需要非常强大的计算集群，是个人或小团队难以承受的。

可改进的几个点 (Potential Improvements)

1. 推理效率: 推理成本是原始模型的 $k$ 倍，这是一个主要缺点。未来的研究可以探索更高效的集成方法，例如，如何在单次前向传播中融合多个文档的信息，或者用更小的模型来预处理和压缩检索到的信息。
2. 训练效率: 异步索引更新是训练的巨大瓶颈。可以研究更高效的索引更新策略，如增量更新，或者探索不需要静态索引的检索方法。
3. 集成权重: 当前的集成权重 $\lambda(d,x)$ 仅依赖于检索器的初始分数。可以设计更动态的权重方案，比如结合LLM对每个文档的“置信度”进行二次加权。
4. 可解释性: 论文作者也提到，模型何时依赖参数知识、何时依赖检索知识是不透明的。增强模型的可解释性，让用户知道答案的来源，是未来一个重要的研究方向。
5. 检索粒度: REPLUG检索的是整个文档。可以探索更细粒度的检索，如段落、句子甚至短语，这可能为模型提供更精确、噪声更少的信息。

可以被引用的一些结论 (Citable Conclusions)

1. 黑盒检索增强是可行且有效的：本文首次证明，在不访问模型内部参数的情况下，通过外部检索能显著提升超大规模（>100B）黑盒语言模型的性能。
2. 语言模型自身是最好的监督者：REPLUG LSR的核心洞见是，可以用LLM本身来指导检索器的训练，让检索器学会去寻找对该LLM最有帮助的知识，实现了检索器和语言模型的“适配”。
3. 在多样的模型和任务上表现稳定：REPLUG框架具有普适性，能够稳定地提升包括GPT系列、BLOOM、OPT在内的多种模型家族在语言建模、知识问答和复杂推理等多个任务上的表现。
4. 显著提升少样本学习能力：在开放域问答的少样本（few-shot）设定下，REPLUG LSR增强的Codex模型甚至超越了使用更多样本进行微调的先进白盒模型（如Atlas），创造了新的SOTA记录。
5. 检索对于长尾知识至关重要：定性分析表明，REPLUG对于包含**稀有实体（rare entities）**的文本尤其有效，这证明了检索在弥补模型参数化知识短板方面的重要作用。
6. 用更小的模型达到更优的效果：通过REPLUG增强，一个175B参数的模型（Codex）在MMLU任务上的表现可以媲美一个三倍大于它、且经过指令微调的模型（540B的Flan-PaLM），展示了检索在提升模型效率上的巨大潜力。