目录
摘要
这篇论文提出了一种名为"Rewrite-Retrieve-Read"的新型检索增强框架,旨在改进大型语言模型(LLMs)在知识密集型任务中的表现。传统方法采用"检索-读取"的两步流程,而本文创新性地在检索前增加了查询重写步骤,形成"重写-检索-读取"的三步流程。核心思想是通过优化检索查询本身来弥合输入文本与所需知识之间的差距,而非像以往研究那样专注于调整检索器或读取器。工作流程分为三个关键阶段:首先利用LLM或可训练的小型语言模型(称为"重写器")对原始查询进行重写;然后使用网络搜索引擎检索相关文档;最后将重写后的查询和检索到的文档一起输入冻结的LLM读取器生成最终答案。为进一步优化流程,作者提出了一种可训练方案,采用T5-large作为重写器,通过强化学习根据LLM读取器的反馈进行训练。实验在开放域问答(HotpotQA、AmbigNQ、PopQA)和多项选择问答(MMLU)任务上进行,使用ChatGPT和Vicuna-13B作为读取器,结果显示查询重写能持续提升性能,验证了该框架的有效性和可扩展性。
Abstract
This paper introduces a new retrieval-enhanced framework called “Rewrite-Retrieve-Read,” aimed at improving the performance of large language models (LLMs) in knowledge-intensive tasks. Traditional methods use a two-step “retrieve-read” process, but this paper innovatively adds a query rewriting step before retrieval, forming a three-step “rewrite-retrieve-read” process. The core idea is to bridge the gap between the input text and the required knowledge by optimizing the retrieval query itself, rather than focusing on adjusting the retriever or reader as previous studies have done. The workflow consists of three key stages: first, using an LLM or a trainable small language model (called the “rewriter”) to rewrite the original query; then using a web search engine to retrieve relevant documents; and finally feeding the rewritten query and retrieved documents into a frozen LLM reader to generate the final answer. To further optimize the process, the authors propose a trainable approach, using T5-large as the rewriter and training it through reinforcement learning based on feedback from the LLM reader. Experiments were conducted on open-domain question answering (HotpotQA, AmbigNQ, PopQA) and multiple-choice question answering (MMLU) tasks, using ChatGPT and Vicuna-13B as readers. The results show that query rewriting consistently improves performance, validating the effectiveness and scalability of the framework.
用于检索-增强大语言模型的查询与重写
Title: Query Rewriting for Retrieval-Augmented Large Language Models
Author: Xinbei Ma, Yeyun Gong, Pengcheng He, Hai Zhao, Nan Duan
Source: Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing
Arxiv: https://arxiv.org/abs/2305.14283
研究背景
大型语言模型虽然在自然语言处理方面展现出惊人能力,但仍面临幻觉和时间错位等问题,这主要源于其对大规模高质量语料库的依赖而缺乏对现实世界的直接感知。检索增强方法通过结合外部知识(非参数知识)和内部知识(参数知识)来缓解这些问题,已成为提高LLM事实性的标准解决方案。传统的"检索-读取"框架中,检索器首先搜索与问题相关的文档,然后LLM接收问题和文档并预测答案。
然而,现有方法存在明显局限。大多数LLM只能通过推理API访问,在流程中充当黑盒冻结读取器的角色,这使得需要完全访问模型参数的早期检索增强方法不再可行。近期研究多集中于LLM导向的适配,如训练密集检索模型以适应冻结的语言模型,或设计检索器与读取器之间的交互方式。但这些方法都忽视了查询本身的适配问题——检索查询要么直接来自数据集,要么由黑盒生成决定,始终固定不变,导致输入文本与真正需要查询的知识之间存在不可避免的差距。
黑盒API的限制
现在很多强大的LLM(如ChatGPT)只提供"输入-输出"的API接口,研究者无法直接调整模型内部参数。这导致早期那些需要修改LLM参数的检索增强方法(比如让LLM在训练时学习如何利用检索结果)无法继续使用——因为现在LLM对我们来说是个"黑盒子",只能调用不能修改。
对于现有解决方案的局限性,研究者们最近主要尝试两种替代方案:
训练专用检索器:专门训练一个检索模型,让它学会为特定LLM提供最合适的外部文档
优化交互方式:设计更聪明的提示词(prompt)让LLM更好地使用检索结果
但现有方法都忽略了一个根本问题:检索查询本身的质量。
目前要么:直接使用原始问题作为检索词(比如把"量子纠缠是什么?"直接扔给搜索引擎)或者让黑盒LLM自己决定检索词(但无法优化这个过程)
这会导致"你说的话"和"系统真正需要查的内容"之间存在偏差。就像你想问"如何预防感冒",但直接搜索这句话可能不如搜索"增强免疫力 维生素C 临床研究"来得有效。
举个具体例子:
假设问题是:“《星际穿越》中那个后来成为科学家的女孩最后怎样了?”直接检索可能效果差,因为包含太多电影细节。
理想的重写可能是:“星际穿越 Murph角色结局 NASA科学家”
但现有系统缺乏主动优化查询的这个环节。
这种"问题表达"与"实际需要查找的知识"之间的不匹配,就是作者指出的"不可避免的差距"。而本文的创新点正是通过增加"查询重写"步骤来填补这个空白。
这一研究空白正是本文工作的出发点。作者观察到,现有检索增强方法过分关注检索器或读取器的调整,而忽略了查询优化这一关键环节。特别是在处理复杂查询(如多跳问题)时,直接使用原始问题作为检索查询往往效果不佳,因为这些问题可能包含冗余信息或缺乏关键检索词。此外,随着LLM规模不断扩大,如何高效地将小型可训练模块与黑盒LLM集成也成为一个重要研究问题。
方法论
上图直观对比了传统的 retrieve-then-read 方法和本文提出的 rewrite-retrieve-read 框架,并展示了可训练重写器的方案。
(a)标准 Retrieve-Then-Read 流程
流程步骤:
输入(Input):原始问题或文本(如 “What is the capital of France?”)。
检索(Retrieve):直接使用原始输入作为查询,从外部知识库(如维基百科)或搜索引擎(如Bing)检索相关文档。
读取(Read):将检索到的文档 + 原始问题一起输入LLM(如ChatGPT),生成最终答案(如 “Paris”)。 核心问题:
查询未优化:直接使用原始输入检索,可能效果不佳(例如,复杂问题、模糊表述或包含冗余信息)。
(b) LLM作为冻结重写器的 Rewrite-Retrieve-Read
流程步骤:
输入(Input):原始问题。
重写(Rewrite):使用一个冻结的LLM(如ChatGPT)对原始问题重写,生成更适合检索的查询。
检索(Retrieve):用重写后的查询检索文档。
读取(Read):将检索结果 + 原始问题输入LLM生成答案。
改进点: 查询适配性增强,重写后的查询更精准,检索效果更好,直接利用现成LLM的重写能力。
局限性: 依赖LLM的提示工程(prompt design),重写质量不稳定,并且黑盒LLM无法针对性优化。
(c)可训练重写器的 Rewrite-Retrieve-Read
流程步骤:
输入(Input):原始问题。
重写(Rewrite):使用一个可训练的小型LM(如T5)作为重写器,生成优化查询。
训练方法:
预热训练(Warm-up):用LLM生成的伪数据(如正确重写样本)监督训练。
强化学习(RL):以LLM的答案质量作为奖励,进一步优化重写器(PPO算法)。
检索(Retrieve):用训练后的重写器生成查询并检索。
读取(Read):LLM生成最终答案。
改进点: 通过强化学习,使重写器适配特定任务和LLM;小型重写器比直接调用LLM更轻量。
关键设计: 奖励函数(Reward)基于LLM输出的答案质量(如EM、F1分数)和检索命中率; KL散度正则化防止重写器偏离初始语言模型太远。
本文提出的"重写-检索-读取"框架包含三个核心组成部分:查询重写、检索和读取。在技术实现上,作者探索了两种主要方案:基于冻结LLM的重写方案和基于可训练重写器的方案。
基于冻结LLM的重写方案
基于冻结LLM的重写方案采用少量示例提示(few-shot prompting)的方式,引导LLM对原始查询进行重写。具体而言,提示模板包含指令、示范和输入三部分,指令简明扼要,示范则来自训练集中的1-3个随机示例,主要用于说明任务特定的输出格式。这种方法无需额外训练,直接利用LLM的推理能力生成更适合检索的查询。例如,对于开放域问答任务,提示LLM逐步思考并生成所需知识的搜索引擎查询;对于多项选择任务,则提示生成能更好回答给定问题的网页搜索查询。
基于可训练重写器的方案
更具创新性的是可训练重写器方案。作者采用T5-large(770M参数)作为基础模型,通过两阶段过程进行训练:预热训练和强化学习。
重写器预热训练(Rewriter Warm-up)
在预热阶段,首先构建伪数据集——使用LLM重写训练集中的原始问题,并筛选那些能使LLM读取器做出正确预测的样本作为训练数据。重写器通过标准对数似然目标进行监督训练,使其初步掌握查询重写的基本模式。
目标:让重写器初步学会生成合理的查询。
伪数据生成:
- 用LLM(如ChatGPT)重写训练集中的原始问题 x x x,生成候选查询 x ~ \tilde{x} x~。
- 筛选能使LLM读者输出正确答案的样本,构成预热数据集:
D T r a i n = { ( x , x ~ ) ∣ y ^ = y } D_{Train} = \{(x, \tilde{x}) \mid \hat{y} = y\} DTrain={(x,x~)∣y^=y}
训练目标(最大似然估计):
L w a r m = − ∑ t log p θ ( x ~ t ∣ x ~ < t , x ) \mathcal{L}_{warm} = -\sum_{t} \log p_{\theta}(\tilde{x}_t \mid \tilde{x}_{<t}, x) Lwarm=−t∑logpθ(x~t∣x~<t,x)- x ~ t \tilde{x}_t x~t:查询的第 t t t个token
- x ~ < t \tilde{x}_{<t} x~<t:已生成的前 t − 1 t-1 t−1个token
- θ \theta θ:重写器参数
每一项 log p θ ( x ~ t ∣ x ~ < t , x ) \log p_{\theta}(\tilde{x}_t \mid \tilde{x}_{<t}, x) logpθ(x~t∣x~<t,x) 表示在给定原始问题 x x x 和已生成部分 x ~ < t \tilde{x}_{<t} x~<t 的条件下,模型对下一个正确token x ~ t \tilde{x}_t x~t 的预测概率。最小化NLL等价于让模型对真实token分配更高概率。
查询重写器模块:
class QueryRewriter:
def __init__(self, model_name="t5-large"):
self.model = T5ForConditionalGeneration.from_pretrained(model_name)
self.tokenizer = T5Tokenizer.from_pretrained(model_name)
def warm_up_train(self, dataset):
# dataset: List[{"original": str, "rewritten": str}]
optimizer = torch.optim.Adam(self.model.parameters(), lr=3e-5)
for batch in dataloader:
inputs = self.tokenizer(
batch["original"],
return_tensors="pt",
padding=True
)
labels = self.tokenizer(
batch["rewritten"],
return_tensors="pt",
padding=True
).input_ids
outputs = self.model(
input_ids=inputs.input_ids,
attention_mask=inputs.attention_mask,
labels=labels
)
loss = outputs.loss
loss.backward()
optimizer.step()
def rewrite(self, query):
inputs = self.tokenizer(query, return_tensors="pt")
outputs = self.model.generate(**inputs)
return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
强化学习(Reinforcement Learning)
强化学习阶段则采用近端策略优化(PPO)算法进一步调整重写器。将重写器优化问题建模为马尔可夫决策过程,其中状态空间由词汇表和序列长度限定,动作空间等同于词汇表,转移概率由策略网络(即重写器模型)决定。奖励函数基于LLM读取器的预测质量设计,包含精确匹配(EM)、F1分数和命中指示器等指标,同时加入KL散度正则化防止模型偏离初始策略太远。价值网络从策略网络初始化,采用广义优势估计(GAE)计算优势函数。通过这种设计,重写器能够学习生成那些能引导检索器找到更有用文档、进而帮助LLM读取器做出更准确预测的查询。
目标:通过LLM读者的反馈进一步优化重写器。
马尔可夫决策过程(MDP)建模:
- 状态 s t s_t st:当前生成的部分查询 [ x , x ~ ^ < t ] [x, \hat{\tilde{x}}_{<t}] [x,x~^<t]
- 动作 a t a_t at:选择下一个token x ~ ^ t \hat{\tilde{x}}_t x~^t
- 策略 π θ \pi_{\theta} πθ:重写器模型 G θ G_{\theta} Gθ
- 奖励 R R R:基于LLM读者的答案质量
奖励函数设计:
R ( s t , a t ) = R l m ( x ~ ^ , y ) − β KL ( π θ ∥ π 0 ) R(s_t, a_t) = R_{lm}(\hat{\tilde{x}}, y) - \beta \text{KL}(\pi_{\theta} \| \pi_{0}) R(st,at)=Rlm(x~^,y)−βKL(πθ∥π0)- R l m R_{lm} Rlm:LLM读者的表现(如EM/F1分数)
- KL \text{KL} KL项:防止重写器偏离初始策略 π 0 \pi_{0} π0(预热后的模型)
- β \beta β:平衡系数(动态调整)
策略优化(PPO算法):
- 优化目标:
max θ E x ~ ^ ∼ p θ ( ⋅ ∣ x ) [ R ( x , x ~ ^ ) ] \max_{\theta} \mathbb{E}_{\hat{\tilde{x}} \sim p_{\theta}(\cdot|x)}[R(x, \hat{\tilde{x}})] θmaxEx~^∼pθ(⋅∣x)[R(x,x~^)] - 具体损失函数:
L θ = − 1 ∣ S ∣ T ∑ τ ∈ S ∑ t = 0 T min ( k t , θ A θ ′ , clip ( k t , θ , 1 − ϵ , 1 + ϵ ) A θ ′ ) \mathcal{L}_{\theta} = -\frac{1}{|\mathcal{S}|T} \sum_{\tau \in \mathcal{S}} \sum_{t=0}^{T} \min(k_{t,\theta} A^{\theta'}, \text{clip}(k_{t,\theta}, 1-\epsilon,1+\epsilon) A^{\theta'}) Lθ=−∣S∣T1τ∈S∑t=0∑Tmin(kt,θAθ′,clip(kt,θ,1−ϵ,1+ϵ)Aθ′)
k t , θ = p θ ( a t ∣ s t ) p θ ′ ( a t ∣ s t ) k_{t,\theta} = \frac{p_{\theta}(a_t \mid s_t)}{p_{\theta'}(a_t \mid s_t)} kt,θ=pθ′(at∣st)pθ(at∣st)- A θ ′ A^{\theta'} Aθ′:优势函数(GAE估计)
- θ ′ \theta' θ′:旧策略参数(采样时固定)
- clip \text{clip} clip:防止策略更新过大
- 优化目标:
价值函数训练:
L ϕ = 1 ∣ S ∣ T ∑ τ ∈ S ∑ t = 0 T ( V ϕ ( s t ) − R t ) 2 \mathcal{L}_{\phi} = \frac{1}{|\mathcal{S}|T} \sum_{\tau \in \mathcal{S}} \sum_{t=0}^{T} (V_{\phi}(s_t) - R_t)^2 Lϕ=∣S∣T1τ∈S∑t=0∑T(Vϕ(st)−Rt)2- V ϕ V_{\phi} Vϕ:价值网络(从策略网络初始化)
- R t R_t Rt:累计奖励
总损失函数:
L p p o = L θ + λ v L ϕ \mathcal{L}_{ppo} = \mathcal{L}_{\theta} + \lambda_v \mathcal{L}_{\phi} Lppo=Lθ+λvLϕ- λ v \lambda_v λv:价值损失权重
核心创新点:
- 两阶段训练:预热(模仿学习)→ RL(精细优化)
- KL正则化:避免重写器生成不合理查询
- 轻量化设计:小型重写器(T5-large)适配黑盒LLM
class RLTraining:
def __init__(self, rewriter, llm_reader):
self.rewriter = rewriter
self.llm_reader = llm_reader # Black-box LLM (e.g. ChatGPT API)
ppo_config = PPOConfig(
batch_size=16,
learning_rate=2e-6,
kl_divergence_coeff=0.2
)
self.ppo_trainer = PPOTrainer(
model=self.rewriter.model,
config=ppo_config
)
def compute_reward(self, queries, answers, gold_answers):
# 实现论文中的奖励函数
em_reward = exact_match(answers, gold_answers)
f1_reward = f1_score(answers, gold_answers)
hit_reward = 1 if answer_in_retrieved_docs() else -1
return em_reward + 1.0*f1_reward + 1.0*hit_reward
def train_step(self, queries, gold_answers):
# 生成重写查询
rewritten_queries = [self.rewriter.rewrite(q) for q in queries]
# 检索文档
retrieved_docs = bing_search(rewritten_queries)
# LLM生成答案
llm_inputs = [f"{doc} {q}" for doc,q in zip(retrieved_docs, queries)]
answers = chatgpt_api(llm_inputs)
# 计算奖励
rewards = self.compute_reward(rewritten_queries, answers, gold_answers)
# PPO更新
self.ppo_trainer.step(queries, rewritten_queries, rewards)
在检索环节,作者使用Bing搜索引擎作为检索器,避免了构建和维护搜索索引的开销,并能访问最新知识。检索采用两种方式:一是拼接搜索引擎返回的相关片段;二是获取网页URL并解析全部文本,然后使用BM25算法保留与查询相关性较高的部分以减少文档长度。读取环节则采用冻结的LLM(如ChatGPT或Vicuna-13B)进行阅读理解,同样使用少量示例提示的方式进行预测。
创新性
本文的核心创新点在于首次提出了从查询重写角度改进检索增强LLM的框架,具有多重理论和方法创新价值。
最显著的贡献是突破了传统"检索-读取"范式的局限,开创性地增加了查询重写环节,形成了"重写-检索-读取"的新范式。这一创新源于对检索过程中"查询-知识"差距的深刻洞察——原始输入往往不是最优的检索查询,特别是在处理复杂、模糊或专业性强的任务时。通过引入重写环节,系统能够主动调整查询形式,更精准地表达信息需求,从而显著提升后续检索和读取的效果。
在技术层面,作者提出了一个精巧的可训练方案,将小型语言模型作为重写器与黑盒LLM读取器协同工作。这一设计既保持了LLM的强大能力,又通过小型可训练模块实现了流程优化,计算效率高且实用性强。特别是采用强化学习将重写器与LLM读取器的表现直接关联,创造性地解决了黑盒模型难以端到端训练的挑战。这种"小模型引导大模型"的思路为整合可训练模块与商业LLM API提供了新范式。
实验设计上也体现了创新性。作者不仅比较了传统检索与查询重写的效果,还深入分析了不同配置下的性能差异,包括使用不同检索方法(片段vs BM25)、不同读取器(ChatGPT vs Vicuna)等。特别有价值的是对重写查询的案例分析,直观展示了查询优化如何影响检索结果和最终答案质量。例如,在多跳问题中,好的重写能够聚焦关键实体和关系;在模糊问题中,重写可以添加明确的约束条件;在专业性问题中,重写可以突出术语和背景。
此外,使用网络搜索引擎作为检索器也是一个务实而创新的选择。相比专业构建的知识库,这种方法知识覆盖面广、实时性强,能够缓解LLM的时间错位问题。作者也坦诚讨论了这一选择的利弊,为后续研究提供了宝贵经验。
实验结果
论文在开放域问答和多项选择问答两大类任务上进行了全面实验验证,涵盖了多种难度和类型的问题。
在开放域问答任务中,HotpotQA包含需要多跳推理的复杂问题,AmbigNQ提供自然问题的消歧版本,PopQA则包含更多长尾分布知识。实验结果显示,查询重写在所有数据集上均带来性能提升。特别值得注意的是,在HotpotQA上标准检索反而损害了读取器性能,说明复杂问题直接作为查询会引入噪声;而增加重写步骤后性能显著提高,证明重写能有效优化多跳问题的检索效果。在AmbigNQ和PopQA上,标准检索已有所帮助,而查询重写能进一步提升性能,其中LLM重写器表现最佳,训练后的T5重写器次之但优于标准检索。
多项选择问答任务使用MMLU数据集,涵盖人文、STEM、社会科学和其他四大类考试题目。实验发现,使用ChatGPT作为读取器时,查询重写在大多数类别上提高了分数(社会科学除外);使用Vicuna作为读取器时,查询重写在所有类别上都获得了更大提升。这表明更强大的读取器拥有更多参数记忆,更需要通过优化检索来补充外部知识。
作者还对训练过程和检索结果进行了深入分析。强化学习曲线显示,随着训练迭代,所有数据集的指标总体呈上升趋势,说明RL阶段能够弥补预热训练的不足。检索分析则表明,查询重写显著提高了检索命中率(AmbigNQ上从76.4%提升到82.2%),且检索效果的改善幅度大于读取器的提升幅度,这与相关研究发现一致。
案例研究生动展示了查询重写如何改变检索结果和预测性能。一个好的重写能够:(1)重组问题结构,将关键信息前置;(2)消除歧义,避免错误解读;(3)简化复杂背景,突出核心概念。例如,在多跳问题中,好的重写会聚焦关键人物关系;在模糊问题中,重写可以修正时间表述;在专业问题中,重写会强调术语和背景。
局限性
尽管取得了显著成果,作者也坦诚指出了工作的几点局限性。
首先是通用性与专用性之间的权衡问题。增加训练过程虽然提高了特定任务性能,但相比少量示例提示的直接迁移能力有所下降。这意味着该方法在需要快速适应新任务的场景中可能不如纯提示方法灵活。
其次,与新兴的"LLM智能体"研究方向相比,本文框架采用单轮检索策略,而智能体方法通常需要多次调用LLM进行多轮检索和交互。虽然单轮设计效率更高,但在处理极其复杂的问题时可能不如多轮方法全面。作者明确表示其动机是改进单轮检索-读取框架,而非探索更复杂的交互模式。
第三,使用网络搜索引擎作为检索器虽然便利,但也带来一些挑战。商业API需要订阅,且网络知识难以控制,可能返回不一致、冗余甚至有害内容,影响LLM读取器。相比之下,基于专业筛选知识库的神经密集检索器可能获得更可控的检索结果。
此外,还有一些潜在限制未在文中明确讨论但值得思考。例如,查询重写质量高度依赖预热阶段生成的伪数据质量;强化学习训练过程计算成本较高;不同领域和任务可能需要设计特定的奖励函数;框架对搜索引擎API的依赖可能影响可复现性等。
总结
《Query Rewriting for Retrieval-Augmented Large Language Models》这篇论文提出了一种创新且实用的检索增强框架,通过引入查询重写步骤显著提升了大型语言模型在知识密集型任务中的表现。其核心价值在于突破了传统方法仅关注检索器或读取器优化的局限,开创性地从查询适配角度解决问题,为检索增强LLM研究提供了新视角。方法论上,论文的亮点在于巧妙结合了冻结LLM的强大能力和小型可训练模块的灵活性,特别是通过强化学习将重写器与黑盒LLM的反馈相连接,为解决类似集成问题提供了范例。实验设计全面,结果分析深入,不仅验证了框架的有效性,还揭示了不同配置下的性能特点,为后续研究提供了丰富参考。从更广的视角看,这项工作反映了LLM研究的一个重要趋势:如何在保持商业LLM黑盒特性的同时,通过可训练模块和外部工具扩展其能力边界。查询重写的思想不仅适用于检索增强场景,也可推广到其他工具使用场景,如代码解释器、专业模型调用等,具有广泛的应用潜力。当然,如同作者指出的,该方法仍有改进空间,特别是在平衡通用性与专用性、处理复杂多轮交互、提高检索结果可控性等方面。未来研究可以探索更智能的多轮重写策略、结合专业知识库的混合检索方式,以及更高效的训练算法等方向。总的来说,这篇论文在检索增强语言模型领域做出了实质性贡献,提出的"重写-检索-读取"框架和可训练方案既有理论创新又具实用价值,为提升LLM的事实性和可靠性提供了有效工具,对相关研究和应用都具有重要启发意义。