← 返回 2026-04-13

过程奖励智能体:用于引导知识密集型推理 Process Reward Agents for Steering Knowledge-Intensive Reasoning

Jiwoong Sohn, Tomasz Sternal, Kenneth Styppa, Torsten Hoefler, Michael Moor 📅 2026-04-10 👍 6 2026-07-13 08:36
医疗推理 推理时搜索 束搜索 检索增强生成 过程奖励模型

用在线过程奖励智能体在推理时实时评分并触发检索,提升医疗推理。

前置知识

过程奖励模型 (PRM)

对推理轨迹中每一个中间步骤而非仅最终答案打分的奖励模型,常用于 step-level 监督和树搜索引导。

本文的核心创新是对 PRM 的重塑:从事后打分改为在线、可触发的实时控制,所以必须先理解传统 PRM 的定位。

链式思考与自洽性 (CoT + Self-Consistency)

CoT 让模型逐步推理;Self-Consistency 采样多条轨迹并以多数投票决定答案,是本文最主要的对比基线。

PRA 用束搜索替代自洽性的简单投票,所以需要理解它要改进的对照方法。

束搜索与推理时扩展 (Beam Search / Inference-Time Scaling)

推理时不更新参数而增加算力,束搜索在每步保留 top-B 个候选轨迹,是 PRA 实施过程奖励引导的载体。

PRA 的整个推理管线建立在束搜索 (B=4, b=16) 上,缺此概念难以理解其树剪枝逻辑。

检索增强生成 (RAG) 与条件检索

RAG 在生成前从外部知识库检索相关文档并拼接到上下文;条件检索则是由模型按当前步骤需求动态决定是否触发,而非每步都检索。

PRA 把检索从策略模型解耦到奖励智能体,并通过二分类 head 决定是否触发,是论文核心机制。

研究动机

医疗等知识密集领域的推理很难可靠化:每一步都涉及多源证据(论文、指南、机构协议)的综合,中间步骤不像数学或代码那样可机械验证。现有 PRM(如 Med-PRM、Med-S3)都是事后 (post hoc) 打分,等整条轨迹生成完才介入,导致错误在中段累积却无法实时纠正;标准 RAG 把文档一股脑塞进上下文,既不能保证模型在正确时刻关注正确证据,也提供不了逐步纠错机制。Self-Consistency 在 GPQA、Lancet 等难题上甚至会因为多数错误投票而掉点。结果就是小模型 (≤4B) 在 MedQA 上的能力被严重低估。

本文的目标是构建一个能在推理时实时介入、为冻结的策略模型 $\pi$ 提供按步奖励的代理 $G_\phi$,使小模型也能在多个医疗基准上达到新 SOTA,并在不重训策略的前提下泛化到不同 backbone。具体目标包括:在 Qwen3-4B-Instruct 上把 MedQA 准确率推到 81.9%,在 6 个 OOD 基准平均提升 5.2 分,并在 0.5B–8B 的未见策略上保持提升,证明过程奖励可以独立于策略部署。

与已有工作不同的是,现有方法要么把奖励当作事后评估器 (Med-PRM/Med-S3),要么把检索当作固定管道 (RAG),都没有把检索与验证合并为一个可学习的在线控制器。本文把检索和打分都委托给一个独立的 Process Reward Agent (PRA),用单一模型的两个 token-level readout 同时决定「搜不搜」与「这步对不对」,从而首次实现检索-验证-剪枝的闭环实时控制。策略模型完全不读检索文档,保证 PRA 可独立重训并在 backbone 之间即插即用。

核心方法

PRA 把推理建模成在线控制问题 $G_\phi: (q, \pi, D) \mapsto \hat{y}_q$。冻结的策略 $\pi$ 每生成一个候选 step,PRA 的两个 head 立即工作:动作头 $\mu_\phi^{act}$ 决定要不要调用检索器 $\rho$,奖励头 $\mu_\phi^{rwd}$ 在 (可能含检索文档) 的条件下为该步打分。整条轨迹放进束搜索 (beam width $B=4$,branching factor $b=16$,共 64 候选) 里,每步按累计奖励 $\sum_t \hat{r}_t^{(j)}$ 排序剪枝。三方模型 (策略/PRA/检索器) 通过一个全局队列做 stage-level 批处理,把不同问题、不同束、不同深度的轨迹混批跑,保持 GPU 高利用率。

