💻 computer science

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

The paper proposes PAIR, a prefix-aware internal reward model that combines a frozen hidden-state probe with a lightweight attention-based head to generate dense, low-cost step-level rewards for multi-turn agent optimization, effectively overcoming the limitations of sparse outcome rewards and the degradation of probes under prefix contamination.

Original authors: Wonjoong Kim, Yeonjun In, Sangwu Park, Dongha Lee, Chanyoung Park

Published 2026-05-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Wonjoong Kim, Yeonjun In, Sangwu Park, Dongha Lee, Chanyoung Park

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot assistant to solve a complex puzzle, like booking a flight, reserving a hotel, and ordering dinner all in one go. The robot has to take many steps, making decisions and using tools along the way.

The big problem with current teaching methods is that the robot only gets a "Good job!" or "Try again" at the very end of the whole process. If the robot messes up step 1, it doesn't know it until it finishes step 10 and fails the whole task. It's like playing a video game where you only get a "Game Over" screen at the end, with no hint about which jump you missed.

This paper introduces a new way to teach these robots called PAIR (Prefix-Aware Internal Reward Model). Here is how it works, using simple analogies:

The Problem: The "Confused Detective"

To give the robot feedback during the process (step-by-step), researchers tried looking inside the robot's brain (its internal computer states) to see if it was doing the right thing.

They found two types of "brain signals":

The "Storyteller" (Hidden States): This signal checks if the current step makes sense with what happened before. It's like a detective who only cares if the current clue fits the story they've already written.
- The Flaw: If the robot made a mistake earlier (contaminated the story), the "Storyteller" gets confused. It thinks, "Oh, this new step fits the wrong story perfectly, so it must be good!" It accidentally rewards the robot for continuing down a wrong path.
The "Structural Architect" (Attention Patterns): This signal looks at how the robot is focusing its attention. It doesn't care about the story; it cares about the structure. Is the robot looking at the right tools? Is it ignoring the irrelevant noise?
- The Strength: Even if the story is messed up, this signal can still tell if the robot is looking at the right things. It's robust.
- The Weakness: On a perfect, clean story, it's not as sharp as the "Storyteller."

The Solution: The "Two-Stage Coach" (PAIR)

The authors realized that neither signal is perfect on its own. The "Storyteller" is great when things are going well, but fails when mistakes happen. The "Architect" is steady but less precise on clean tasks.

So, they built PAIR, a two-stage coach:

Stage 1: The Belief Check. The coach first asks the "Storyteller" (the hidden state probe): "Does this step fit the current story?" This gives a quick, initial score.
Stage 2: The Reality Check. The coach then asks the "Architect" (the attention probe): "Wait, looking at how you're focusing, are you actually solving the problem, or just following a broken story?"

How they work together:

If the story is clean: The "Architect" agrees with the "Storyteller." The coach accepts the score.
If the story is broken (contaminated): The "Storyteller" says, "This looks good because it fits the mistake!" But the "Architect" says, "No, you're looking at the wrong things!" The coach listens to the "Architect" and corrects the score, telling the robot, "Actually, that step was bad, even if it felt right."

Why This is a Big Deal

Before PAIR, to get step-by-step feedback, you had to do one of three expensive things:

Run the whole game 100 times to see what works (very slow and wasteful).
Call a human or a super-smart AI to judge every single step (very expensive and slow).
Know the exact answer beforehand to grade the steps (impossible for many real-world tasks).

PAIR changes the game because:

It's Free: It uses the robot's own brain signals. It doesn't need to call anyone else or run extra simulations.
It's Instant: It happens in the blink of an eye while the robot is thinking.
It's Honest: It catches mistakes even when the robot is confidently following a wrong path.

The Result

When they tested this on robots trying to use real-world tools (like booking flights or searching databases), PAIR helped them learn much faster and solve more problems than any other method. It's like giving a student a teacher who can spot a mistake while they are writing the essay, rather than waiting until the essay is finished to say, "You failed."

In short, PAIR is a smart, internal "lie detector" for AI agents that helps them correct their own mistakes in real-time, without needing expensive external help.

Technical Summary: PAIR (Prefix-Aware Internal Reward Model for Multi-Turn Agent Optimization)

1. Problem Statement

Current Large Language Models (LLMs) struggle with complex, multi-stage agent tasks requiring sequences of actions (e.g., API calls, tool invocations) and environmental feedback. While Group Relative Policy Optimization (GRPO) is a leading reinforcement learning (RL) algorithm for these tasks due to its simplicity and efficiency, it suffers from sparse outcome rewards. In standard GRPO, a single scalar reward is assigned only upon trajectory completion, making it difficult to perform credit assignment across intermediate steps. This sparsity leads to slow convergence and poor sample efficiency.

Existing attempts to densify rewards face significant limitations:

Full Rollouts: Propagating outcome rewards backward requires completing full trajectories, generating wasteful rollouts from already-failed paths.
External Judges: Using external LLMs (e.g., GPT-4) to evaluate each step incurs high API costs and latency, rendering it impractical for large-scale online RL.
Intrinsic Rewards: Methods relying on ground-truth answers (e.g., teacher-forced log-probabilities) require runtime access to gold labels, which is often unavailable or operationally costly.
External Reward Models: Separately trained Process Reward Models (PRMs) introduce overhead in training and maintenance.

