[论文笔记]Query Rewriting for Retrieval-Augmented Large Language Models

引言

今天带来论文Query Rewriting for Retrieval-Augmented Large Language Models的笔记。

本篇工作从查询重写的角度介绍了一种新的框架，即重写-检索-阅读，而不是以前的检索-阅读方式，用于检索增强的LLM。关注的是搜索查询本身的适应性，因为输入文本与检索所需的知识之间难免存在差距。

首先提示一个LLM生成查询，然后使用网络搜索引擎检索上下文。此外，为了更好地使查询与冻结的模块对齐，作者提出了一个可训练的流程方案。采用一个小型语言模型作为可训练的重写器，以适应LLM阅读器。重写器通过强化学习使用LLM阅读器的反馈进行训练。

1. 总体介绍

LLM仍然面临着幻觉问题和时间不对齐问题。这影响了LLM的可靠性，并阻碍了更广泛的实际应用。现有的研究已经证明，将外部知识(非参数化知识)与内部知识(参数化知识)结合起来，可以有效缓解幻觉问题，特别是对于知识密集型任务。事实上，检索增强的LLM已经被证明非常有效，被视为缓解原始LLM生成中的事实性缺陷的标准解决方案。检索增强方法被应用于选择相对段落作为语言模型的外部上下文，这是检索-然后阅读的框架。以开放域问答任务为例，检索器首先为一个问题搜索相关文档。然后LLM接收问题和文档，然后预测答案。

由于大多数LLM只能通过推理API进行访问，它们在流程中扮演冻结的黑盒子阅读器的角色。这使得以前需要完全访问的检索增强方法不再可行。最近的检索增强语言模型的研究更加倾向于面向LLM的适应性。现有方法忽视了查询的适应性，即检索然后阅读流程的输入。检索查询要么来自数据集中的原始数据，要么直接由黑盒子生成确定，因此一直是固定的。然而，输入文本和实际需要查询的知识之间难免存在差距。这限制了性能，并给检索能力增强和提示工程带来了负担。

考虑到这个问题，本文提出了一个新的检索增强框架——重写-检索-阅读，可以进一步调整以适应LLM。在检索器之前，添加了一个步骤来重写输入，填补给定输入和检索需求之间的差距，如图1所示。采用现成的工具，即互联网搜索引擎作为检索器，避免了搜索索引的维护，并可以获取最新的知识。重写步骤的动机是从输入文本中澄清检索需求。
作者还提出了一个可训练方案，用于重写-检索-阅读框架(图1c)。黑盒检索器和阅读器形成一个冻结系统。采用了一个小型可训练的语言模型来执行重写步骤，称为重写器。通过使用LLM性能作为奖励进行强化学习，重写器学习适应检索查询以提高下游任务中的阅读器。

总之，提出的新颖检索增强方法，重写-检索-阅读，是第一个为冻结的检索器和LLM阅读器调整输入文本的框架。引入了一个可调节的方案，使用一个小的可训练模型，在资源消耗较少的情况下实现了性能提升。

2. 相关工作

2.1 检索增强

语言模型需要外部知识来缓解事实性缺陷。检索增强被视为标准有效的解决方案。借助检索模块，相关段落被提供给语言模型作为原始输入的上下文。因此，像常识或实时新闻这样的事实信息通过上下文化的阅读理解有助于输出预测。

早期的研究使用稀疏检索器或密集检索器放在预训练语言模型(PrLM)之前。神经检索器和阅读器都是可训练大小的PrLM，如BERT或BART。因此，整个检索-阅读框架是一个可调节的端到端系统，检索到的上下文可以被视为中间结果。为了优化检索和阅读理解，提出了平滑两步框架的方法。近期的研究表明，随着模型和数据规模的迅速增加，检索仍然是一个强大的增强方法。另一方面，与大规模语言模型相比，检索增强可以弥补参数规模不足的问题。例如，通过联合训练检索器和阅读器，Atlas显示了与540B PalM相当的零样本性能，但体积小50倍。

以互联网为知识库 更相关于作者的工作，搜索引擎可以扮演检索器的角色，利用互联网作为外部知识的来源。Komeili等通过互联网搜索与对话历史相关的信息来进行对话响应生成。SeeKeR使用单个Transformer进行迭代式的搜索查询生成，然后进行知识提取以进行对话生成和句子补全。对于大规模模型，网络搜索仍然显示出有效的知识增强、事实检查和LLM代理增强。

