Q-RAG：通过基于价值的嵌入器训练实现长上下文多步检索 Q-RAG: Long Context Multi-step Retrieval via Value-based Embedder Training
用强化学习微调嵌入器而非LLM，在10M tokens长上下文上实现SOTA多步检索。
前置知识
检索增强生成（RAG）
RAG是一类通过从外部语料库或长上下文中检索最相关的片段，再将其与查询一起输入大语言模型（LLM）以生成答案的技术。其核心思想是让LLM聚焦于更短、更相关的信息，从而缓解静态知识、长上下文效率低和幻觉等问题。RAG通常包含检索器（embedder）和生成器（LLM）两个组件。
Q-RAG本质上是一种多步RAG方法，理解RAG的基础流程是把握Q-RAG将'检索器'也变成RL训练对象的必要性所在。
马尔可夫决策过程（MDP）与Q学习
MDP由状态空间S、动作空间A、转移函数p、奖励函数r和折扣因子γ构成。Q学习通过学习状态-动作值函数Q(s, a)来估计在状态s下采取动作a的长期收益。Q-RAG将多步检索形式化为有限时域MDP，每个episode从查询q开始，每一步选择一个文本片段，奖励基于最终是否收集到所有支撑事实。
Q-RAG的核心数学建模完全建立在MDP框架之上，状态是'查询+已选片段'，动作是'选择一个未选片段'，奖励只在episode终止时给出。
Soft Q-Learning（最大熵强化学习）
Soft Q-Learning在标准Q学习的基础上增加了熵正则项，鼓励策略保持探索性。其值函数定义为Q^π(s, a) = r(s, a) + γV^π(s')，V^π(s) = E_{a~π}[Q^π(s, a) - α log π(a|s)]，其中α是温度参数。Q-RAG采用soft Q-learning，因为长上下文多步检索环境状态空间巨大，需要持续探索。
Q-RAG相比传统Q-learning引入了熵正则化以应对长上下文的探索难题，并配合Boltzmann策略π(a|s) ∝ exp(Q(s,a)/α)进行动作采样。
Rotary Position Embedding（RoPE）
RoPE是一种通过旋转矩阵对token的位置信息进行编码的方法，它将绝对/相对位置信息融入attention中的query/key向量，使得模型能够感知序列中的位置关系。Q-RAG将RoPE应用于action嵌入器以编码片段的文档位置，并提出了一种新的相对位置编码以编码'相对于已选事实'的位置。
Q-RAG的时序推理机制依赖于对片段位置信息的编码，其相对位置映射ρ_t(i)是支撑长程叙事问答能力的关键。
PQN（Parallel Q-Network）
PQN是Gallici等人提出的简化时序差分学习方法，相比DQN移除了replay buffer，支持on-policy并行训练。由于Q-RAG中replay buffer需要重新嵌入所有文档片段来估计后续状态的V/Q值，会显著拖慢训练并增加显存，PQN的无replay特性正好契合这一需求。
Q-RAG的训练算法基于PQN而非经典DQN，这一选择直接关系到'单卡A100即可训练'的工程可行性。
研究动机
现有RAG方法大多依赖单步检索，在BabiLong等benchmark上面对需要多跳推理和时序推理的复杂问题时表现严重不足，例如Mamba2-Hybrid在4K上下文的NIAH-3rd任务仅得95.7分，长上下文退化更明显。三类多步检索方案各有明显短板：第一类是GraphReader、HippoRAG等知识图谱方法，需要LLM处理整个上下文以构建图，推理极慢；第二类是Search-o1等LLM Agent方法，对噪声片段敏感、容易破坏后续查询；第三类是Search-R1、R1-Searcher、RAG-RL等近期工作，它们通过PPO或GRPO微调LLM本身以生成中间搜索查询，虽然性能更好，但需要至少8×A100级别的GPU进行训练，对绝大多数研究者和实践者而言代价过高，限制了与专有大模型的兼容性。
本文的目标是本文的目标是提出一种资源高效的多步RAG方法，绕开'必须微调LLM'这一昂贵路径，转而用强化学习微调一个轻量的嵌入器模型，使其在长上下文（最长10M tokens）中学会多步检索。该方法需要在BabiLong（常识与时序推理）、RULER（针海检索）、HotPotQA和MuSiQue（开放域多跳QA）四类benchmark上达到SOTA或具有强竞争力的结果，同时保证训练只需单卡A100 80GB（约12小时），推理阶段可与任意规模LLM（含闭源专有模型）灵活搭配。
与已有工作不同的是，Q-RAG的独特切入角度是把检索过程本身建模为在文本片段嵌入潜空间中的马尔可夫决策过程，使用基于价值的强化学习（soft Q-learning配合PQN训练算法）直接在状态-动作内积Q(s, a) = ⟨E_s(s), E_a(a, i)⟩的空间中学习检索策略。这一设计有三个不同于现有方法的独到之处：第一，仅微调embedder，LLM保持冻结，因此可与任何规模LLM组合；第二，提出了新的相对位置编码ρ_t(i)显式编码片段相对已选事实的位置，使检索器能在长程叙事中做时序推理；第三，论文证明了这种内积分解的Q函数具有通用逼近能力（Theorem 1）并给出L2逼近速率d^{-s/(D+1)}（Theorem 2），为方法提供了理论保证。
核心方法
Q-RAG的直觉是：把'多步检索'看作'在文本片段嵌入空间中搜索证据'的顺序决策问题。具体来说，给定一个长文档C={c^{(1)}, ..., c^{(m)}}被切成m个片段，初始查询q，agent需要按步骤选择T个片段，构成最终状态s_T=[q, a_0, ..., a_{T-1}]，使得下游LLM仅看这T个片段就能正确回答q。技术上，作者将检索建模为有限时域MDP，使用soft Q-learning（熵正则化的值函数）配合PQN训练算法，并通过两个embedder的内积⟨E_s(s), E_a(a, i)⟩来参数化Q函数；其中状态embedder E_s将'查询+已选片段序列'编码为d维向量，action embedder E_a将'候选片段内容+其位置'编码为d维向量，二者内积即为该候选片段在当前状态下的Q值。论文还引入了相对位置编码ρ_t(i)使action embedder能感知'该候选片段相对于已选事实的位置'，从而支持时序推理。整个训练基于稀疏终态奖励（若最终状态包含全部支撑事实F*则r=1，否则r=0），采用λ-return配合EMA目标网络进行稳定更新。
Q-RAG的核心创新是与现有方法在两个层面上的本质区别。第一个层面是训练对象的选择：Search-R1/R1-Searcher/RAG-RL/ReSearcher等近期工作固定embedder、用GRPO微调LLM以生成中间文本查询；Q-RAG则反其道而行，固定LLM、用时序差分强化学习微调embedder。这使得Q-RAG能以约一个数量级的算力完成训练（单卡A100 vs 8×A100），且推理时能与任意LLM（包括GPT-4、Claude等闭源模型）灵活组合。第二个层面是Q函数的参数化与时序推理机制：将Q(s, a)分解为两个embedder的内积，使动作选择退化为高效的向量相似度计算（而非Beam Retriever那样每步都需要transformer重排序）；并提出基于'已选事实区间'的相对位置编码ρ_t(i) = j·δ + ℓ·(i-b_j)/(b_{j+1}-b_j)，使检索器在长程叙事中能区分'两个事实之间'的片段，从而支持时序推理。这两点共同构成了与所有现有方法的本质区别。
方法步骤详情
Q-RAG方法包含以下完整步骤：(1) 数据准备：将长文档C切分为非重叠片段c^{(i)}，构造数据集D={(C, q, y)}，若提供支撑事实F*则用其作为监督信号，否则可由LLM生成答案y用作奖励。(2) 环境构建：在episode开始时，agent嵌入所有片段得到动作嵌入集{E_a(c^{(i)}, i)}，初始状态s_0=[q]。(3) episode循环：第t步时，先用状态embedder编码s_t得E_s(s_t)，计算所有候选动作的Q值Q_θ(s_t, a)=⟨E_s(s_t), E_a(a, i)⟩，按Boltzmann策略π(a|s) ∝ exp(Q(s, a)/α)采样选择a_t，将a_t加入状态得s_{t+1}=ord([q, a_0, ..., a_t])，并从动作集A_t中移除a_t。(4) 时序编码：对每个候选片段，将其绝对位置i替换为相对位置ρ_t(i)，其中ρ_t(i)=j·δ + ℓ·(i-b_j)/(b_{j+1}-b_j)，δ=10为区间间隔、ℓ=9为区间内分辨率。(5) 训练目标：用PQN on-policy框架计算λ-return G^λ_t = (1-λ)Σ_n λ^{n-1} G_{t:t+n} + λ^{T-t-1} G_t，其中G_{t:t+n} = Σ_{k=1}^n γ^{k-1} r_{t+k} + V_θ'(s_{t+n})，V_θ'通过soft值公式V_θ'(s)=α log Σ_a exp(Q_θ'(s, a)/α)用目标网络计算。(6) 损失与更新：L_Q = E[(Q_θ(s_t, a_t) - G^λ_t)^2]，对θ=(θ_1, θ_2)求梯度更新，并用EMA更新目标网络θ' ← τθ + (1-τ)θ'。(7) 推理：给定新查询q和长文档C，agent运行T步（如HotPotQA用T=2，RULER/BabiLong用T=4）选择片段，将这些片段连同q输入冻结的LLM得到最终答案。
技术新颖性
Q-RAG的技术新颖性体现在四个层面。其一是首次将时序差分强化学习（soft Q-learning + PQN）系统性地应用于embedder微调以解决多步检索任务，区别于以往固定embedder微调LLM的范式；其二是内积Q函数Q_θ(s, a)=⟨E_s(s), E_a(a, i)⟩的参数化方式，使动作选择退化为高效的内积计算，理论上有通用逼近保证（Theorem 1：内积参数化类A_Φ在C(K, ℝ)中稠密）以及L2逼近速率d^{-s/(D+1)}（Theorem 2），其中D=d_x+d_y+1为总维度，s为Sobolev正则性指标；其三是相对位置编码ρ_t(i)的设计，这一编码将文档按已选事实的索引分为k+1个区间，明确编码每个候选片段'位于哪个区间'以及'在区间内的相对位置'，从而支持'两个事实之间'这类时序推理需求；其四是PQN的无replay buffer设计避免了在长上下文环境中的重新嵌入开销，这是工程可行性的关键。
实验结果
Q-RAG在四类benchmark上均取得领先或极具竞争力的结果。在BabiLong（图2a）上，Q-RAG在1M至10M tokens的上下文上取得最高平均准确率，超过Titans、Atlas、ARMT、RMT等专门长上下文方法以及专有LLM和LLM agent；在最难子任务QA3（图2b）上，Q-RAG几乎不随上下文增长而退化，而其他方法的退化越来越严重。在RULER（表1）上，Q-RAG在所有NIAH子任务上取得100%或接近100%的成绩，4K训练即可泛化到1M tokens无明显损失；在多跳QA子任务上Q-RAG在所有上下文长度（4K到1M）均显著优于所有baseline，例如1M上下文时Q-RAG取得52（SH EM）/61（MH EM），而Beam Retriever仅29/39。在HotPotQA和MuSiQue（表2）上，Q-RAG取得HotPotQA Fact F1 0.93、Ans F1 0.76，超过Search-R1（R1-Searcher类）和GraphReader；在MuSiQue OOD上取得Fact F1 0.71、Ans F1 0.52，显著优于Beam Retriever（0.61/0.40）。消融实验（表3）证明了每个组件的必要性：Q-RAG在QA3 1K上下文取得97.8% F1，去掉Soft-Q后降至95.9%，去掉Target Network后骤降至79.2%（且方差高达26%），用SFT替代RL仅得20.33%，完全不做微调仅15.52%。检索步数敏感性实验（表4）显示步数从2增加到3时效果提升明显（HotPotQA EM从0.832升至0.935），但之后趋于饱和，提示最终答案质量主要依赖于正确片段的召回而非片段数量。早停实验（图4、表5、表6）显示基于Q值阈值的早停在HotPotQA和BabiLong QA2上AUC分别为0.96和0.97，最佳阈值约0.2时Fact F1最优。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| BabiLong QA1-QA5平均 | 答案准确率（1M-10M tokens） | Q-RAG：1M-10M tokens全程最高平均准确率（图2a） | Titans、Atlas、ARMT、RMT等均出现显著退化 | 在1M-10M长上下文区间相对其他长上下文专用方法保持最优 |
| BabiLong QA3（最难子任务） | 答案准确率（需≥3步推理+时序感知） | Q-RAG几乎无退化 | 其他长上下文方法退化严重（具体数值未细分报告） | 在最难子任务上扩大领先优势 |
| RULER NIAH 4K | 1st/2nd/3rd needle平均准确率 | Q-RAG：100/100/100 | LongRoPE2-8B：100/100/99；Mamba2-Hybrid：100/100/95.7；Beam Retriever：100/100/98 | 全面达到100%，4K训练即可泛化到1M tokens |
| RULER NIAH 128K | 1st/2nd/3rd needle平均准确率 | Q-RAG：100/100/100 | LongRoPE2-8B：100/100/99 | 在128K上下文保持100%准确率 |
| RULER NIAH 1M | 1st/2nd/3rd needle平均准确率 | Q-RAG：100/100/100 | LongRoPE2未在1M上下文报告 | 唯一在1M上下文达到100%的模型 |
| RULER Multi-hop QA（1M tokens） | Single-hop/Multi-hop Exact Match | Q-RAG：52/61 | Beam Retriever：29/39（4K上下文） | 在1M上下文下显著超过Beam Retriever（后者无法扩展到此规模） |
| HotPotQA | Fact F1 / Fact EM / Ans F1 / Ans EM | Q-RAG：0.93 / 0.89 / 0.76 / 0.59 | Beam Retriever：0.97/0.94/0.77/0.61；Search-R1：0.81/0.66/0.65/0.52；RAG-RL：0.82/-/0.69/- | 接近Beam Retriever，超过Search-R1/RAG-RL等基于LLM微调的方法 |
| MuSiQue (OOD) | Fact F1 / Fact EM / Ans F1 / Ans EM | Q-RAG：0.71 / 0.55 / 0.52 / 0.37 | Beam Retriever：0.61/0.36/0.40/0.27；Search-R1：0.71/0.55/0.51/0.39 | OOD泛化显著优于Beam Retriever，与Search-R1持平 |
| BabiLong QA3 消融 | 支撑事实检索F1（1K上下文） | Q-RAG：97.8 ± 0.17 | w/o Soft-Q：95.9；w/o Target：79.2；w/ SFT：20.33；w/o FT：15.52 | 完整Q-RAG相比各消融变体均有显著提升，特别是相对SFT提升77.5个百分点 |
局限与改进
作者承认的局限性包括：(1) Q-RAG目前仅使用支撑事实作为奖励信号，未探索基于LLM的奖励设计（如exact match或F1答案奖励），文中明确将'使用LLM反馈作为奖励'列为未来工作。(2) Q-RAG的细粒度能力仍受限于其基础embedder（如multilingual-e5-large、contriever等小模型）的语义理解能力，对极复杂的组合推理帮助有限。(3) 在MuSiQue OOD上，Q-RAG的Ans F1仅0.52，说明即便检索到的证据正确，最终LLM（QwQ-32B）仍可能综合失误。(4) 训练时每个episode最多需要T步前向传播，对超长上下文（10M tokens）的工程开销仍较大。本人观察到的额外局限包括：(5) Q-RAG在Table 1的SH QA子任务上逊于LongRoPE2-8B（4K时62 vs 79、128K时55 vs 56），说明在短上下文的多跳QA上专门的长上下文方法仍占优势；(6) 当retrieval budget大于真正需要的步数时，会引入噪声片段（F1下降），虽然文中说明下游LLM对此有一定容忍度；(7) 状态embedder在每步都要重新处理累积的'查询+已选片段'，当T较大或片段较长时存在O(K·t_state)开销；(8) 相对位置编码引入了两个超参数δ和ℓ，目前未给出其在不同任务下的系统调参结果。
独立分析的弱点
Q-RAG存在的弱点可归纳为四点。第一是奖励信号的稀疏性：当前实现中episode只有在选择完所有支撑事实后才得到r=1的稀疏终态奖励，这导致长episode（如T=4）的训练信号极弱，特别在训练初期几乎完全靠随机探索，作者可以考虑引入中间步骤的dense reward（如每个正确片段0.5分）或shaping reward以加速收敛。第二是对闭源LLM评估的依赖：Q-RAG的下游答案生成强依赖于QwQ-32B或Qwen3-4B，在RULER的SH QA上62 vs LongRoPE2-8B的79说明Q-RAG检索到的证据需要更强的生成器才能转化为更好答案，可考虑结合re-ranker或更强的生成模型。第三是OOD泛化仍有提升空间：MuSiQue OOD的Ans F1仅0.52，说明从HotPotQA迁移到MuSiQue时存在明显的domain shift，可通过domain adaptation或在多领域混合数据上联合训练来缓解。第四是状态表征的简单拼接：当前状态s_t=ord([q, a_0, ..., a_{t-1}])仅做顺序拼接，没有显式建模已选片段之间的逻辑依赖，对于需要复杂推理链的任务可能不足，可考虑引入结构化状态表示（如片段间的因果图）。
未来方向
作者在结论部分明确指出三个未来方向：(1) 使用LLM反馈作为奖励信号，例如让LLM评估检索证据的相关性和完整性，从而提供更密集的奖励；(2) 强化embedding space中的组合与时序推理能力，例如通过结构化约束或对比学习使embedder能编码更复杂的关系；(3) 探索与生成模型的更紧密集成，在保持效率与可扩展性的同时让检索器与生成器协同优化。基于本文成果可延伸的方向还包括：(4) 将Q-RAG的多步检索范式推广到结构化数据（如知识图谱、表格）的检索；(5) 探索更复杂的早停策略或多agent协作机制，让多个Q-RAG agent分工检索不同类型的证据；(6) 研究Q-RAG与端到端长上下文模型（如Recurrent Memory Transformer）的混合架构，用RMT做粗筛再用Q-RAG做精排；(7) 将Q-RAG扩展到多模态长上下文（如长视频、长文档图像）的检索任务。
复现评估
Q-RAG的复现性非常好，作者明确在GitHub（https://github.com/griver/Q-RAG）开源了全部代码、数据和训练说明，并在Hugging Face上提供了预训练checkpoint。所有实验使用的backbone embedder均为公开可下载的模型（multilingual-e5-large、Alibaba-NLP/gte-multilingual-base、facebook/contriever），超参数和训练schedule在附录F中完整给出：AdamW优化器、学习率1.5×10^{-5}、β1=0.9/β2=0.98、weight decay 5×10^{-4}、线性warm-up 1000步、梯度裁剪ℓ2阈值2.0、梯度累积8步、batch size 12×8=96、γ=0.99、α=0.05、λ=0.5、τ=0.02、动作最大长度220 tokens。算力需求方面，单个模型的端到端训练在单卡A100-80GB上不超过12小时，远低于Search-R1/R1-Searcher的8×A100配置。每个数据集的具体chunk size和T设置列在Table 8中：HotPotQA用220 tokens chunk + T=2、MuSiQue用110 tokens + T=4、BabiLong/RULER用64 tokens + T=4。整体复现难度为低-中，主要工作量是配置环境和下载数据，模型实现本身相对简洁。
论文图表