The core challenge addressed by this work is: Can we generate dense, step-level reward signals using only the agent's own internal states, without external models, ground-truth dependencies, or full-trajectory rollouts?

2. Methodology

2.1 Empirical Motivation: The Failure of Naïve Probing

The authors first investigate whether internal correctness probes (trained to predict step correctness from hidden states) can serve as rewards. They introduce the concept of prefix contamination, where an agent generates an erroneous intermediate step, corrupting the prefix for subsequent steps.

Through systematic analysis on the Qwen-2.5-7B-Instruct backbone, they uncover three critical findings:

Hidden-State Degradation: Hidden-state probes (e.g., Last-Token, Mean-Pooled) degrade severely under prefix contamination. They track belief-consistency (coherence with the corrupted prefix) rather than grounded correctness (actual task utility). Consequently, they often reward error-perpetuating turns and penalize repair turns.
Attention Robustness: Attention-based features remain relatively robust to contamination. They capture structural behavior patterns (e.g., whether the model attends to task-relevant tokens) rather than semantic coherence with the prefix.
Complementarity: Hidden states excel on clean prefixes but fail on contaminated ones; attention features are robust but underperform on clean prefixes. A naïve concatenation of both signals fails because the inverted signal from the hidden state cancels out the correct signal from attention.

2.2 Proposed Method: PAIR

To leverage the complementary nature of these signals, the authors propose PAIR (Prefix-Aware Internal Reward Model), a two-stage architecture:

Stage 1: Belief-Consistency Estimator (Frozen): A pre-trained logistic regression probe maps hidden-state features (specifically the last-token hidden state) to a belief-consistency score ( $s_{bc}$ ). This probe is trained on clean data and frozen during RL training. It captures the standard probe's ability to detect correctness when the prefix is clean.
Stage 2: Attention-Based Correction: A lightweight logistic regression head takes the attention features (statistics such as peak magnitude, spread, prefix ratio, and self-ratio) concatenated with $s_{bc}$ $s_{b c}$ as input. It predicts a final score ( $s_{final}$ $s_{f ina l}$ ).
- Mechanism: On clean prefixes, the correction head learns to output $s_{final} \approx s_{bc}$ . On contaminated prefixes, the attention features detect structural anomalies, and the head learns to adjust $s_{final}$ away from the biased $s_{bc}$ toward grounded correctness.
Reward Integration: The final score is transformed via a temperature-clipped sigmoid and combined with a momentum-based reward mechanism. This contrasts the current step's score against the trajectory's running mean to recover group-relative variance, which is essential for GRPO advantage estimation.

3. Key Contributions

Empirical Discovery: The paper provides the first systematic study demonstrating that internal correctness probes fail under prefix contamination in multi-step settings, specifically tracking belief-consistency rather than grounded correctness.
Mechanistic Insight: It establishes that hidden states encode belief-consistency while attention patterns encode structural routing behavior. These signals are complementary, explaining why simple concatenation fails and guiding the design of adaptive architectures.
Methodological Innovation: PAIR is proposed as a two-stage model that adaptively combines hidden-state and attention signals. It preserves high performance on clean trajectories while correcting for bias under contamination.
Practical Efficiency: PAIR operates without external LLM calls, runtime ground-truth access, or separate reward model training. It incurs negligible inference cost (probe-level) by reusing activations already computed during the policy's forward pass.

4. Experimental Results

The authors evaluate PAIR on two multi-step agent benchmarks: GTA (General Tool Agents) and ToolBench, using Qwen-2.5-7B-Instruct as the policy model.

Step-Level Prediction: PAIR achieves the highest AUROC (0.9217 on GTA, 0.8129 on ToolBench) among all internal probing baselines, outperforming hidden-state probes, attention probes, and unsupervised methods (e.g., Lookback Lens, Head Entropy).
RL Performance: In GRPO training, PAIR achieves the highest success rates on both datasets (0.2489 on GTA, 0.4498 on ToolBench), outperforming the second-best baseline (LLM-as-a-Judge) on ToolBench and matching it on GTA.
Cross-Domain Transfer: When trained on GTA and evaluated on HotpotQA (without fine-tuning), PAIR ranks first among 16 methods, demonstrating superior transferability compared to rollout-based or ground-truth-dependent methods.
Efficiency: PAIR operates at probe-level latency (~0.2 ms/step), which is approximately 2,500× faster than LLM-as-a-Judge API calls and orders of magnitude faster than tree-search methods.

5. Significance and Claims

The paper claims that PAIR addresses the four primary limitations of existing dense reward mechanisms simultaneously:

No Full Rollouts: Rewards are computed per step without waiting for trajectory completion.
No External Calls: It relies solely on the agent's internal representations, eliminating API costs and latency.
No Ground-Truth Dependency: The probe is trained offline and applied at runtime without needing gold answers for every step.
No Separate Reward Model: It does not require training a separate, large-scale reward network.

The authors position PAIR as a practical solution for enabling dense, step-level reward signals in multi-turn agent RL training. By resolving the failure mode of hidden-state probes under prefix contamination through an adaptive, attention-corrected architecture, PAIR makes internal state analysis a viable and efficient alternative to external reward sources for optimizing complex agent behaviors.