Q-RAG: Long Context Multi-step Retrieval via Value-based Embedder Training

Paper: arXiv:2511.07328 Code: griver/Q-RAG Code reference: pqn-qa @ b821cdcf (2026-04-14)

1. 核心结论与定位

Q-RAG 的核心贡献不是把长上下文全部塞进更大的 LLM，也不是训练一个会写搜索 query 的小 LLM，而是把多步检索本身建模成一个 value-based RL 问题，并只微调 retriever/embedder。它把“下一步应该取哪个 chunk”看成离散动作选择：状态是问题加已经取到的证据，动作是剩余候选 chunk，Q 值由 state embedding 与 action/chunk embedding 的内积给出。这样做的直接好处是，检索器可以多步收集互相依赖的证据，而回答 LLM 仍然只读少量被检索出的 chunk；原始上下文从 4K 到 10M tokens 增长时，回答 LLM 的上下文长度不随原文长度线性增长。

这篇论文最适合归入 LLM & VLM / Long-Context & Streaming：它解决的是长上下文下的多跳、多步、超长检索问题，方法侧虽然用了 RL，但目标不是通用 post-training，而是长上下文 RAG 的检索阶段。它和传统 dense retriever 的差异在于：传统 RAG 多数是 single-step scoring，一次性按 query 与 chunk 相似度排序；Q-RAG 则允许第 1 步拿到一个事实后，第 2 步把这个事实并入状态，再去找后续事实。它和 Search-R1、IM-RAG 一类 LLM-agent 检索的差异在于：后者通常要让 LLM 生成查询或中间推理，训练/推理成本高；Q-RAG 把可训练对象限制在 embedder agent 上，论文声称单个模型训练不超过 12 小时、单张 A100-80GB 即可完成。

一句话理解：Q-RAG 把“长文档里逐步找证据”转成 soft Q-learning over chunks，用 Q 值训练 embedder，让 retrieval trajectory 自己学会先找哪个证据、再找哪个证据，而不是靠 LLM 反复读全文或生成搜索计划。

1.1 问题背景

RAG 的常见瓶颈有三类。第一，LLM 直接处理长上下文会遇到 attention dilution、推理成本高、上下文窗口限制和幻觉问题；即使窗口扩到很长，关键事实也可能被大量无关文本稀释。第二，单步检索对多跳问答不够：一个问题可能需要先找到人物位置、再根据时间顺序找到后续事件，或者先定位 key，再找与该 key 相关的多个 value。第三，多步检索的已有做法常常依赖小 LLM fine-tuning 或 agentic search：这些方法可解释性强，但会把训练和推理成本转移到 LLM 上，且不方便直接使用更大的 reader LLM。

Q-RAG 的切入点是：多步检索真正需要学习的是“证据选择策略”，不一定要训练一个语言模型来产生自然语言搜索指令。若能在 embedding 空间里直接学习状态—动作价值函数，就可以把长文档检索变成一个轻量 RL agent：状态编码器读 question + selected chunks，动作编码器读候选 chunk + 位置，Q 值决定下一步取哪个 chunk。这个设计保留了 dense retrieval 的可扩展性，又给 single-step retriever 补上了 trajectory-level credit assignment。

1.2 主要贡献

Value-based embedder training：使用 maximum-entropy TD learning/PQN 训练状态 embedder 与动作 embedder，Q 函数写成 $Q_{θ} (s, a^{i}) = ⟨ E_{s} (s; θ_{1}), E_{a} (a^{i}, i; θ_{2})⟩$ 。
多步检索而非 LLM 微调：检索 agent 负责迭代收集 chunk，reader LLM 保持 frozen；论文在 HotPotQA、MuSiQue、RULER、BabiLong 上验证该策略。
时间/相对位置编码：针对叙事类任务，作者提出相对位置映射，让候选 chunk 的位置相对于已选事实而不是只相对于文档开头编码。
长上下文可扩展性：BabiLong 与 RULER 的实验覆盖到 1M、10M tokens；RULER NIAH 子任务里 Q-RAG 在 4K 到 128K 都达到 100 NIAH Avg，1M 仍为 99.7。

2. 方法总览：把多步检索写成 RL 环境

Figure 1 解读：图中的 retriever agent 并不是一次性返回 top-k，而是在一个 episode 内反复选择 chunk。每一步会把当前 question 与已经选择的 chunk 拼成 state，然后对所有尚未选择的候选 chunk 计算 Q 值。Q 值高的 chunk 更可能被选入 memory；选中后，state 更新、动作集合移除该 chunk，下一步继续。最后 reader LLM 只接收这些被选出的 chunk，而不是完整长上下文。

2.1 MDP 定义

