PaperlayerView exact revision on arXiv

Source-linked AI summary

Process Reward Agents for Steering Knowledge-Intensive Reasoning

Jiwoong Sohn, Tomasz Sternal, Kenneth Styppa, Torsten Hoefler, Michael Moor

arXiv:2604.09482v2cs.AI

TL;DR

Knowledge-intensive medical reasoning lacks reliable intermediate-step verification, while prior retrieval-augmented rewards score completed trajectories post hoc. PRA gives frozen policies online, domain-grounded step rewards for search-based decoding, improving benchmark performance and transferring across unseen backbones without policy updates.

  • Problem

    Intermediate medical reasoning steps often require synthesizing external evidence and are not locally verifiable, while prior retrieval-augmented process rewards operate post hoc on completed trajectories.

  • Method

    PRA decouples retrieval and step-wise evaluation from a frozen policy, selectively acquiring evidence and steering search by ranking and pruning partial trajectories during generation.

  • Results

    PRA consistently outperforms strong baselines across medical reasoning benchmarks, reaching 81.9% accuracy on MedQA with Qwen3-4B-Instruct and improving unseen frozen policies by up to 25.7%.

  • Takeaways & Limitations

    Online, step-wise rewards provide a stable, transferable mechanism for improving medical reasoning while allowing frozen reasoners and domain-specific reward modules to remain decoupled.

  • Takeaways & Limitations

    Process reward models face a critical challenge in generalizing across policies because off-policy scoring can degrade under distributional mismatch.

Abstract

from arXiv · show

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), an inference-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 81.9% 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.

1. Introduction

Medical reasoning requires intermediate-step verification across evolving, layered evidence, but existing process rewards evaluate completed traces post hoc. PRA instead provides online, retrieval-grounded step rewards to steer frozen policies during generation, with strong benchmark results and transfer across models.

  • Intermediate medical reasoning steps are difficult to verify mechanically, yet clinically correct decisions must remain defensible throughout the reasoning trace.
  • Synthesizing evolving scientific literature, guidelines, and institutional protocols makes repeated backbone fine-tuning costly, while context-level retrieval does not ensure timely evidence use or error correction.
  • Prior retrieval-augmented process reward approaches score complete trajectories post hoc, limiting intervention before errors accumulate and restricting alternative-path exploration.
  • PRA observes partial traces, optionally searches external medical evidence, and assigns local rewards that evaluate intermediate steps before errors propagate.
  • PRA formulates retrieval-grounded step-wise evaluation as online control, decoupling evidence search from a frozen policy and enabling inference-time branching and pruning.
  • 81.9% accuracy on MedQA with Qwen3-4B-Instruct establishes state-of-the-art performance at the 4B scale under a matched policy sampling budget.
  • Up to 25.7% MedQA accuracy improvement transfers to unseen frozen policies spanning 0.5B to 8B parameters without policy updates.

2. Related Work

Medical reasoning systems commonly rely on post-training or retrieved context, while online retrieval coupled with step-wise verification remains underexplored. PRA addresses this gap amid broader process-reward challenges involving noisy supervision and cross-policy generalization.

  • Medical reasoning must combine expanding biomedical literature with contextual constraints such as guidelines and common practice.
  • Existing approaches improve medical reasoning through post-training or by injecting retrieved documents directly into the policy context.
  • Process reward models assign rewards to intermediate reasoning steps, making them suitable for tree-search inference frameworks.
  • Automated intermediate-step labels can be noisy because correct final answers may coexist with incorrect reasoning steps.
  • Off-policy PRM scoring often degrades under distributional mismatch, creating a critical cross-policy generalization challenge.
  • Medical retrieval-augmented process rewards typically operate after complete traces, leaving online retrieval-grounded evaluation underexplored under policy shift.

3. Process Reward Agents

