PAIR: Prefix-Aware Internal Reward Model for Multi-turn Agent Optimization¶

会议: NeurIPS 2026
arXiv: 2605.17877
OpenReview: https://openreview.net/forum?id=0Ut88MX4Oj
代码: https://github.com/wonjoong-kim/PAIR
领域: 对齐/RLHF
关键词: 强化学习, 过程奖励模型, 智能体优化, 探针技术, 注意力模式

一句话总结¶

提出两阶段前缀感知内部奖励模型 PAIR，利用冻结隐状态探针估计信念一致性并结合多层注意力结构特征校正前缀污染偏差，在近零边际计算成本下为多轮智能体 GRPO 提供致密单步奖励。

研究背景与动机¶

在大语言模型智能体执行多步复杂任务（如工具调用与环境交互）时，强化学习特别是群体相对策略优化（GRPO）已成为主流训练手段。然而在多轮交互场景中，GRPO 高度依赖稀疏的结果奖励（仅在最终答案生成后提供单点标量），导致中间步骤的信用分配极其困难。训练信号无法精确辨别轨迹中哪个环节促成了成功或造成了失败，导致策略探索效率低下、收敛缓慢。现存的单步致密奖励方案均伴随沉重代价：树搜索回溯需要大量完整 Rollout，造成严重算力浪费；每步调用外部大模型充当裁判（LLM-as-a-judge）带来高昂的 API 延迟与调用开销；基于内在优势的似然评估则强制依赖运行时黄金标签；而外部过程奖励模型（PRM）也需专门维护独立网络。

通过探测模型自身内部表征提取正确性信号是一种兼具致密性与零额外调用开销的理想路径。然而，经典探针研究均默认输入前缀是干净的。在多轮智能体的强化学习探索中，策略不可避免会在中间步骤产生推理错误、工具误用或观测误读，导致后续步骤所依赖的前缀受到严重污染。深入分析表明，传统的隐状态线性探针在受污染前缀下性能急剧崩溃，因为自回归预训练目标促使隐状态追踪的是与前缀的「信念一致性」（belief-consistency）而非客观的「真实正确性」（grounded correctness）；相反，多层注意力分布作为结构性路由特征对前缀污染表现出高度鲁棒性，但在干净前缀下判别力略逊于隐状态。

为了打破这一两难困境，必须在保留隐状态对干净上下文敏锐判别力的同时，引入注意力特征纠正受污染上下文下的奖励反转偏差。核心 idea：构建解耦的两阶段奖励架构，先由干净前缀上预训练并冻结的隐状态探针输出信念一致性先验，再由轻量级多层注意力校正头动态预测真实正确性，并结合动量奖励整形机制消除组内方差折叠。

方法详解¶

整体框架¶

PAIR（Prefix-Aware Internal Reward Model）旨在完全依赖智能体策略模型自身的单次前向激活值，为多轮交互的每个步骤输出连续的正确性奖励 $r_t$。在执行多轮动作生成时，智能体策略 $\pi_\theta$ 在步骤 $t$ 根据前缀状态 $p_t = (x, t_1, a_1, o_1, \dots, o_{t-1})$ 自回归生成思考与工具调用动作。模型在前向传播的同时导出最后 token 的隐藏状态 $\mathbf{h}_t$ 以及跨层注意力统计特征 $\mathbf{a}_t$。

整个评估由两个轻量级逻辑回归阶段串联完成：Stage 1 的信念一致性估计器读取隐藏状态，输出置信度得分 $s_{bc}$；Stage 2 的注意力校正头接收拼接向量 $[\mathbf{a}_t; s_{bc}]$，输出对齐到真实客观正确性的最终得分 $s_{final}$。最后，通过温度剪裁与轨迹级滑动均值对比形成动量奖励，驱动 GRPO 策略参数更新。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["多轮智能体前向激活<br/>隐状态与多层注意力矩阵"] --> B["信念一致性估计器<br/>冻结线性探针计算一致性得分"]
    B --> C["注意力结构校正头<br/>融合注意力特征调整偏置"]
    C --> D["动量奖励整形<br/>温度剪裁与滑动均值对比"]
    D --> E["GRPO 策略更新<br/>组内相对优势信用分配"]

关键设计¶

1. 信念一致性估计器：解耦语言连贯性与客观任务正确性

隐状态探针在受污染前缀下往往将「顺着错误前缀继续胡说」判为高分，而将「纠正前期错误的修复步骤」判为低分。为了将这一固有表征先验显式剥离，PAIR 在 Stage 1 训练一个轻量级逻辑回归探针，直接将智能体生成动作末尾 token 的顶层隐藏状态 $\mathbf{h}_t \in \mathbb{R}^{d_{\text{model}}}$ 映射为标量一致性得分： $$s_{bc} = \sigma(\mathbf{w}_1^\top \mathbf{h}_t + b_1)$$ 该分类器在干净前缀数据上充分拟合后完全冻结参数 $(\mathbf{w}_1, b_1)$。它承担基线探测器的角色，准确捕获语义流利度与局部连贯性，为后续校正提供稳定的锚点。