Q-RAG 的环境可以抽象为一个 episodic MDP。给定问题 $q$ 和被切分成 chunks 的文档 $C = {c_{1}, \dots, c_{N}}$ ，第 $t$ 步状态 $s_{t}$ 包含问题和已选 chunk 序列 $M_{t} = {c_{i_{1}}, \dots, c_{i_{t - 1}}}$ 。动作集合 $A_{t}$ 是尚未被选择的 chunks。agent 选择 $a_{t} = c_{i}$ 后，环境把该 chunk 加入 memory，得到 $s_{t + 1}$ ，并从可选动作中删除 $c_{i}$ 。episode 的最大长度是预设检索步数 $T$ ，不同数据集使用不同 $T$ ：HotPotQA 主实验 $T = 2$ ，MuSiQue $T = 4$ ，BabiLong 与 RULER $T = 4$ ；早停实验则允许 HotPotQA 最多 $T = 5$ 并用 Q-value threshold 决定是否继续。

奖励主要来自 ground-truth supporting facts。代码中的 GroundTruthFeedback 会维护已经发现的 supporting fact 集合：每次新选到尚未发现的 supporting fact，可以给 per-fact reward；当已选 chunks 覆盖全部 ground-truth supporting facts 时，给 completion reward，并可用 penalize_extra_steps=True 对多取无关 chunk 的 trajectory 降低终止奖励。这个 reward 设计对应论文表格里的 Fact EM/F1：EM 关心是否找全事实，F1 同时惩罚噪声 chunks。

2.2 Q 函数：state embedder × action embedder

论文的核心建模是把 Q 函数写成两个 embedder 的内积： $Q_{θ} (s, a^{i}) = ⟨ E_{s} (s; θ_{1}), E_{a} (a^{i}, i; θ_{2})⟩ .$ 其中 $E_{s}$ 编码当前状态，也就是 question 加 memory 中已有证据； $E_{a}$ 编码候选 chunk 内容和文档位置 $i$ 。这让 Q-RAG 在形式上接近 dense retrieval：它仍然是 embedding similarity，但相似度的 query 不再是固定问题，而是会随已经检索到的证据动态变化。也就是说，第一步可能靠问题本身定位到某个实体，第二步的 state embedding 已经包含这个实体事实，因此可以把下一跳证据排到更高。

选择动作时，Q-RAG 使用 Boltzmann/softmax policy： $π (a_{t} ∣ s_{t}) = \frac{e x p (( Q _{θ} ( s _{t} , a _{t} ) - q ) / α )}{\sum _{a \in A_{t}} e x p (( Q _{θ} ( s _{t} , a ) - q ) / α )}, q = max_{a \in A_{t}} Q_{θ} (s_{t}, a) .$ 这里减去最大值 $q$ 是数值稳定技巧， $α$ 控制探索强度。论文和代码都强调 $α$ 会随学习率 schedule 退火：训练初期允许探索多个候选 chunk，后期趋向 greedy。这个机制对长文档尤其重要，因为候选动作非常多，纯 greedy 容易被局部相似但不支持最终答案的 chunk 吸住；纯随机则会浪费 episode 步数。

2.3 Soft Q 与 $λ$ -return