2.2 与黑盒LLM的合作

ChatGPT、Codex、PaLM（等大型语言模型具有令人印象深刻的自然语言处理能力和可扩展性。这导致越来越多地在各种NLP任务中采用LLM。然而，在大多数情况下，LLM只能作为黑盒使用，原因是：(1)像ChatGPT这样的模型不是开源的；(2)大规模参数需要计算资源，而用户并不总是能够负担得起。这种限制意味着除了输入和输出文本之外，用户无法访问其他信息。

现有研究已经证明，通过精心设计的交互方法可以更好地利用LLM的能力。GenRead提示LLM生成上下文，而不是部署检索器，显示出LLM可以通过提示检索内部知识。ReAct和Self-Ask结合了CoT和与Web API的交互。仅依靠提示构建，ReAct为交互任务提供了新的基线。Demonstrate-Search-Predict定义了LLM和检索器之间的复杂流程。与ReAct不同，DSP除了多跳拆分和检索之外，还集成了演示引导的提示。

尽管在零样本或少样本设置中表现出有希望的性能，但LLM的行为有时需要调整。一种可行的方法是在LLM之前或之后附加可训练的小模型。这些小模型作为系统参数的一部分可以进行微调优化。RePlug提出了对冻结LLM的密集检索器进行微调，以用于检索-阅读流程。检索器在LLM的监督下进行训练，以检索适合LLM的文档。为了实现相同的目的，Directional Stimulus Prompting部署一个小模型为LLM提供刺激，该模型根据LLM的奖励进行更新。

与上述激发人的工作不同，作者提出的流程在检索-阅读模块之前包含一个查询重写步骤。进一步提出了一个可训练方案，使用一个小的重写模型，通过重构搜索查询来增强检索增强的LLM。

3 方法

作者提出了重写-检索-阅读(Rewrite-Retrieve-Read，3R)的流程，从查询重写的角度改进检索增强的LLM。图1展示了一个概述。

3.1 重写-检索-阅读

一个有检索增强的任务可以表示如下。给定一个知识密集型任务的数据集 $D = \{(x, y)_i\}，i = 0, 1, 2, . . . , N$ ，其中 $x$ (例如一个问题)是流程的输入， $y$ 是期望的输出(例如正确答案)。流程包括三个步骤：(1)查询重写：根据原始输入 $x$ 生成一个与所需知识相关的新查询 $\tilde x$ 。(2)检索：搜索相关的上下文文档。(3)阅读：理解输入和上下文 $[d o c, x]$ ，并预测输出 $\hat y$ 。

一种直接但有效的方法是让LLM重写查询以搜索可能需要的信息。作者使用少样本提示来鼓励LLM思考，输出可以是一个或多个查询进行搜索。

3.2 可训练方案

然而，完全依赖冻结的LLM也存在一些缺点。推理错误或无效的搜索会影响性能。另一方面，检索到的知识有时可能会误导和损害语言模型。为了更好地与冻结模块对齐，可以添加一个可训练的模型，并通过将LLM阅读器的反馈作为奖励来进行调整。

作者进一步提出利用一个可训练的小型语言模型来接管重写步骤，如图1右侧所示。可训练的模型以预训练的T5-large(770M)作为初始化，记为可训练重写器 $G_θ$ 。重写器首先在伪数据上进行预热训练，然后通过强化学习进行持续训练。

3.2.1 重写器预热

查询重写任务与T5等序列生成模型的预训练目标有很大的区别。首先，为查询重写任务构建了一个伪数据集。受最近的蒸馏方法的启发，提示LLM重写训练集中的原始问题 $x$ ，并收集生成的查询 $\tilde x$ 作为伪标签。然后对收集到的样本进行过滤：从LLM阅读器中获得正确预测的样本被选入预热数据集 $D_{Train}=\{(x, \tilde x)|\hat y = y\}$ 。重写器 $G_θ$ 在 $D_{Train}$ 上进行微调，使用标准对数似然作为训练目标，表示为
$\mathcal{L}_{warm} = -\sum_t \log p_\theta(\hat {\tilde x} |\tilde x_{<t},x) \tag 1$
预热后的重写器模型表现出适度的性能，这取决于伪数据的质量和重写器的能力。由于高度依赖人工编写的提示行， $\tilde x$ 可能不是最优的。重写器规模相对较小的限制也会影响预热后的性能。