2. 注意力结构校正头：利用路由特征纠偏受污染表征

注意力分布反映了模型在生成当前步骤时将信息从上下文何处汇聚路由，体现了计算的结构属性。即使在语义受污染的情境下，模型是否回溯关键任务指令和真实环境观测依然能清晰反映在注意力模式中。Stage 2 提取所有 Transformer 层中每个注意力头的 4 种统计标量：最大注意力峰值（max_attn）、分布离散度（std_attn）、前缀注意力占比（prefix_ratio）以及本轮自注意力占比（self_ratio），构成高维向量 $\mathbf{a}_t \in \mathbb{R}^{4 \times H \times L}$。校正头将注意力向量与冻结的标量得分 $s_{bc}$ 拼接后进行预测： $$s_{final} = \sigma(\mathbf{w}_2^\top [\mathbf{a}_t; s_{bc}] + b_2)$$ 在干净前缀下，注意力特征不引发显著偏移，使 $s_{final} \approx s_{bc}$ 自动保持高基线表现；而在受污染前缀下，注意力特征能敏锐捕捉到结构失真并纠正 $s_{bc}$ 的偏差，使得误导性延续被压低、纠错修复动作得到高额正向评分。

3. 动量奖励整形：缓解组内方差折叠并维持探索梯度

直接将探针概率值 $s_{final, t}$ 作为强化学习单步奖励会导致 GRPO 在同一 Prompt 下各轨迹得分趋同，引起组内奖励方差急剧缩减、策略梯度失效。为此，PAIR 引入温度剪裁与对称滑动动量奖励机制。首先对原始预测得分进行平滑与边界剪裁： $$\tilde{s}_{final, t} = \mathrm{clip}\left(\sigma\left(\frac{1}{T} \mathrm{logit}(s_{final, t})\right), \varepsilon, 1-\varepsilon\right)$$ 设定温度 $T=2$，阈值 $\varepsilon=0.05$。随后维护该轨迹历史步数的滑动均值 $\bar{s}_{<t}$（$t=1$ 时无偏置），并将当前步相对于历史均值的相对增益映射到 logit 空间计算单步奖励 $r_t$： $$r_t = \sigma\left(\mathrm{logit}(\tilde{s}_{final, t}) + \alpha \cdot (\tilde{s}_{final, t} - \bar{s}_{<t})\right)$$ 其中放缩系数 $\alpha=5$。当当前动作质量优于历史水平时获得正向 logit 奖励增强，反之予以惩罚。由于在 logit 空间相加后经过 Sigmoid 压缩，已确定高置信度的动作不会过度抖动，而处于决策模糊区（得分接近 0.5）的探索动作对相对进步极其敏感，从而保障了稳定的组内学习信号。

损失函数 / 训练策略¶

PAIR 探针采用离线单次监督训练方式，完全不依赖在线强化学习期间的反向传播： 1. 数据构建：在多轮基准上提取轨迹，对中间步骤注入推理错误、工具误用、观测误读及陈旧记忆四类污染，构造包含干净与受污染前缀且配有单步黄金标签的配对数据集。 2. 两阶段优化：Stage 1 与 Stage 2 均采用带 $\ell_2$ 正则化的逻辑回归（惩罚系数 $C=0.01$）优化二元交叉熵损失，输入特征经过 StandardScaler 标准化。 3. 在线推断部署：在 GRPO 训练循环中，策略模型开启 output_hidden_states=True 和 output_attentions=True。探针以前向推理副产物作为输入，计算单步奖励仅涉及两组向量内积与标量运算，单步计算延迟仅为 $\sim 0.2$ ms，实现零额外大模型调用的致密单步奖励。

实验关键数据¶

主实验¶

在代表真实复杂 API 调用的 GTA 基准（229 条复杂长程任务）与 ToolBench 基准（涵盖 RapidAPI 工具的 1000 条多轮轨迹子集）上，以 Qwen2.5-7B-Instruct 为策略模型进行 GRPO 训练评测。对比基线涵盖仅结果奖励（Outcome-only）、基于 GPT-4o-mini 的每步裁判（LLM-as-a-judge）、树搜索优势估计（Tree-GRPO、AT2PO）、内生优势估算（IGPO、TIPS）及常规单模态内部探针。