PRA is an inference procedure around a frozen reasoning policy and fixed domain knowledge base: it selectively retrieves evidence, scores partial traces, and uses step rewards to control tree search.

  • 3.1. Problem Formulation: The formulation assumes verifiable answers through a correctness function that returns one for a matching answer and zero otherwise.
  • 3.1. Problem Formulation: A frozen policy autoregressively generates cumulative reasoning traces, with the final step defined as the completed trace’s answer.
  • 3.1. Problem Formulation: The inference procedure takes a question, frozen policy, and fixed domain knowledge base as inputs and outputs a final answer.
  • 3.2. Process Reward Agents: PRA separates reasoning from evidence acquisition through action-controller and reward-readout components implemented as separate token-level readouts.
  • 3.2. Process Reward Agents: When the controller selects search, relevant documents are retrieved; when it selects reward, the evidence set is empty, and the scorer evaluates the latest step.
  • 3.2. Process Reward Agents: Step-wise rewards steer inference-time tree search by ranking and pruning candidate trajectories online.
  • 3.2. Process Reward Agents: PRA supports knowledge-base updates, backbone substitution without retraining, and selective retrieval as an inference-time compute–accuracy trade-off.
  • 3.3. PRA-Guided Tree Search: Beam search maintains B traces, expands each into b candidates, scores all candidates, and retains the top-B by cumulative reward.

4. Experiments

PRA is evaluated across in-distribution and unseen medical benchmarks, policy backbones, retrieval settings, and inference-time scaling. Across these evaluations, online step-wise rewards consistently improve reasoning performance and enable generalization without policy retraining.

  • Experimental setup: PRA is evaluated on held-out MedQA and six out-of-distribution medical benchmarks using a fixed retrieval configuration and matched sampling budgets.The out-of-distribution datasets include Medbullets, MedMCQA, MMLU-Med, GPQA, and clinical cases from The Lancet and The New England Journal of Medicine.
  • Main results: PRA consistently outperforms baselines across seven medical benchmarks, improving over RAG + SC by 5.2 points on average.The comparison includes direct answering, chain-of-thought, retrieval-augmented generation, and self-consistency variants.
  • Inference-time scaling: PRA continues benefiting from additional inference compute, whereas self-consistency shows little improvement beyond eight samples.The paper attributes this difference to PRA applying rewards during generation rather than aggregating only completed samples.
  • Generalization: Across six unseen benchmarks, PRA outperforms the strongest baseline by an average of 5.2 points.These results measure generalization to medical reasoning benchmarks not seen during training.
  • Generalization: PRA transfers from Qwen3-4B training to frozen policy models spanning 0.5B to 8B parameters without policy retraining.On Qwen2.5-0.5B-Instruct, MedQA accuracy increases from 28.4 to 54.1, a 90.5% relative improvement over chain-of-thought.

5. Analysis

The analysis shows that PRA’s gains depend on combining trained reward modeling with online, process-level control, while selective retrieval exposes an accuracy–search-cost trade-off. Margin-shift patterns further indicate that retrieved evidence affects correct and incorrect trajectories differently across reasoning steps and question difficulty.

  • Ablation on Inference: 81.9 accuracy is achieved by PRA, which combines reward-agent training and retrieval for effective inference-time scaling with beam search.Under single-sample decoding, CoT and RAG are comparable, while search improves scaling and trained PRA adds a clear gain.
  • Ablation on Inference: Online process-level rewards outperform outcome-level or post hoc process-level aggregation when guiding reasoning during generation.The comparison uses the same trained reward agent and varies only reward granularity and timing.
  • Search–Accuracy Trade-off: Selective retrieval lets PRA trade retrieval cost against answer accuracy by triggering search when its step-wise search score exceeds threshold θdep.The threshold sweep identifies Pareto-frontier operating points for different search budgets.
  • Margin Shift Analysis: Correct traces show larger margin shifts at later reasoning steps, whereas incorrect traces show decreasing shifts later in the trajectory.At the final step, retrieved evidence has little effect, consistent with minimal substantive reasoning there.
  • Margin Shift Analysis: Correct traces exhibit especially large margin shifts on harder questions, with the shift decreasing as policy solve rate increases.Difficulty is defined by the fraction of policy-generated reasoning samples reaching the correct answer.

