一文读懂 RAG-Gym：用过程监督优化推理与搜索智能体 精华

传统的检索增强生成（RAG）架构依赖于静态检索，这使得它们在处理需要顺序信息搜索的复杂问题时，效果大打折扣。虽然智能推理和搜索提供了一种更具适应性的方法，但现有的大多数方法在很大程度上都依赖于提示工程。

为了解决这一问题，本文介绍了 RAG-Gym，这是一个统一的优化框架，它通过在每个搜索步骤进行细粒度的过程监督，来提升信息搜索智能体的性能。

核心贡献

• 引入 RAG-Gym，这是一个利用过程监督优化智能 RAG 的统一框架。
• 提出 ReSearch，这是一种全新的智能体架构，它将答案推理和搜索协同起来，比现有的基线模型取得了更优异的性能。
• 证明了使用经过训练的过程奖励模型作为验证器，可以显著提升搜索智能体的性能。
• 全面分析了过程监督的来源、奖励模型的可迁移性，以及智能 RAG 中的缩放定律。​

RAG-Gym 框架

一、概述

RAG-Gym 将知识密集型问答任务构建为一个嵌套的马尔可夫决策过程（MDP）。

过程奖励数据是通过在每个时间步随机抽样动作候选，并使用外部注释器选择最佳动作来收集的。

RAG-Gym 中实现了不同的过程监督方法。

二、知识密集型问答作为嵌套 MDP

1. 外部 MDP 的构建：旨在归纳各种针对知识密集型问题的语言智能体设计，具体过程如下：

• 状态空间 S：在每个时间步 t，状态 st ∈ S 由原始问题 Q 和信息搜索历史 Ht 组成，表示为：st = (Q,Ht) ，其中 Ht = {(q1,D1), · · · , (qt−1,Dt−1)}是信息搜索查询 q1, · · · , qt−1 以及它们对应的由环境返回的检索文档 D1, · · · , Dt−1 的历史记录。状态空间 S 是所有可能状态的集合：

其中 I 是问题空间，Aq 是所有可能搜索查询的空间，D 是所有文档的空间。 - 动作空间 A：在每个时间步 t，动作 at 可以是一个搜索查询，也可以是对原始问题的预测答案。因此，我们将动作空间 A 定义为 A = Aq ∪ Ap ，其中 Aq 是所有可能查询的集合，Ap 表示可能答案的集合。 - IR 环境：RAG-Gym 中外部 MDP 的环境由一个信息检索（IR）系统驱动，该系统将搜索查询 qt 作为输入，并返回一组相关文档 Dt 作为输出。IR 系统可以表示为从 Aq 到 P(D)的映射，其中 P(D)是 D 的幂集。检索过程由底层文本检索器和特定的检索设置（例如返回文档的数量）决定。 - MDP 工作流程：对于给定的问题 Q，MDP 从初始状态 s1 = (Q, ∅)开始。在每个步骤 t，at 是从智能体的策略 πf(θ)(· | st)中采样得到的，其中 πf(θ) : S → Δ(A)定义了给定状态下的动作分布。智能体的策略由 θ（表示基础语言模型的参数）和特定于智能体的函数 f（表示如何利用基础 LLM 策略）组成。然后，通过添加(qt,Dt)来更新历史记录为 Ht+1 ，状态转换为 st+1 = (Q,Ht+1) 。否则，如果 at ∈ Ap ，则认为当前情节完成，MDP 终止。 - 奖励：对于外部 MDP，一个情节的奖励由最终预测的正确性决定。状态-动作对(st, at)的即时奖励为：

外部 MDP 中智能体的目标是最大化轨迹上的预期累积奖励：

• 通过过程监督改进搜索智能体：通过纳入过程奖励，RAG-Gym 能够更有效地对大语言模型（LLM）进行微调，使标记生成与高质量的搜索行为保持一致。 -过程奖励数据的收集：数据收集流程从轨迹采样开始，语言智能体根据其当前策略生成一系列动作。在轨迹的每个步骤中，会提出多个候选动作，并根据预定义的评估标准选择最佳动作。我们采用基于排名的评估框架，而不是分配数值分数，以确保一致性。然后执行选定的动作，轨迹转移到下一个状态。这个过程反复迭代，直到轨迹终止。为了确保质量，只有最终答案正确的轨迹才会被保留，这由结果奖励来判断。 -使用过程监督调整智能体-监督微调（SFT）：从过程奖励中选择的动作用于训练语言智能体。正式地说，SFT 的目标是最小化给定状态下选定动作的负对数似然：

  