把「搜不搜」和「这步对不对」合并成同一冻结参数的 PRA 模型的两个 token-level 输出,从而把检索、验证、剪枝三件事耦合成一个可微调模块;策略模型完全冻结、永远不读检索文档,因此 PRA 可以脱离策略独立重训并在不同 backbone 间即插即用。本质区别于 Med-PRM/Med-S3:那些方法只做事后打分,本文做的是 step-wise 在线控制;本质区别于标准 RAG:本文用 $\Delta m$ 边际位移学会有选择地触发检索,而非每步都检索。

方法步骤详情

训练阶段:用 Qwen3-235B-Instruct 做教师,在 MedQA 训练集 10,178 题上各采 8 条 Qwen3-4B 轨迹,对每个中间 step 分别在「有/无检索文档」条件下让教师输出 log p(0), log p(1);定义边际 $\Delta m = m - m_d$ 作为检索必要性标签,$\epsilon_{global}$ 取训练集 $|\Delta m|$ 的中位数 (50% step 标为需检索);推理标签取 argmax 的 0/1 单 token。微调 Qwen3-4B-Instruct 在两个输出槽分别预测检索动作和推理奖励。推理阶段:所有问题进入全局队列,按 stage (策略生成/检索/PRA 评分) 批处理;每个 beam 维护 B 条 partial trace,每步用策略扩展 b 个候选得 $B \times b$ 个新 trace,PRA 对每个同时采样动作与奖励;检索动作 1 时用 MedCPT 在 (问题+最近两步) 上召回 200 候选并 rerank 保留 top 64;按累计奖励保留 top-B,其余剪枝。

技术新颖性

三点技术新颖性:(1) 同一模型用两个输出 slot 同时做动作决策与奖励打分,避免两套独立模型的开销与不一致;(2) 引入「边际位移 $\Delta m$」作为检索必要性监督信号,让 PRA 学会有选择地触发检索,论文图 3 给出 Pareto 前沿;(3) stage-level 全局批处理把不同问题不同束的 trace 混合跑,论文报告在可变长推理与早期终止下仍维持高 GPU 利用率。整体上把 PRM 从「事后评估器」重定义为「在线控制器」。

Overview of our approach. A Process Reward Agent (PRA) observes the reasoning trace generated by the frozen policy model (reasoner), decides when to search for external evidence, and assigns step-level rewards.
Figure 1: Overview of our approach. A Process Reward Agent (PRA) observes the reasoning trace generated by the frozen policy model (reasoner), decides when to search for external evidence, and assigns step-level rewards.
Search–accuracy trade-off on MedQA. We sweep the search threshold and report accuracy versus search frequency; the Pareto frontier highlights the best operating points for a given search budget.
Figure 3: Search–accuracy trade-off on MedQA. We sweep the search threshold and report accuracy versus search frequency; the Pareto frontier highlights the best operating points for a given search budget.

实验结果

主表 (Table 1) 显示在 Qwen3-4B-Instruct 上 PRA 平均 72.1 分,比最强基线 RAG+SC (66.9) 高 5.2 分,并以 81.9% 成为首个突破 80% 的 4B 模型。Figure 2 推理时扩展曲线显示 PRA 从 1 样本到 64 样本持续上升,Self-Consistency 在 8 样本后基本饱和。Table 2 跨模型泛化表明 PRA 仅用 Qwen3-4B 训练就能即插即用到 Llama-3.1-8B (+15.3)、Llama-3.2-3B (+23.1)、Qwen2.5-3B (+25.4)、Llama-3.2-1B (+25.0)、Qwen2.5-0.5B (+25.7)。Table 3 消融证明训练奖励智能体是主要收益。Table 4 显示 post-hoc 用 PRA 打分 (无论 outcome / min / max / avg) 都明显弱于在线 process-level 控制 (74.3–77.6 vs 81.9)。Figure 4/5 边际位移分析验证 PRA 学到「真证据」。