3.2.2 强化学习

强化学习暂略

4. 实现

Rewriter 给LLM带少样本示例的提示去重写查询的上下文学习。提示遵循[instruction, demonstrations, input]的公式，其中输入为 $x$ 。指令是直接的，示例是来自训练集的1-3个随机示例，并在所有运行中保持不变，主要用于任务特定的输出格式说明，即对于HotpotQA的简短短语作为答案，对于MMLU的选项作为答案。对于第3.2节中的训练方案，将T5作为重写器进行微调。

Retriever 使用必应搜索引擎作为检索器。它不需要像密集检索器那样构建候选索引，也不需要像教科书(textbook)那样的候选项。但它允许广泛的知识范围和最新的事实性。通过Bing API，检索分为两种方法进行。(1)对于所有检索到的网页，连接由Bing选择的相关句子片段。这种方法类似于在浏览器中使用搜索引擎，输入查询并按Enter键，然后收集搜索结果页面上显示的文本。(2)对于检索到的网页，请求URL和解析器获取所有文本。这类似于在搜索结果页面上点击。然后使用BM25保留与查询具有更高相关性分数的文档，缩短文档长度。

Reader 阅读器是一个冻结的LLM，采用ChatGPT和Vicuna-13B。它执行阅读理解和预测，具有少量的上下文学习。在作者的提示中，遵循简要的指导和演示，输入为 $x$ 或 $[doc，\hat{\tilde x}]$ 。

5. 实验

在开放领域问答方面的实验结果见表2。对于三个数据集，查询重写始终能够通过冻结的重写器和可训练的重写器带来性能提升。

使用复杂问题作为查询无法弥补参数化知识，而是带来了噪音。表明多跳问题不适合作为网络搜索引擎的查询。通过添加重写步骤，分数有所提高。

6. 分析

6.1 训练过程

训练过程包括两个阶段，预热和强化学习。图2展示了在强化学习过程中的训练迭代中的度量分数。由于重写模型在RL之前已经在伪数据上进行了预热，因此“0迭代”处的分数表示从预热训练中获得的能力。可以观察到曲线在所有数据集上都呈上升趋势，并伴有一些波动。

6.2 检索结果

作者提出的方法是一个流水线框架，而不是端到端系统。查询重写首先影响检索到的上下文，然后上下文对阅读器的输出产生影响。因此，问答指标是间接测量。通过检索指标“命中率”更仔细地观察检索到的上下文和阅读器的能力。计算命中率以衡量检索到的上下文是否包含正确答案。

表4显示了在AmbigNQ上的分数。第二行的分数是在选定的样本中计算的，这些样本的检索上下文命中了正确答案（在标准检索-然后阅读设置下）。这些分数显示了具有检索增强的阅读器的近似上限能力，简称为“上限”分数。与无检索设置（第一行）相比，检索的有效性得到了证明。

指标显示，使用BM25进行内容选择比片段召回更好的文档，而查询重写在两种设置上都取得了进展。检索器的命中率提高比阅读器的提高更为显著。

6.3 案例研究

图3：直观说明的示例。Q0表示原始输入，Q1来自LLM重写器，Q2来自经过训练的T5重写器。Hit表示检索器召回答案，而Correct表示阅读器的输出。

我们看图3中的示例，比较了原始问题和查询。在示例1中，原始问题询问Lady Mary-Gaye Curzon的最小女儿与两位特定演员合作出演的电影。查询1和查询2都将关键词电影放在前面，紧随Lady Mary-Gaye Curzon的最小女儿。通过两者，可以检索到女演员Charlotte Calthorpe及其电影信息，并包含了答案。第二个示例是LLM重写器生成的查询失败，但T5生成的查询获得了正确答案。查询1中的数字2000被误解，而查询2将200和电影放在一起，避免了无意义的检索。示例3是关于多项选择的案例。查询简化了背景并增强了关键词社区规划者。检索到的上下文主要是关于社区规划介绍，其中答案环境出现了多次。