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

arXiv: https://arxiv.org/abs/2605.17877 · HTML: https://mlsanity.com/paper/2605.17877/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Reinforcement learning for reasoning, one of 17 directions in LLM reasoning.
- Builds on 29 of its 34 references that are in the map.
- First posted on arXiv 18 May 2026.

## Placed in

- LLM reasoning / [Reinforcement learning for reasoning](https://mlsanity.com/lr/reinforcement-learning-for-reasoning/)
- LLM agents and tool use / [Reinforcement learning for LLM agents](https://mlsanity.com/agents/reinforcement-learning-for-llm-agents/)
- LLM alignment / RLHF / preference optimisation / [Process reward models for reasoning](https://mlsanity.com/alignment/process-reward-models-for-reasoning/)

## Builds on (most cited first, 29 of 29)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [ToolLLM: Facilitating Large Language Models to Master 16000+ Real-world APIs](https://mlsanity.com/paper/2307.16789/) (2023)
- [Language Models (Mostly) Know What They Know](https://mlsanity.com/paper/2207.05221/) (2022)
- [A Survey on LLM-as-a-Judge](https://mlsanity.com/paper/2411.15594/) (2024)
- [Understanding R1-Zero-Like Training: A Critical Perspective](https://mlsanity.com/paper/2503.20783/) (2025)
- [Math-Shepherd: Verify and Reinforce LLMs Step-by-step without Human Annotations](https://mlsanity.com/paper/2312.08935/) (2023)
- [Discovering Latent Knowledge in Language Models Without Supervision](https://mlsanity.com/paper/2212.03827/) (2022)
- [The Lessons of Developing Process Reward Models in Mathematical Reasoning](https://mlsanity.com/paper/2501.07301/) (2025)
- [Lookback Lens: Detecting and Mitigating Contextual Hallucinations in Large Language Models Using Only Attention Maps](https://mlsanity.com/paper/2407.07071/) (2024)
- [VisualPRM: An Effective Process Reward Model for Multimodal Reasoning](https://mlsanity.com/paper/2503.10291/) (2025)
- [SWEET-RL: Training Multi-Turn LLM Agents on Collaborative Reasoning Tasks](https://mlsanity.com/paper/2503.15478/) (2025)
- [Let's reward step by step: Step-Level reward model as the Navigators for Reasoning](https://mlsanity.com/paper/2310.10080/) (2023)
- [Critique-GRPO: Advancing LLM Reasoning with Natural Language and Numerical Feedback](https://mlsanity.com/paper/2506.03106/) (2025)
- [Tree Search for LLM Agent Reinforcement Learning](https://mlsanity.com/paper/2509.21240/) (2025)
- [GTA: A Benchmark for General Tool Agents](https://mlsanity.com/paper/2407.08713/) (2024)
- [AgentPRM: Process Reward Models for LLM Agents via Step-Wise Promise and Progress](https://mlsanity.com/paper/2511.08325/) (2025)
- [SPA-RL: Reinforcing LLM Agents via Stepwise Progress Attribution](https://mlsanity.com/paper/2505.20732/) (2025)
- [Latent Space Chain-of-Embedding Enables Output-free LLM Self-Evaluation](https://mlsanity.com/paper/2410.13640/) (2024)
- [Entropy-Regularized Process Reward Model](https://mlsanity.com/paper/2412.11006/) (2024)
- [Beyond the Final Answer: Evaluating the Reasoning Trajectories of Tool-Augmented Agents](https://mlsanity.com/paper/2510.02837/) (2025)
- [TIPS: Turn-Level Information-Potential Reward Shaping for Search-Augmented LLMs](https://mlsanity.com/paper/2603.22293/) (2026)
- [Demystifying Group Relative Policy Optimization: Its Policy Gradient is a U-Statistic](https://mlsanity.com/paper/2603.01162/) (2026)
- [Dynamic and Generalizable Process Reward Modeling](https://mlsanity.com/paper/2507.17849/) (2025)
- [AT2PO: Agentic Turn-based Policy Optimization via Tree Search](https://mlsanity.com/paper/2601.04767/) (2026)

## Built on by (most cited first, 5 of 5)

- [It's the Problem, Not the Path: Budget and Difficulty Confounds in LLM Reasoning Trajectories](https://mlsanity.com/paper/2609.03436/) (2026)
- [Before the Rollout Ends: Early Terminal Reward Prediction for Long-horizon Coding Agents](https://mlsanity.com/paper/2609.31995/) (2026)
- [Search-G1: Grounded Search Agents via Representation-Based Intrinsic Rewards](https://mlsanity.com/paper/2608.07531/) (2026)
- [ProCredit: From Outcome Rewards to Progress Credit in Agentic Reinforcement Learning](https://mlsanity.com/paper/2609.27532/) (2026)
- [My FAULT: Self-Diagnosis as Credit Assignment in Self-Evolving Agentic Reinforcement Learning](https://mlsanity.com/paper/2610.01161/) (2026)

## Abstract (arXiv)

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.