6. Conclusion

PRA guides frozen reasoning models through knowledge-intensive tasks using online, domain-grounded process rewards and consistently outperforms reasoning and retrieval baselines. Its gains extend across unseen policy backbones and arise primarily from applying process rewards online, while selective search can reduce retrieval under an accuracy–cost frontier.

  • Conclusion: PRA guides frozen reasoning models with online, step-wise, domain-grounded process rewards without modifying the policy.The framework treats inference-time reasoning as controllable search over partial traces.
  • Conclusion: PRA consistently outperforms strong reasoning and retrieval baselines across medical benchmarks, including state-of-the-art 4B-scale MedQA performance and robust out-of-distribution gains.The conclusion also reports generalization across unseen policy backbones and underutilized reasoning capacity in smaller models.
  • Conclusion: Ablations indicate that gains arise primarily from applying process-level rewards online during generation rather than from post hoc scoring alone.The conclusion also characterizes selective search as a way to reduce retrieval while preserving performance along a Pareto frontier.

Impact Statement

The paper targets reliable and verifiable reasoning traces in high-stakes healthcare applications, where individual reasoning steps must meet a high standard for trust and appropriate reliance on AI systems.

  • Impact Statement: PRA is intended to improve the reliability and verifiability of medical reasoning traces for high-stakes healthcare use.The impact statement emphasizes that individual steps must meet a high bar to support trust and appropriate reliance on AI systems.

A. Table of Notations

This appendix material identifies the manuscript’s notation summary, without supplying substantive methodological or empirical findings.

  • Table of Notations: Table 5 provides a summary of the notation used throughout the main manuscript.The supplied passages identify the notation table but do not enumerate its symbols.

B. Stage-Level Batching

PRA-guided beam search manages each question as a trace with staged partial reasoning, while a global queue batches traces by their pending operation.

  • Each Trace maintains a beam of partial reasoning traces and a stage tag among REASON, REWARD, SEARCH, and DONE.
  • The global queue drains active traces, partitions them by stage, and dispatches each partition as one batched operation.Operations are sent to the policy, reward agent, or retriever regardless of each trace’s step index.
  • Batched dispatch targets the policy π, retriever ρ, or reward agent µϕ according to the pending stage.

C. Additional Training Details

The authors fine-tune Qwen3-4B-Instruct for PRA label prediction and use distinct search-label regimes for the main experiments and trade-off analysis.

  • Qwen3-4B-Instruct is fine-tuned to predict the reasoning and search labels described in Section 4.1.
  • Training uses a 3 × 10−5 learning rate, cosine decay, 100 warmup steps, 0.01 weight decay, batch size 16, 3 epochs, and bfloat16 precision.
  • The main experiments use an always-search setting with the search label fixed to 1 at every reasoning step.
  • The Search–Accuracy Trade-off analysis uses search labels derived from the margin-shift criterion described in Section 4.1.

D. Prompt Templates

The experiments use structured prompts for frozen-reasoner generation, step-level evaluation, and PRA beam-search inference, with standardized formats for reasoning and answer parsing.

  • Prompt set: Figures 7, 8, and 9 contain the prompts used throughout the experiments, including the PRA prompt for beam-search inference and supervised fine-tuning.
  • Policy prompt: The frozen-reasoner policy prompt requires step-by-step reasoning, one option per step, numbered Step {number}: formatting, and a standardized final-answer phrase.The user prompt presents the question and lettered options before an Answer: field, while retrieved documents are excluded from the policy prompt.
  • Teacher prompt: The teacher prompt supplies retrieved documents, the question, the correct answer, and the trace, then scores only the last reasoning step as 0 or 1.A second pass omits documents to estimate the no-retrieval margin used to construct search labels.
  • PRA prompt: The PRA prompt includes the documents block only on the post-retrieval
Loading 2604.09482v2…