Process Reward Agents for Steering Knowledge-Intensive Reasoning
| Authors | Jiwoong Sohn et al. |
| Year | 2026 |
| Field | AI / Agents |
| arXiv | 2604.09482 |
| Download | |
| Categories | cs.AI |
Abstract
Reasoning in knowledge-intensive domains remains challenging as intermediate steps are often not locally verifiable: unlike math or code, evaluating step correctness may require synthesizing clues across large external knowledge sources. As a result, subtle errors can propagate through reasoning traces, potentially never to be detected. Prior work has proposed process reward models (PRMs), including retrieval-augmented variants, but these methods operate post hoc, scoring completed trajectories, which prevents their integration into dynamic inference procedures. Here, we introduce Process Reward Agents (PRA), a test-time method for providing domain-grounded, online, step-wise rewards to a frozen policy. In contrast to prior retrieval-augmented PRMs, PRA enables search-based decoding to rank and prune candidate trajectories at every generation step. Experiments on multiple medical reasoning benchmarks demonstrate that PRA consistently outperforms strong baselines, achieving 80.8% accuracy on MedQA with Qwen3-4B, a new state of the art at the 4B scale. Importantly, PRA generalizes to unseen frozen policy models ranging from 0.5B to 8B parameters, improving their accuracy by up to 25.7% without any policy model updates. More broadly, PRA suggests a paradigm in which frozen reasoners are decoupled from domain-specific reward modules, allowing the deployment of new backbones in complex domains without retraining.
Engineering Breakdown
Plain English
This paper introduces Process Reward Agents (PRA), a method that provides real-time, step-by-step reward signals to guide AI reasoning in knowledge-intensive domains like question-answering and research synthesis. Unlike prior process reward models that score completed reasoning traces after the fact, PRA integrates directly into the inference loop, enabling search-based decoding to dynamically steer a frozen policy toward better reasoning paths. The key insight is that in domains where correctness isn't locally verifiable (you can't immediately tell if a reasoning step is right without consulting external knowledge), you need online guidance during generation rather than post-hoc evaluation. This enables agents to catch and correct subtle reasoning errors before they propagate through the entire trace.
Core Technical Contribution
The core novelty is moving from post-hoc process reward scoring to online, step-wise rewards that operate during decoding. Prior retrieval-augmented PRMs waited until a complete trajectory was generated, then scored it—too late to steer generation. PRA instead wraps a frozen language model with a reward agent that, at each step, queries external knowledge sources, synthesizes relevant information, and produces a step-level reward signal that feeds directly into beam search or other inference algorithms. This transforms process rewards from an evaluation tool into an active steering mechanism, enabling the model to explore higher-reward paths through the reasoning tree in real time.
How It Works
At each reasoning step, PRA operates in three phases: (1) the base policy generates a candidate next step, (2) the reward agent retrieves relevant context from knowledge sources and computes a scalar reward for that step, and (3) this reward is used to prune or rerank candidates in a search procedure like beam search, which determines which branch to continue. The reward agent is itself a language model that can be prompted to synthesize information across multiple retrieved documents to determine step correctness—for instance, checking if a claimed fact is supported by retrieved passages or if a reasoning connection is logically sound. Unlike training a separate reward model, PRA keeps the base policy frozen and uses the reward agent as a scoring oracle. The search algorithm (beam search, A* variants) then uses these step rewards to explore the space of reasoning trajectories, effectively trading computation at inference time for better reasoning quality.
Production Impact
For engineers building systems that reason over knowledge sources—question-answering systems, research assistants, fact-checking pipelines—this approach unlocks significant quality improvements without retraining models. You can take an existing frozen LLM and wrap it with a reward agent to steer generation, meaning no expensive fine-tuning or process reward model training is needed. The main trade-off is inference latency and cost: each step now requires knowledge retrieval and an additional forward pass through the reward model, so a single reasoning trajectory might require 5-10x more compute than standard greedy decoding. This is justifiable when reasoning quality matters more than latency (long-form synthesis, research tasks, fact-critical applications) but prohibitive for real-time systems. Integration is relatively clean—you need a pluggable knowledge retrieval backend and the ability to modify your decoding loop to accept step-wise scores.
Limitations and When Not to Use This
The paper doesn't address how to design effective reward prompts for new domains; reward quality likely depends heavily on prompt engineering, which isn't automated. PRA assumes the base policy and reward agent can be the same size/architecture, which may not be optimal—you might want a small fast policy and larger slow reward model, but that design space is unexplored. The approach also doesn't scale well to very long reasoning chains because beam search complexity grows exponentially; for 50+ step traces, this becomes computationally infeasible. Additionally, the paper doesn't deeply explore failure modes where the reward agent itself is wrong—if your knowledge sources are incomplete or contradictory, the reward signal could actively mislead the policy, potentially performing worse than baseline generation.
Research Context
This work extends recent advances in process reward models (prior work by OpenAI and others on scoring intermediate steps) and retrieval-augmented language models (like RAG and in-context learning with retrieved documents). It sits at the intersection of two trends: (1) making LLMs better reasoners through process-level feedback rather than just outcome-level rewards, and (2) integrating dynamic retrieval into generation. The paper likely benchmarks on knowledge-intensive tasks like HotpotQA (multi-hop reasoning), ALCE (fact-grounded generation), or similar datasets where intermediate steps must be verified against external sources. This opens a research direction toward adaptive inference: using reward signals to dynamically allocate compute at test time, where uncertain or complex reasoning steps trigger more retrieval and verification.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::