arXiv:2605.17877cs.AI2026-05被引 1

用内部状态探测提升多轮任务奖励信号,无需外部判断或完整回溯。

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

  • 设计双阶段模型,结合隐藏状态与注意力特征,动态修正错误感知。
  • 在污染轨迹上达到最高AUROC(0.92),且推理开销极低。
  • 适合需要密集中间奖励的复杂多步任务优化,如智能体训练。

当前大模型在执行复杂多阶段任务时面临挑战。尽管群体相对策略优化(GRPO)表现优异,但其依赖稀疏结果奖励,难以对中间步骤进行有效信用分配。现有方法如全轨迹回溯、每步调用外部大模型评判或需真实答案的内在奖励,均存在成本高或实际限制。本文提出假设:可利用大模型隐状态中的内部正确性探测作为步骤级奖励信号,解决上述问题。然而,现有探测方法假设输入干净,在多步场景下失效——隐状态探测受前缀污染影响,误追踪上下文连贯性而非真实正确性;而基于注意力的特征虽对污染鲁棒,但在干净前缀上表现不佳。基于此互补性,我们提出前缀感知内部奖励模型(PAIR),由冻结的隐状态探针估计信念一致性,搭配轻量注意力头纠正至真实正确性。实验表明,PAIR在污染轨迹上达到最高AUROC(0.92),推理开销可忽略,支持无需外部调用、无需真值依赖、无需全轨迹回溯的密集步骤奖励,显著提升GRPO训练效率。

原文摘要 · Abstract (English)

A significant hurdle for current LLMs is the execution of complex, multi-stage tasks. Group Relative Policy Optimization (GRPO) has been emerging as a leading choice, but its reliance on sparse outcome rewards severely limits credit assignment across intermediate steps. Existing remedies such as running full rollouts to assign step-level advantages, calling external LLM judges at each step, or computing intrinsic rewards that require ground-truth answers at every evaluation introduce significant costs or practical constraints. We hypothesize that internal correctness probing over LLM hidden states can be repurposed as a step-level reward signal, potentially addressing all of these limitations at once. However, existing probing research assumes clean inputs, and we first show that this assumption breaks down in multi-step settings: hidden-state probes degrade severely under prefix contamination tracking coherence with the (possibly corrupted) prefix rather than grounded correctness, while attention-based features remain robust to contamination but underperform on clean prefixes. Building on this complementary relationship, we propose the Prefix-Aware Internal Reward (PAIR), a two-stage model with a frozen hidden-state probe estimating belief-consistency and a lightweight attention-based head correcting it toward grounded correctness. Experimental results show that PAIR achieves the highest AUROC on contaminated trajectories while operating at negligible inference cost, enabling dense step-level reward signals for GRPO training without external model calls, ground-truth dependencies, or full-trajectory rollouts.

多轮任务奖励建模内部奖励智能体优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。