Research map / Reinforcement learning / RL from feedback
PAIR: Prefix-Aware Internal Reward Model for Multi-Turn Agent Optimization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes PAIR, a two-stage internal reward model that corrects a frozen hidden-state correctness probe with an attention-based head, giving dense step-level rewards for GRPO on multi-turn agent tasks without external calls or ground truth.
Problem
Sparse outcome rewards in GRPO hinder step-level credit assignment for multi-turn agents; existing dense alternatives need full rollouts, external LLM judges, ground-truth log-probs, or separate reward models. Hidden-state correctness probes also degrade under prefix contamination.
Approach
Hidden-state probes track belief-consistency with a (possibly corrupted) prefix, while attention features track structural, grounded signals. PAIR freezes a hidden-state probe to produce s_bc and trains a logistic correction head on attention features plus s_bc to output a grounded correctness score used as a momentum-shaped step reward.
Key results
- PAIR reaches AUROC 0.9217 on GTA and 0.8129 on ToolBench for step-level correctness, the highest among 10 probing and unsupervised baselines.
- In GRPO, PAIR (momentum reward) reaches success rate 0.2489 on GTA and 0.4498 on ToolBench, above LLM-as-a-judge (0.2387 and 0.4203) and other baselines.
- Hidden-state probes score at or below chance on an adversarial diagnostic set where belief-consistency and grounded correctness are opposed.
Surprising findings
- Hidden-state probes invert the reward at repair turns, actively penalizing error correction under contaminated prefixes.
- Naive concatenation of hidden and attention features yields near-chance performance on the diagnostic set.
Limitations
- Contaminated training data is synthetic, generated by GPT-4o-mini, and relies on GPT-4o-mini for contamination and judging.
- Evaluation uses small benchmarks (GTA has 229 samples), and PAIR does not achieve lowest ECE in all cases.
- Probe-based rewards require offline labelled correctness data, and performance degrades with contamination distance.
Experimental setup
Datasets: GTA ToolBench HotpotQA
Benchmarks: GTA ToolBench HotpotQA
Metrics: AUROC ECE success rate
Baselines: Outcome-only GRPO LLM-as-a-judge (GPT-4o-mini) Last-token probe Mean-pooled probe Multi-layer probe Attention probe Hidden+Attn probe Multi-Attn probe Lookback Lens Head Entropy CoE-C IGPO
Builds on: GRPO Lookback Lens Chain-of-Embedding IGPO Tree-GRPO
Closest papers by meaning
- Group-in-Group Policy Optimization for LLM Agent Training (NeurIPS 2025 · similarity 0.89)
- MURPHY: Feedback-Aware GRPO with Retrospective Credit Assignment for Multi-Turn Code Generation (NeurIPS 2026 · similarity 0.88)
- Reflective Prompted Policy Optimization: Trajectory-Grounded Revision and Salience Bias (NeurIPS 2026 · similarity 0.88)
- Step-by-Step Optimization-like Reasoning in LLMs over Expanding Search Spaces (NeurIPS 2026 · similarity 0.88)
- InfoPO: Information-Driven Policy Optimization for User-Centric Agents (ICML 2026 · similarity 0.88)
- Agentic Reinforced Policy Optimization (ICLR 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.