Q-RAG 采用 maximum-entropy RL 的 soft value 形式。给定 target network $θ^{'}$ ，离散动作空间下 state value 可写成： $V_{θ^{'}} (s_{t}) = α lo g \sum_{a \in A_{t}} exp (Q_{θ^{'}} (s_{t}, a) / α) .$ 训练目标不是只用一步 TD，而是用 $λ$ -return： $G_{t}^{λ} = (1 - λ) \sum_{n = 1}^{T - t - 1} λ^{n - 1} G_{t : t + n} + λ^{T - t - 1} G_{t},$ 其中 $G_{t : t + n} = \sum_{k = 1}^{n} γ^{k - 1} r_{t + k} + V_{θ^{'}} (s_{t + n})$ 。直觉上，单步 TD 对检索 trajectory 的 credit assignment 太短视：一个 chunk 本身可能没有立刻触发 completion reward，但它为下一步找到关键 supporting fact 铺路。 $λ$ -return 在 Monte Carlo 和 TD 之间折中，让早期 chunk 能从后续成功中获得信用。论文默认训练细节给出 $γ = 0.99$ 、 $α = 0.05$ 、 $λ = 0.5$ 、 $τ = 0.02$ ，这些值来自论文 Appendix 的训练细节；公开代码的 configs/algo/pqn_e5_hotpotqa.yaml 也使用同一组关键超参。

2.4 Target network 与双头 Q

公开实现中，rl/agents/pqn.py 的 PQN 会实例化 TextQNet 作为 critic，同时构造 TextVNet 和 ActionEmbedTarget 作为 target 分支。policy_apply() 先用当前 policy 在可选 chunks 中采样或 argmax，再用 target value net 计算下一状态的 soft value。update() 从 rollout batch 中反向递推 target：先用最后一步的 reward 和 target Q 初始化 return，再从后往前调用 _get_target() 混合 immediate reward、target value 和上一轮 return，最后用当前 critic 的两个 Q 头分别拟合同一个 target。代码里的 loss 是两个 MSE 的平均：0.5*MSE(qf_1,target)+0.5*MSE(qf_2,target)，这比只用一个 head 更接近 clipped/double-Q 的稳定化思想。

target 更新不是硬拷贝，而是 EMA：论文写作 $θ^{'} \leftarrow τ θ^{'} + (1 - τ) θ$ ；代码里 TextVNet.update() 和 ActionEmbedTarget.update() 调用 soft_update(..., decay=tau)。注意实现中的 tau 语义需要按 soft_update 函数解释；笔记中只把它作为“慢更新 target 分支”的角色使用，不把公式中的系数方向过度绑定到实现细节。

2.5 相对位置编码：为什么对 BabiLong/叙事任务重要

Figure 2 解读：QA3 是 BabiLong 中更强调多跳时间推理的任务。单个 chunk 往往只说明某个事件片段，正确答案依赖事件顺序。Q-RAG 在这类任务上几乎不随 context length 增长而明显退化，说明它不只是“找语义相似句子”，还需要利用已选事实的顺序关系继续检索。

论文的相对位置编码针对如下问题：叙事文本里“Mary went to the kitchen”可能出现多次，候选 chunk 的文本内容相同或相近，仅靠语义 embedding 不能判断哪个事件是当前问题所需事实。若已经选到一个事实，下一步候选 chunk 应该知道自己位于该事实之前、之后，还是两个已选事实之间。因此作者定义已选 chunk 索引集合 $S_{t} = {i_{1} < \dots < i_{k}}$ ，它把文档分成 $k + 1$ 个区间；相对位置映射 $ρ_{t}$ 给每个候选 chunk 分配新的实数位置，既标识它落在哪个区间，又保持区间内部的相对顺序。动作 embedder 使用 RoPE/位置编码读入这个相对位置。

这部分的关键 intuition 是：多步检索的 state 不应该只告诉模型“已经有什么事实”，还应该告诉模型“这些事实在原文里处于什么时间/位置关系”。对普通 HotPotQA，chunk 位置可能不是主要信息；对 BabiLong 的故事和 RULER 的 needle 位置，位置是支持事实的一部分。公开代码中 BabiLong 配置 configs/envs/babilong.yaml 使用 positions_processor: relative，而 HotPotQA+MuSiQue 默认使用 positions_processor: none，这与论文动机一致：只有当任务需要时间顺序或位置推理时，才显式使用相对位置处理。

2.6 Plan Q-RAG

论文还报告了 Plan Q-RAG。它不是主方法的必要组成，而是一个计划增强变体：先引入一个“plan”或 query decomposition 风格的辅助信息，再使用 Q-RAG 检索。表 2 里 Plan Q-RAG 在 HotPotQA 上 Fact F1/EM 为 0.95/0.91，略高于 Q-RAG 的 0.93/0.89；但在 MuSiQue OOD 上 Q-RAG 的 Fact F1/EM 为 0.71/0.55，反而高于 Plan Q-RAG 的 0.69/0.53。我的理解是：planning 对 in-domain 热点数据有帮助，但如果 plan 生成带来错误先验，纯 value-based trajectory 可能更稳。

2.7 复杂度

论文 Appendix 讨论复杂度：answerer LLM 只看检索出的 chunks，因此相对原始 context length $N$ ，reader 的时间/显存复杂度是 $O (1)$ ；真正随 $N$ 增长的是 retrieval agent 对所有 chunks 的 embedding/scoring。动作 embedder 对整个文档做一遍 chunk embedding，成本近似 $O (N)$ ；state embedder 每步处理 question + memory，成本与检索步数 $T$ 和 memory 长度有关。这个取舍很清楚：Q-RAG 没有让长文档处理“免费”，而是把主要成本从 quadratic attention over full context 迁移到 linear retrieval scoring，再把 reader LLM 的输入压缩到少量证据。

3. 训练与推理流程

3.1 训练流程伪代码（按论文 + 代码对齐）

# Inputs: question q, document chunks C, max retrieval steps T
# Networks: state embedder Es, action embedder Ea, target networks Es', Ea'
for update in training:
    batch = []
    for env in parallel_envs:
        state = reset(q, C)                  # state = question + selected chunks
        action_embeds = Ea(C, positions)     # candidate chunk embeddings
        target_action_embeds = Ea'(C, positions)
        trajectory = []
        for t in range(T + 1):
            q_values = dot(Es(state), action_embeds[available])
            action = sample_softmax(q_values / alpha)  # or random during warmup
            next_state, selected_chunk, reward, done = env.step(action)
            target_v = alpha * logsumexp(dot(Es'(state), target_action_embeds[available]) / alpha)
            trajectory.append((state, action, reward, target_v, done))
            state = next_state
            if done:
                break
        batch.extend(trajectory)
    targets = lambda_return(batch, gamma=0.99, Lambda=0.5)
    loss = mse(Q1(state, action), targets) / 2 + mse(Q2(state, action), targets) / 2
    optimize(loss)
    ema_update(target_networks, online_networks, tau=0.02)
    anneal_alpha_with_lr()

这段伪代码有三个要点。第一，candidate action embeddings 可以在 episode 内缓存并更新位置，因此长文档的主要成本集中在候选 chunk 编码。第二，环境 reward 不需要 LLM judge；公开实现的主 reward 是 supporting facts 覆盖。第三，训练时的 policy 是 stochastic softmax，评估时可以 return_argmax=True，即取最大 Q 值动作。

3.2 推理流程

state = question
memory = []
for t in range(T):
    q_values = score_all_available_chunks(state, document_chunks)
    chunk = argmax(q_values)
    memory.append(chunk)
    state = concatenate(question, memory)
answer = reader_llm(question, memory)

如果使用早停，推理流程会在每一步看当前最大 Q 值是否低于 threshold。Appendix 的 HotPotQA 早停表显示，threshold 从 -0.1 提高到 1.1 时，episode length 从 4.99 逐渐降到 0；较平衡的阈值在 0.0 到 0.3 左右。例如 threshold=0.0 时 average episode length 为 2.82，Fact EM/F1 为 0.954/0.843，Ans EM/F1 为 0.592/0.761；threshold=0.2 时 episode length 进一步降到 2.13，Fact EM/F1 为 0.883/0.917，Ans EM/F1 为 0.587/0.752。也就是说，早停可以减少冗余检索，但过高阈值会过早停止，导致 facts 覆盖下降。

Figure 3 解读：早停实验把 Q 值当成“继续检索是否值得”的置信信号。低阈值对应保守停止：几乎不会错过支持事实，但会继续取噪声；高阈值对应激进停止：retrieval steps 少，但更容易没取全证据。这个分析说明 Q-RAG 的 Q 值不仅能排序 chunk，也能作为动态 budget 控制信号。

3.3 训练细节

论文 Appendix F 给出主训练配置：AdamW，学习率 $1.5 \times 1 0^{- 5}$ ， $β_{1} = 0.9$ ， $β_{2} = 0.98$ ， $ϵ = 1 0^{- 6}$ ，weight decay $5 \times 1 0^{- 4}$ ；warmup 1000 steps，然后线性衰减到初始学习率的 10%；gradient clipping 最大 $ℓ_{2}$ norm 为 2.0；base mini-batch size 12，gradient accumulation 8，因此有效 batch size 96。目标和算法超参为 $γ = 0.99$ 、 $α = 0.05$ 、 $λ = 0.5$ 、 $τ = 0.02$ ，action representation 最大长度 220 tokens。公开代码中 HotPotQA E5 配置与这些值基本一致：configs/algo/pqn_e5_hotpotqa.yaml 使用 alpha=0.05、Lambda=0.5、tau=0.02、lr=1.5e-05、weight_decay=0.0005；configs/training.yaml 是默认 BabiLong/Contriever 起点，不能直接当作所有实验的最终数值。

4. 实验结果

4.1 RULER：NIAH 近乎满分，QA 子任务更稳

RULER 表 1 覆盖 Single-needle、Multi-keys、Multi-values、Multi-query，以及 single-hop/multi-hop QA。Q-RAG 在 NIAH 子任务上的结果非常强：4K、16K、32K、128K 的 NIAH Avg 全部是 100；1M 时仍有 99.7。对比 LongRoPE2-8B，4K NIAH Avg 为 99.7、16K 为 98.8、32K 为 98.9、128K 为 96.7；Q-RAG 在这些长度上保持 100。对 QA 子任务，Q-RAG 的 multi-hop QA 在 4K/16K/32K/128K/1M 分别为 67/64/65/65/61，显著高于 Beam Retriever 的 39/35（只报告到 16K）以及 LongRoPE2-8B 在同长度附近的 60/58/55/50。single-hop QA 则在 4K 到 1M 从 62 下降到 52，说明超长长度下 QA 类任务仍有退化，但退化比 NIAH 更集中在 answerability 和跨跳证据组合上。

我的解读：RULER NIAH 的高分证明 Q-RAG 的 retriever 可以在超长 distractor 中稳定定位 needle；QA 子任务的下降说明“找出所有相关 chunk”还不是完整问答能力，reader LLM 的 reasoning、chunk 内容是否足够完整、以及训练集到 RULER QA 的分布偏移都会影响最终答案。

4.2 Open-domain QA：HotPotQA in-domain 与 MuSiQue OOD

表 2 比较 HotPotQA 与 MuSiQue OOD。HotPotQA 上，Beam Retriever 的 fact retrieval 最强：Fact F1/EM 为 0.97/0.94，Ans F1/EM 为 0.77/0.61；Plan Q-RAG 为 0.95/0.91 和 0.76/0.60；Q-RAG 为 0.93/0.89 和 0.76/0.59。也就是说，在 HotPotQA in-domain，Q-RAG 并没有超过 Beam Retriever，但已非常接近，且不需要 Beam Retriever 式的 sequence scoring 训练。

MuSiQue OOD 上趋势不同。Q-RAG 的 Fact F1/EM 为 0.71/0.55，Ans F1/EM 为 0.52/0.37；Beam Retriever 为 0.61/0.36 和 0.40/0.27；Search-R1 为 0.71/0.55 和 0.51/0.39。Q-RAG 在 fact retrieval 上与 Search-R1 持平，在 answer F1 上略高，但 answer EM 略低。这个结果支持论文的一个隐含观点：value-based embedder 不是只记住 HotPotQA 的模板，而能迁移到另一个多跳 QA 数据集；不过最终 answer EM 仍受 reader 和生成格式影响。

Setting	Method	HotPotQA Fact F1/EM	HotPotQA Ans F1/EM	MuSiQue OOD Fact F1/EM	MuSiQue OOD Ans F1/EM	Avg Ans F1/EM
Fine-tuned on HotPotQA	Plan Q-RAG	0.95 / 0.91	0.76 / 0.60	0.69 / 0.53	0.51 / 0.36	0.64 / 0.48
Fine-tuned on HotPotQA	Q-RAG	0.93 / 0.89	0.76 / 0.59	0.71 / 0.55	0.52 / 0.37	0.64 / 0.48
Fine-tuned on HotPotQA	Beam Retriever	0.97 / 0.94	0.77 / 0.61	0.61 / 0.36	0.40 / 0.27	0.59 / 0.44
Fine-tuned on HotPotQA	Search-R1	0.81 / 0.66	0.65 / 0.52	0.71 / 0.55	0.51 / 0.39	0.58 / 0.46

4.3 BabiLong：上下文到 10M tokens 的长期稳定性

Figure 4 解读：BabiLong 的平均任务曲线体现了 Q-RAG 的核心卖点：当上下文长度从 1M 扩展到 10M，许多 recurrent memory、long-context 或 agentic baseline 的曲线会明显下降，而 Q-RAG 的实线保持在高位。这里的优势不是 reader LLM 读了 10M tokens，而是 retriever 把 reader 输入稳定压缩成少量 facts。

BabiLong 是这篇论文最能展示“long-context retrieval”价值的 benchmark。它用 PG19 等长文本作为噪声，把 bAbI 支持事实藏到长上下文中，QA1–QA5 需要不同数目的 supporting facts 和时间推理。论文强调 QA3 是最难子任务之一，需要较长 reasoning chain 和 temporal awareness；Q-RAG 在 QA3 上几乎没有随 10M context 明显退化。结合相对位置编码，可以把结果理解为：state/action embedding 学到的不只是 lexical matching，也学到了“已经拿到某个事件后，下一个证据应该位于何处”的顺序依赖。

4.4 Ablation：Soft-Q、target network、SFT baseline

Figure 5 解读： $α$ 控制 entropy regularization/exploration。过小会更接近 greedy，可能早期探索不足；过大则会给低价值动作过多概率。论文用 QA2/QA3 曲线说明中等温度更稳。

Figure 6 解读： $λ$ 控制 TD target 的 horizon。多步检索里 reward 可能延迟出现，因此完全 one-step 过短视；但过高也可能引入高方差。默认 $λ = 0.5$ 是稳定性和 credit assignment 的折中。

表 3 在 BabiLong QA3 上给出支持事实 retrieval F1。完整 Q-RAG 在 1K/4K/32K/128K/1M 分别为 97.8±0.17、97.4±0.14、97.1±0.08、96.8±0.08、96.5±0.16。去掉 Soft-Q 后仍较强但下降到 95.9、95.5、94.5、94.0、93.3；去掉 target network 则大幅不稳定，均值约 79/78/77/77/75，标准差高达 26–28。SFT 多步 RAG 和无 fine-tuning 多步 RAG 都只有 15–21 左右，说明只用 supporting facts 做监督序列学习没有解决这个环境中的 exploration 与 delayed reward。

4.5 Runtime / inference cost

Figure 7 解读：runtime 图展示 Q-RAG 的成本随 context 增长主要体现在线性检索侧，而 reader LLM 的输入长度固定。对实际系统来说，这意味着可以把长上下文处理拆成两个 budget：retriever 扫描全部 chunks，reader 只读 top trajectory。相比 GraphReader 这类要让 LLM 反复处理图/全文的 agent，Q-RAG 的 reader 端更可控。

5. 代码阅读与实现映射

Code reference: pqn-qa @ b821cdcf (2026-04-14)

Paper Concept	Source File	Key Class/Function	Notes
主训练循环	`train_q_rag.py`	`evaluate`, main training loop	Hydra 组合 config，实例化 `PQN`、`QAEnv`、`ParallelTextEnv`；rollout 后调用 `agent.update()`。
Value-based agent / PQN	`rl/agents/pqn.py`	`PQN.__init__`, `select_action`, `_get_target`, `update`	读取 `gamma/alpha/Lambda/tau`，维护 critic、policy、target value/action embedder；反向计算 $λ$ -return。
Q 网络	`rl/q_module.py`	`TextQNet.forward`	state embedder 和 action embedder 输出向量，两个 Q head 返回 `logits_1/logits_2`。
Soft policy	`rl/q_module.py`	`TextQNetPolicy.forward`	对可用动作 mask 后，在 top-k/available 动作上按 softmax 采样；评估时返回 argmax。
Soft target value	`rl/q_module.py`	`TextVNet.forward`	用 target state/action embeddings 计算两个 `alpha*logsumexp(Q/alpha)` value。
Target action embedder	`rl/q_module.py`	`ActionEmbedTarget.update`	用 EMA/soft update 同步 action embedder target。
QA retrieval 环境	`envs/qa_env.py`	`QAEnv.reset`, `QAEnv.step`	state 是问题和已选 chunks；step 后从 feedback model 得到 reward/terminated。
Ground-truth reward	`rl/feedback/feedback.py`	`GroundTruthFeedback.reward`	新 supporting facts 可给 step reward；找全 supporting facts 给 completion reward；可惩罚 extra steps。
BabiLong 位置处理配置	`configs/envs/babilong.yaml`	`positions_processor: relative`	BabiLong 使用 relative position processor，符合 temporal reasoning 设定。
OpenQA/RULER 配置	`configs/envs/hotpotqa+musique.yaml`, `configs/envs/niah.yaml`	`max_steps`, `feedback_model`	HotPotQA+MuSiQue 训练用 combined dataset；NIAH 任务用 `GroundTruthFeedback`。
LLM 前过滤	`prompts_and_metrics/chunk_filtering.py`	`QValueChunkFilter`, `EarlyStopChunkFilter`	支持基于 Q-value 或 support fact 的 chunk filtering/early stopping。

5.1 代码与论文一致处

代码和论文在主机制上高度一致。PQN.__init__ 从 config 中读入 gamma、alpha、Lambda、tau，然后实例化 state/action embedder、critic、policy、target value net 与 target action embedder。TextQNet.forward 中，state embedding 来自 state_embed(input_ids, attention_mask)，action embedding 来自 action_embed(input_ids, attention_mask, positions)；随后返回两个 logits，这对应论文的 inner product Q approximation，只是实现中使用双头/双值以增强稳定性。

select_action_batch 调用 policy_apply()：当前 policy 根据 Q 值选动作，同时 target value net 计算 target Q/value。训练早期 random=True 时会用 TextRandomPolicy 随机探索。train_q_rag.py 中 parallel_env.rollout(..., random=(step < 2 * cfg.learning_start)) 表明 warmup 阶段确实不完全依赖当前 Q policy。update() 里从最后时间步向前构造 return；_get_target() 体现了 $r + γ [(1 - λ) V + λ G]$ 的递推思想。优化后，代码会更新 scheduler、按学习率比例退火 alpha，并更新 target nets。

5.2 代码中需要注意的复现细节

第一，仓库 README 明确说代码仍处于 active debugging/refactoring，复现时应固定 commit。本文使用 pqn-qa@b821cdcf。第二，论文 Appendix 的训练参数与仓库默认 config 不完全等价：configs/training.yaml 是默认训练入口，使用 steps_count=50000、batch_size=64、accumulate_grads=2、max_action_length=256；而 HotPotQA E5 论文设置在 configs/algo/pqn_e5_hotpotqa.yaml，其中 lr、weight decay、alpha、lambda、tau 更贴近 Appendix F。因此记录训练配置时不能只看 configs/training.yaml。第三，BabiLong/RULER 用 facebook/contriever，OpenQA 用 multilingual-e5-large 或 Alibaba-NLP/gte-multilingual-base，这些 backbone 的 embedding 维度和最大长度不同，chunk size 也不同。

5.3 每个关键组件的实现级伪代码

Q-value scorer：

def q_values(state_memory, candidate_chunks):
    s = state_embed(state_memory.input_ids, state_memory.attention_mask)
    a = action_embed(candidate_chunks.input_ids,
                     candidate_chunks.attention_mask,
                     positions=candidate_chunks.position)["rope"]
    q1 = inner_product_head_1(s, a)
    q2 = inner_product_head_2(s, a)
    return q1, q2

Soft policy：

def select_action(state, action_embeds, alpha, evaluate=False):
    logits = dot(policy_state_embed(state), action_embeds)
    logits[~state.available_mask] = very_small
    if evaluate:
        return argmax(logits)
    probs = softmax((logits - max(logits)) / alpha)
    probs[~available_or_top_mask] = 0
    return Categorical(probs).sample()

Ground-truth feedback：

def reward(obs, info):
    pred_idx = set(obs["pred_idx"])
    sf_idx = set(info["sf_idx"])
    new_facts = pred_idx.intersection(sf_idx) - found_facts
    step_reward = len(new_facts) * per_fact_reward
    found_facts.update(new_facts)
    if not completed and sf_idx.issubset(found_facts):
        completed = True
        if penalize_extra_steps:
            completion = (0.5 + 0.5 * len(sf_idx) / len(pred_idx)) * completion_reward
        else:
            completion = completion_reward
    return step_reward + completion

这些伪代码说明 Q-RAG 的实现没有使用复杂 LLM judge 或外部工具调用；训练信号主要依赖数据集标注的 supporting facts。这也是它能在单张 A100 上训练的原因之一，但也意味着方法需要任务具备可监督的支持事实或可构造 reward。

6. 局限性与风险

6.1 对 supporting-fact supervision 的依赖

Q-RAG 的主 reward 需要知道 ground-truth supporting chunks。HotPotQA、MuSiQue、BabiLong、RULER 都能提供或构造这种监督，但真实企业文档、开放网页搜索、代码库问答不一定有标注的支持事实。若没有 supporting facts，就需要替代 reward，例如答案正确性、LLM judge、contrastive preference 或 weak supervision；这些 reward 会重新引入噪声、成本和 credit assignment 问题。因此 Q-RAG 更像一种有监督/可监督 long-context retriever training recipe，而不是完全无监督 RAG 系统。

6.2 检索成功不等于答案成功

表 2 中 Q-RAG 在 MuSiQue OOD 的 Fact F1/EM 强于 Beam Retriever，但 answer EM/F1 的优势没有同等幅度；RULER QA 子任务中 NIAH Avg 接近满分，但 QA 分数仍随长度下降。这说明最终效果受三层影响：retriever 是否找全证据，chunk 是否保留了足够上下文，reader LLM 是否能在这些 chunks 上推理。Q-RAG 强化的是第一层。若 chunks 太短，supporting fact 可能被切断；若 reader LLM 对多跳推理弱，检索再好也不能保证最终答案。

6.3 线性扫描仍然有成本

Q-RAG 避免了 reader LLM 对原始长上下文的 quadratic attention，但 retrieval agent 仍需要对所有 chunks 打分或预编码。对 10M tokens 文档，chunk 数量很大；若每个 query 都要从头计算所有 chunk embeddings，服务端成本仍不可忽略。实际部署应缓存 document/action embeddings，只在 state embedding 和 dot-product scoring 上做在线计算；多用户、多文档、多版本文档的缓存一致性会成为系统工程问题。

6.4 位置编码可能任务相关

相对位置编码对叙事和时间顺序任务非常合理，但在百科型 open-domain QA 中，事实位置常常是任意的。若强行把位置作为重要信号，可能让模型学习到数据集特定偏差。代码配置也反映了这一点：BabiLong 使用 relative position，HotPotQA+MuSiQue 默认不使用。复用 Q-RAG 到新领域时，应先判断“事实顺序”是否是语义的一部分，再决定是否启用相对位置。

6.5 公开代码状态

GitHub README 标注代码库仍在 debugging/refactoring。虽然论文是 ICLR 2026 conference paper，代码也公开，但生产复现应固定 commit、记录 configs、保存 pretrained checkpoint config，而不是追随默认分支的最新状态。尤其是 README 中提到 config priority 为 CLI args > configs/testing.yaml > pretrained_path/config.yaml，如果只看当前 YAML 而不看 checkpoint 自带 config，可能误读实验参数。

7. 与相关工作的关系

Q-RAG 和 Beam Retriever 都是多步 retrieval，但训练信号和推理方式不同。Beam Retriever 训练 sequence scoring 或 beam search 风格的 trajectory model，HotPotQA in-domain 表现很强；Q-RAG 用 value function 训练 chunk-level policy，在 MuSiQue OOD 和长上下文任务上更稳。Q-RAG 和 Search-R1 都利用 RL 思想，但 Search-R1 更偏 LLM search/reasoning agent，成本和可扩展性取决于 LLM；Q-RAG 的 RL 发生在 embedding 空间，reader LLM 不参与策略训练。Q-RAG 和 GraphReader/HippoRAG 的差异更明显：后者构建图或知识结构，常常需要 LLM/复杂 pipeline 处理全文；Q-RAG 的表示更简单，牺牲一部分结构显式性，换取训练和推理效率。

与 long-context architecture（RMT、ARMT、Titans、Mamba2、LongRoPE2）相比，Q-RAG 不是扩展 LLM 的上下文机制，而是外置 retriever。它的优势是可以搭配任意 frozen reader LLM，而且 reader 成本不随原文长度增长；缺点是如果答案需要全局聚合、统计、风格模仿或非局部模式，而不是少量 supporting facts，检索压缩会丢信息。换言之，Q-RAG 适合“稀疏证据问答”，不一定适合“全文理解/生成”。

8. 复现与使用建议

如果要复现论文结果，我会按三层检查。第一层是数据和 chunking：HotPotQA chunk size 220、MuSiQue 110、BabiLong/RULER 64；chunk size 与最大检索步数共同决定 reader 输入的信息量。第二层是 reward：确认 GroundTruthFeedback 的 penalize_extra_steps、never_terminate 与论文表格一致，因为 Fact F1 对噪声非常敏感。第三层是 checkpoint config：不要只依赖默认 YAML，训练输出目录里的 config.yaml 才是精确配置；README 也说明了 config priority。

如果要把 Q-RAG 用到自己的长文档问答系统，可以先从离线监督开始：把历史 QA 的引用段落或人工标注证据转成 supporting fact indices，训练 Q-RAG agent；上线时缓存 action embeddings，只在线计算 state embedding 和 Q-values。若没有证据标注，可以先用强 reader/LLM judge 生成 weak supporting facts，但需要做噪声控制。对于多文档检索，可以先粗召回到文档/段落，再让 Q-RAG 在候选集合内做多步 trajectory；否则全库 chunk 级动作空间太大。

9. 读后总结

这篇论文的价值在于把“多步检索”从 LLM agent 叙事里抽出来，重新写成一个可以高效训练的 value function。它的设计非常贴合 long-context QA：长文档中真正需要读的通常是少数支持事实，而支持事实之间有依赖关系；用 state-dependent Q 值逐步找证据，比 single-step top-k 更自然。实验最强的证据来自 RULER/BabiLong 的超长上下文稳定性，以及 HotPotQA→MuSiQue 的 OOD 迁移。局限也清楚：需要 supporting-fact reward，主要适用于稀疏证据任务，真实系统还要解决缓存、chunking、弱监督和 reader failure。总体上，Q-RAG 是一个值得作为 long-context RAG retriever 训练基线的方法，尤其适合那些“证据少、上下文长、多跳依赖强”的 QA 场景。

Paper Notes

探索

Q-RAG: Long Context Multi-step Retrieval via Value-based Embedder Training

Q-RAG: Long Context Multi-step Retrieval via Value-based Embedder Training

1. 核心结论与定位

1.1 问题背景

1.2 主要贡献

2. 方法总览：把多步检索写成 RL 环境

2.1 MDP 定义

2.2 Q 函数：state embedder × action embedder

2.3 Soft Q 与 λ-return

2.4 Target network 与双头 Q

2.5 相对位置编码：为什么对 BabiLong/叙事任务重要

2.6 Plan Q-RAG

2.7 复杂度

3. 训练与推理流程

3.1 训练流程伪代码（按论文 + 代码对齐）

3.2 推理流程

3.3 训练细节

4. 实验结果

4.1 RULER：NIAH 近乎满分，QA 子任务更稳

4.2 Open-domain QA：HotPotQA in-domain 与 MuSiQue OOD

4.3 BabiLong：上下文到 10M tokens 的长期稳定性

4.4 Ablation：Soft-Q、target network、SFT baseline

4.5 Runtime / inference cost

5. 代码阅读与实现映射

5.1 代码与论文一致处

5.2 代码中需要注意的复现细节

5.3 每个关键组件的实现级伪代码

6. 局限性与风险

6.1 对 supporting-fact supervision 的依赖

6.2 检索成功不等于答案成功

6.3 线性扫描仍然有成本

6.4 位置编码可能任务相关

6.5 公开代码状态

7. 与相关工作的关系

8. 复现与使用建议

9. 读后总结

目录

2.3 Soft Q 与 $λ$ -return