Main results on medical reasoning benchmarks. PRA outperforms direct answering, chain-of-thought (CoT), and retrieval-augmented generation (RAG) baselines on MedQA and six OOD benchmarks.
Table 1: Main results on medical reasoning benchmarks. PRA outperforms direct answering, chain-of-thought (CoT), and retrieval-augmented generation (RAG) baselines on MedQA and six OOD benchmarks.
Performance on MedQA under inference time scaling. PRA continues to benefit from additional compute, while Self-Consistency saturates early.
Figure 2: Performance on MedQA under inference time scaling. PRA continues to benefit from additional compute, while Self-Consistency saturates early.
Margin shift across step positions in reasoning trajectories on MedQA, separated by traces with correct and incorrect final answers.
Figure 4: Margin shift across step positions in reasoning trajectories on MedQA, separated by traces with correct and incorrect final answers.
Mean absolute margin shift over reasoning steps across questions of varying difficulty.
Figure 5: Mean absolute margin shift over reasoning steps across questions of varying difficulty.
查看结构化数据
任务指标本文基线提升
MedQA (ID) Accuracy 81.9% RAG + SC 76.7% +5.2
Medbullets (OOD) Accuracy 65.9% CoT + SC 58.7% +7.2
MedMCQA (OOD) Accuracy 66.2% RAG + SC 64.8% +1.4
MMLU-Med (OOD) Accuracy 86.6% RAG + SC 86.2% +0.4
GPQA (OOD) Accuracy 65.1% CoT 60.8% +4.3
Lancet (OOD) Accuracy 70.9% CoT + SC 63.5% +7.4
NEJM (OOD) Accuracy 68.0% RAG + SC 66.9% +1.1
Qwen2.5-0.5B-Instruct (Cross-Model) MedQA Accuracy 54.1% CoT 28.4% +25.7 (相对 +90.5%)

局限与改进

作者在 Impact Statement 中明确指出 PRA 不应被视为可部署的医疗决策系统,仍可能产生幻觉和错误中间步骤。我观察到的额外局限包括:(1) 实验仅在 Qwen3-4B 单一 backbone 上训练 PRA,泛化虽强但没有覆盖非 Qwen 系 70B+ 级别大模型;(2) 训练 teacher (Qwen3-235B) 自身就有事实性偏差,边际位移标签可能继承教师错误;(3) 检索固定为 MedCPT + top 64,未与 BM25、稀疏-稠密混合等更强检索器对比;(4) 64 条候选/题的计算成本远高于单次 SC,部署时延与显存开销未量化;(5) 知识库停留在医学教科书+StatPearls+指南+罕见病语料,未纳入最新论文,因此 OOD 表现仍有上限。

独立分析的弱点

独立分析三个主要弱点:(1) 双头耦合的设计虽然高效,但「检索必要性」与「推理正确性」是两个本质不同的任务,共享参数可能互相干扰;改进方向是解耦为两独立 LoRA head 或两独立小模型。(2) 边际位移 $\Delta m$ 作为二值检索标签会丢失幅度信息,应改为软标签或排序损失;目前 50% 强制搜索率过于粗粒度。(3) Table 3 中未训练 PRA+检索就能与 RAG+SC 打平,说明检索收益与策略无关,论文未对此做对照实验,难以判断「训练过的 PRA」与「未训练 PRA」的差距究竟来自奖励学习还是仅仅来自束搜索结构。改进方向是加入 ablation:相同束搜索结构 + 随机奖励。

未来方向

作者在附录与正文中暗示的方向:(1) 把 PRA 推广到非医疗知识密集领域 (法律、金融);(2) 通过软检索概率或检索文档压缩进一步降低每题成本;(3) 用 RL 训练 PRA 而非监督标签。我建议的延伸:(1) 在多模态 (影像+文本) 推理上替换 MedCPT 为跨模态检索;(2) 把「PRA 即插即用」拓展到推理时可热替换的 Agent 协议,实现按硬件预算动态切换 B 和 b;(3) 引入因果反事实评估,验证 PRA 选择的 step 是否真的「换证据就改答案」,以减少对教师标签的依赖。

复现评估

论文明确承诺开源代码与数据 (process-reward-agents.github.io),附录 C 提供训练超参、附录 D 给出 prompt 模板,使用公开 MedQA、MedMCQA、Medbullets、MMLU-Med、GPQA、Lancet、NEJM 等基准。训练侧 teacher Qwen3-235B 与策略 Qwen3-4B 均为开源权重,MedCPT 也是公开检索器。复现门槛主要在算力:每次评测要跑 64 条候选/题 × 7 个 benchmark × 6 个 OOD policy,加 stage-level 批处理虽提升 GPU 利用率但仍需多卡 A100/H100;PRA 微调数据需先用 235B 教师在 10,178 题上各采 8 条并前向两次 (有/无检索),算力开销显著。中等偏难复现。