数据集	指标	本文 (PAIR)	之前SOTA (LLM-as-a-judge)	提升
GTA	任务成功率 (Success Rate)	0.2489 ± 0.042	0.2387 ± 0.042	+0.0102 (+4.27%)
ToolBench	任务成功率 (Success Rate)	0.4498 ± 0.025	0.4203 ± 0.028	+0.0295 (+7.02%)
GTA	预测正确性 AUROC	0.9217 ± 0.009	0.9082 ± 0.010 (Multi-Layer)	+0.0135
ToolBench	预测正确性 AUROC	0.8129 ± 0.006	0.7821 ± 0.007 (Multi-Layer)	+0.0308

消融实验¶

消融实验对比了 PAIR 架构各组成部分在两类多轮基准中的表现，包括特征解耦设计、未见污染类型泛化度以及动量奖励整形机制的有效性：

配置	GTA 成功率 / AUROC	ToolBench 成功率 / AUROC	说明
Full PAIR (带动量奖励)	0.2489 / 0.9217	0.4498 / 0.8129	完整架构，在受污染前缀与多轮 RL 中达到最高分
Vanilla PAIR (去除动量奖励)	0.2300 / 0.9217	0.4150 / 0.8129	直接使用 $s_{final}$ 作单步奖励，因组内方差退化掉点约 8%
仅隐藏状态探针 (Multi-Layer)	0.1789 / 0.9082	0.3856 / 0.7821	在受污染前缀下因一致性偏差导致奖励反转
仅注意力探针 (Multi-Attn)	0.2218 / 0.9034	0.3614 / 0.7765	干净前缀下判别力不足，影响初期策略收敛
简单拼接 (Hidden+Attn)	0.1432 / 0.8978	0.3784 / 0.7598	隐状态维度压制注意力信号，受污染前缀下性能急剧恶化
排除推理错误类型 (Leave-1-Out)	AUROC: 0.9056	-	泛化到未见污染模式，AUROC 下降控制在 0.025 以内

关键发现¶

隐状态探针在修复轮次表现出致命的「奖励反转」：在对抗性诊断测试集中，当智能体执行客观正确但违背错误前缀的修复步骤时，所有隐状态探针的 AUROC 跌至 0.5 随机概率以下，意味着其在 RL 训练中会主动惩罚自我纠错行为，反向鼓励将错误进行到底。
注意力路由特征具备长程鲁棒性：随着污染步骤与当前评估步的时间跨度 $d$ 扩大（$d=1 \to 7+$），隐状态探针准确率由 0.90 陡降至 0.60 左右；而注意力特征中的 prefix_ratio 和 self_ratio 能直接反映信息聚合来源，在长程污染下几乎不衰减。
动量机制是打通探针与 RL 优化的关键纽带：探针准确率高并不直接等价于 RL 策略收敛好。动量对比机制使得奖励对局部质量边际改进更为敏感，将成功率额外拉升约 8%。

亮点与洞察¶

从表征机制层面揭示了「信念一致性」与「客观正确性」的分离：语言模型的自回归预训练天然促使其隐状态契合上文逻辑，导致探针在受污染上下文中误把连贯性当作真理。这一机制性发现为分析大模型幻觉与自我纠错障碍提供了全新切入点。
两阶段级联架构优雅避免特征主导失衡：直接拼接隐藏状态与注意力特征会导致高维稠密的隐状态淹没稀疏结构特征。Stage 1 先提取一维标量一致性概率再由 Stage 2 校正，以结构约束保障了干净前缀不退化、污染前缀强纠偏。
极致的在线推断效率：利用生成过程中已计算的键值与激活副产物，单步计算仅需 0.2 毫秒，较每步调用外部大模型裁判提速超过 2,500 倍，打破了致密奖励无法低成本规模化应用于多轮强化学习的工程壁垒。

局限与展望¶

对极长上下文的表征衰减仍存瓶颈：尽管显著优于基线，但在交互轮数较多（$d \ge 7$）的长序列中，模型自身的全局注意力被稀释，注意力统计量的校正信噪比仍会出现边际下降。
依赖离线构建的合成错误分布：Stage 2 虽表现出跨错误类型的鲁棒泛化，但若智能体在探索过程中遭遇与四类典型错误机制截然不同的新型异常（如超长递归循环调用），探针的判别边界仍有待持续验证。
未来方向：探索动态自适应的在线探针微调，结合跨环境反馈闭环在线更新校正头，将方法推广至多模态具身智能体交互决策中。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次揭示前缀污染对隐状态探针的破坏机理，并巧妙解耦一致性与客观正确性。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖探针诊断、多轮工具基准、跨领域迁移测试及详细的消融分析。
写作质量: ⭐⭐⭐⭐⭐ 论证链条严密，动机驱动明确，从机制分析到方法落地环环相扣。
价值: ⭐⭐⭐⭐⭐ 极具工程落地与科研启发价值，为低成本推进智能体致密过程强化学习提供了通用范式。