• 其中 D 是带有过程奖励标记的状态-动作对的数据集。 -直接偏好优化（DPO）：引入了一个对比学习框架，该框架同时考虑了选定和未选定的动作。过程奖励数据被重新组织为偏好对(a+t , a−t )，其中 a+t 是对于 st 更偏好的动作，a−t 是较不偏好的替代动作。DPO 的目标是最小化以下损失：



• -过程奖励建模（PRM）：训练一个单独的奖励模型 rϕ(st, at)，以根据收集的数据预测过程奖励。目标是最小化一个对比损失，该损失评估更偏好动作相对于较不偏好动作的质量：

推理与搜索（ReSearch）智能体

推理与搜索（ReSearch）智能体在一个统一的、以答案为驱动的框架中集成了推理和搜索。

• 历史知识总结：给定状态 st ，其中包含原始问题 Q 和历史 Ht = {(q1,D1), . . . , (qt−1,Dt−1)}，智能体首先会将检索到的文档总结为对相应查询的结构化响应，形成一个精炼的知识表示 H′t： H′t = {(q1,m1), . . . , (qt−1,mt−1)} 总结步骤会过滤掉不相关的信息，并减轻长上下文处理的挑战，使智能体在构建答案时能够专注于最相关的事实。
• 答案推理：利用这个精炼的知识 H′t ，智能体接着进行结构化推理，以推断出问题的候选答案。然后，它会检查推理步骤，并判断所有的主张是否都在历史记录中有充分的依据。如果智能体确定答案推理中的所有主张都有检索到的证据支持，它就会将答案作为最终动作输出。否则，它会识别出未经验证的主张，即那些根据现有证据缺乏充分理由的陈述。
• 搜索查询生成：未经验证的主张作为生成下一个搜索查询的基础，这个查询专门用于检索缺失的信息。从这个查询中检索到的文档随后会被添加到 Ht 中，推理过程会反复迭代，直到所有主张都得到验证，或者检索预算用尽。

实验结果

• 过程监督方法的比较：下表展示了使用 Llama-3.1–8B-Instruct 实现的各种智能体，以及它们在 RAG-Gym 中使用不同过程监督方法进行微调后的性能。与零样本学习（ZSL）基线相比，过程监督在所有智能体上都一致地提高了性能，证明了其在增强中间推理和查询生成方面的有效性。在三种过程监督算法中，PRM 总体上取得了最好的结果，比 ZSL 基线的平均 F1 得分高出 25.6%（ReAct）。

  

• ReSearch 与其他智能体的比较：结果还表明，在零样本学习设置和有过程监督的设置中，ReSearch 始终优于其他智能体。在未经调优的情况下，ReSearch 就实现了很强的零样本性能，证明了明确地将答案推理与查询生成对齐的有效性。使用过程奖励模型后，ReSearch 达到了最先进的性能，在不同数据集上的平均 EM 得分达到 54.31%，平均 F1 得分达到 62.41%。
• 奖励模型的可迁移性：下图突出显示了使用基于 Llama-3.1–8B 的过程奖励模型的 GPT-4o-mini 的 ReSearch 智能体的性能提升。使用奖励模型进行动作选择在所有任务中都带来了一致的收益，证明了 PRM 在不同大语言模型中有效选择高质量动作的可迁移性。

分析

• 不同奖励来源的比较：四位领域专家对 200 个 MedQA 问题进行了注释。在其余 800 个用 GPT-4o 注释的训练问题上训练了一个奖励模型，并将其偏好与领域专家的偏好进行了比较。下表显示了领域专家偏好与不同来源奖励估计之间的一致性。使用 GPT-4o 注释训练的奖励模型与人类偏好的一致性最高（85.85%），显著优于 Math-Shepherd 中引入的基于滚动的方法（71.03%）。这表明在这种情况下，GPT-4o 注释与人类推理和决策紧密对齐。

• 训练时间缩放：下图展示了使用在不同数量训练样本上微调的过程奖励模型的 ReSearch 智能体的性能。结果表明，ReSearch 的性能随着训练样本数量的增加而提高，但随着样本量的增长，收益趋于收敛。

• 推理时间缩放：以 ReSearch 为测试智能体，下图展示了我们的推理时间缩放研究结果。在多个基准测试中观察到一致的趋势，即增加采样动作的数量通常会提高性能。

研究结论

本文介绍了 RAG-Gym，这是一个通过过程监督优化推理和搜索智能体的框架，并引入了 ReSearch，这是一种将答案推理与搜索查询生成统一起来的智能体架构。实验表明，RAG-Gym 在知识密集型任务上改进了搜索智能体，ReSearch 始终优于基线模型。研究还证明了使用大语言模型作为过程奖励评判的有效性、训练后的奖励模型在不同大语言模型上的可迁移性，以及 ReSearch 在训练和推理过程中的缩放模式。

本文转载自 ​​柏企科技圈​​，作者：柏企