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

TL;DR

The paper introduces PAIR, a prefix-aware internal reward model that provides dense, step-level rewards for multi-turn agent optimization, improving credit assignment without external calls or ground-truth data.

Contribution

It proposes a novel two-stage model combining hidden-state probing and attention mechanisms to generate robust internal rewards in multi-step tasks.

Findings

PAIR achieves the highest AUROC on contaminated trajectories.

Operates at negligible inference cost.

Enables dense step-level reward signals without external calls.

Abstract

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.