PaperlayerView exact revision on arXiv

Source-linked AI summary

PRO-Step: Step-level Process Reward Optimization for Retrieval-Augmented Generation

MinKeon Kim, Namjun Lee, Jaekwang Kim

arXiv:2609.01658v1cs.CLcs.AI

TL;DR

Multi-hop RAG can propagate early retrieval errors, while outcome-only and existing process rewards may miss flawed intermediate steps that reach correct answers. PRO-STEP evaluates logical validity and evidential grounding with a generative PRM, uses PRM-guided value tree search to build preferences, and applies step-level DPO. Across five QA benchmarks, it achieves the best average EM and F1, while trading some performance on MuSiQue for reproducibility without closed-source teachers.

  • Problem

    Outcome-based RAG optimization misses intermediate retrieval and reasoning errors, while existing process methods can reward spurious successes without jointly checking logical validity and evidential grounding.

  • Method

    PRO-STEP trains a generative PRM for step rationales and correctness labels, uses PRM-guided value tree search to construct preferences, and optimizes the policy with step-level DPO.

  • Results

    PRO-STEP achieves the best average performance across five QA benchmarks, recording 34.5 EM and 44.1 F1.

  • Takeaways & Limitations

    Dense process supervision targets intermediate reasoning flaws and provides an effective approach for optimizing retrieval-augmented reasoning.

  • Takeaways & Limitations

    PRO-STEP trails StepSearch by 1.3 EM on MuSiQue, reflecting a reproducibility trade-off from avoiding closed-source teachers and using less training data.

Abstract

from arXiv · show

Retrieval-Augmented Generation enhances Large Language Models by grounding responses in external knowledge, but multi-hop reasoning remains vulnerable to error propagation, where early retrieval failures confound subsequent steps. Standard outcome-based optimization only rewards the final answer, leaving intermediate retrieval and reasoning errors undetected. While existing process-based methods introduce step-level signals, they still score each step against the final answer, rewarding spurious successes where flawed retrieval coincidentally produces the correct answer. Step-level supervision in RAG requires evaluating both logical validity and evidential grounding at each step. We introduce PRO-STEP: we train a generative PRM that evaluates both dimensions, employ PRM-guided value tree search to construct preference pairs contrasting valid steps against flawed ones, and optimize the policy via step-level Direct Preference Optimization. Experiments on single and multi-hop QA datasets demonstrate that PRO-STEP achieves the best average EM and F1 across five benchmarks. Code, models, and training data are publicly available at https://github.com/keemminnke/PRO-Step.

1 Introduction

RAG improves factuality by grounding language models in external knowledge, but multi-hop tasks remain vulnerable to intermediate retrieval and reasoning errors. PRO-STEP addresses this gap with step-level supervision that evaluates both logical validity and evidential grounding, achieving the best average EM and F1 across five benchmarks.

  • Traditional RAG struggles with multi-step tasks requiring iterative retrieval and dynamic query reformulation.
  • Outcome-based optimization provides feedback only from final-answer correctness, leaving early retrieval failures and intermediate reasoning errors undetected.These errors can propagate through subsequent reasoning steps, while the sparse signal requires more training data to converge.
  • Existing process-based RAG methods can reward spurious successes because they score steps against final outcomes without evaluating logical validity and evidential grounding together.
  • PRO-STEP trains a generative PRM to produce correctness labels and rationales, then uses PRM-guided VTS and process-supervised DPO for policy optimization.
  • PRO-STEP achieves the best average EM and F1 among evaluated methods across five single-hop and multi-hop QA benchmarks.

2 Related Work

RAG research progressed from fixed retrieval toward iterative and autonomous retrieval, but outcome-only rewards remain sparse and process methods can still accept flawed intermediate steps. PRO-STEP extends process supervision by evaluating both reasoning validity and evidence grounding in RAG trajectories.

  • Early RAG systems appended fixed passages, whereas later methods introduced uncertainty-triggered retrieval, query reformulation, and interleaved reasoning-retrieval chains.
  • RL-based systems enabled autonomous retrieval decisions but generally used final-answer correctness as a sparse reward without intermediate quality signals.
  • Recent process-level RAG methods provide step-level rewards and tree search, yet still score intermediate steps against final-answer correctness.
  • PRM research introduced step verification, automated rollout scaling, outcome estimation, and generative rationales for formal reasoning tasks.
  • PRO-STEP adapts generative PRMs to RAG by evaluating logical validity and evidential grounding, then guiding preference construction and policy optimization.

3 PRO-STEP

PRO-STEP is a two-stage framework that trains a generative process reward model and uses its step-level evaluations within value tree search and DPO. The resulting preference data contrasts trajectories using both intermediate correctness and downstream outcome quality.

  • 3 PRO-STEP: PRO-STEP constructs process-supervised data by generating interleaved reasoning-and-retrieval trajectories, annotating steps, and training a generative PRM.
  • 3 PRO-STEP: The PRM generates a rationale and binary correctness label for each step, while retaining later steps so recovery after a local failure remains possible.
  • 3 PRO-STEP: PRM-guided value tree search replaces random rollouts with direct value estimation over states formed from the question and preceding context.
  • 3 PRO-STEP: Terminal nodes receive depth-discounted F1 rewards, while non-terminal nodes use PRM predictions propagated to ancestor values.
  • 3 PRO-STEP: Sibling steps become chosen and rejected preferences when their combined value difference exceeds δ, balancing descendant outcome quality with step correctness.
  • 3 PRO-STEP: Step-level DPO optimizes the policy on these preference tuples while masking retrieved-document blocks from the loss computation.

4 Experiments

Across five QA benchmarks, PRO-STEP achieves the strongest average performance and improves multi-hop reasoning through dense, step-level supervision. Ablations and comparisons attribute gains to PRM-based filtering of flawed trajectories, preference optimization, and adaptive retrieval.

  • Overall results: PRO-STEP records the highest overall average performance, achieving 34.5 EM and 44.1 F1 across the evaluated benchmarks.
  • Process vs. outcome: PRO-STEP outperforms outcome-based Search-R1 by 2.5 EM and 3.4 F1 on average, with gains of 9.2 EM and 8.9 F1 on 2WikiMultiHopQA.It reaches these results with 5k seed questions, compared with roughly 170k training instances reported for Search-R1.
  • Component ablation: Removing PRM scoring reduces average performance from 34.5 to 32.7 EM and from 44.1 to 41.9 F1, with larger drops on multi-hop datasets.The ablation supports PRM supervision as particularly useful for filtering spurious reasoning in multi-hop tasks.
  • Optimization strategies: DPO achieves the best average performance with 34.5 EM and 44.1 F1, while SFT and KTO perform worse under the tested optimization strategies.DPO provides explicit step-wise contrast between valid and flawed reasoning.
  • PRM comparison: Under WMV with K = 128, the RAG-specific PRM reaches 46.5 F1 on 2WikiMultiHopQA, 59.5 F1 on HotpotQA, and 50.8 F1 on PopQA.These scores outperform both MV and VersaPRM across all three datasets, whereas general-purpose PRMs degrade as K increases under BoN.
  • Retrieval behavior: Increasing retrieved documents from k = 1 to k = 3 improves performance by an average of 6.4 F1, with further gains at k = 5.On retrieval depth, PRO-STEP resolves 65.6% of PopQA instances at step one, versus 52.0% on HotpotQA and 36.6% on 2WikiMultiHopQA.
  • Generalization: PRO-STEP generalizes across backbones and scales, improving Llama-3.1-8B-Instruct by 11.1 EM and 9.5 F1 and outperforming Search-R1 at the 3B scale by 4.82 EM and 6.49 F1.The framework also retains most performance without instruction tuning, reaching 32.85 / 42.01 average EM / F1 versus 34.51 / 44.09 for the instruction-tuned backbone.
  • Recovery analysis: PRO-STEP outperforms Search-R1 after both failed and successful initial retrieval, including a 12.1 EM gain on 2WikiMulti and a 29.2 EM gain on Bamboogle.

5 Conclusion

PRO-STEP integrates process-level supervision into RAG by evaluating the logical validity and evidential grounding of intermediate steps. PRM-guided value tree search constructs targeted preference pairs, which DPO uses to improve interleaved reasoning and retrieval.

  • PRO-STEP trains a generative PRM to evaluate both logical validity and evidential grounding for intermediate retrieval-and-reasoning steps.
  • PRM-guided value tree search filters spurious trajectories and constructs step-level preference pairs for policy optimization.
  • Direct Preference Optimization trains the policy to navigate interleaved generation and retrieval using targeted process supervision.

Limitations

The paper identifies limitations involving labeler reliability, uneven dataset performance, small-test-set uncertainty, and a failed trajectory-augmentation strategy.

  • QwQ-32B labeler reliability was audited against Claude Opus 4.7, showing substantial agreement (κ = 0.6104) without systematic permissive bias.The paper uses QwQ-32B for reproducibility without restrictive API licenses.
  • PRO-STEP trails StepSearch by 1.3 EM on MuSiQue, reflecting a trade-off between reproducibility without API access and headroom on the hardest tasks.The training uses 5,000 seed questions and 15,877 pairs, compared with StepSearch’s GPT-4o-distilled 19k corpus.
  • Bamboogle’s 125-instance test set yields wide confidence intervals, and its comparison with ReasonRAG is statistically indistinguishable from noise (p = 0.66).
  • PRM-feedback-guided regeneration degrades average performance by 10.7 EM and 9.8 F1 because regenerated trajectories drift from natural inference behavior.The outcome-margin filter does not detect this stylistic drift.

Ethical Considerations

The paper states that it uses publicly available datasets and training data derived from them, while noting residual LLM-system risks and the need for human oversight.

  • The study uses publicly available HotpotQA, MuSiQue, 2WikiMultiHopQA, PopQA, and Bamboogle datasets for research purposes.
  • Training data is derived solely from existing datasets, and no personally identifiable information is collected or processed.
  • The framework does not introduce new risks beyond residual hallucination and retrieval of outdated information inherent to LLM-based systems.
  • Practitioners deploying the framework in domains should apply appropriate human oversight.

B Data Statistics

The appendix reports the scale and composition of supervision data, the filtering used to construct DPO pairs, and evidence that the filter improves data efficiency without materially changing performance.

  • 109,664 PRM training decisions come from 31,728 trajectories sampled across 2,000 source multi-hop QA questions.Search actions comprise 71.2% of the steps, while internal reasoning accounts for 0.5%.
  • The outcome-margin filter retains 2,866 questions and 15,877 contrastive pairs from 5,000 seed questions.The retained pairs satisfy chosen F1 ≥0.2 and ∆F1 ≥0.2.
  • Removing the filter increases pair count by 2.5× but produces comparable average performance, 34.65 versus 34.51 EM.The filter is adopted for training and data efficiency rather than performance.
  • 75.0% of preference comparisons use the same action type, while step depths peak at 2 and 3 with 45.7% and 29.5% of pairs.The reward margin is at least 0.20 for 98.4% of pairs, including 52.0% with margins of at least 0.50.

C VTS Implementation Details

The appendix details PRM-guided value tree search, including its selection and reward mechanics, parameter choices, ablations, and audits of the supervision labels.

  • VTS balances exploration and exploitation using UCB selection based on mean backed-up value, visit count, and an exploration constant.Unvisited nodes are selected immediately.
  • Terminal F1 is discounted by 0.9 per depth level, while intermediate scores combine backed-up F1 reward with PRM feedback weighted by α = 0.3.The PRM term breaks ties without overriding substantial F1 differences.
  • The main branching configuration improves average performance by 1.9 EM and 2.0 F1 over reduced branching, which remains competitive.
  • Every PRM-signal configuration except α = 0.5 outperforms the no-PRM baseline of 32.68 EM, with α = 0.3 achieving 34.51 EM and 44.09 F1.The results support complementary process and outcome rewards.
  • Human auditing found 84% agreement with QwQ-32B, while frontier-model auditing reported κ = 0.61 and strongest agreement on answer correctness and search quality.Agreement was lower for intermediate reasoning steps.

F Statistical Significance of Main Results

The statistical analyses support PRO-STEP’s stronger average performance against major baselines, including at the 3B scale, while documenting evaluation settings and ablations.

  • Evaluation protocol: Temperature 0 greedy decoding removes sampling stochasticity, so significance is reported over per-question score distributions rather than multiple seeds.The evaluation supplements the main results with bootstrap confidence intervals and paired t-tests.
  • Main significance findings: PRO-STEP significantly outperforms Search-R1 and ReasonRAG in macro-AVG EM and F1 at the reported 95% confidence level.Against StepSearch, F1 improvements are significant; EM confidence intervals marginally include zero, although PRO-STEP leads significantly on PopQA and 2WikiMultiHopQA.
  • 3B-scale results: 31.04 / 40.00 average EM / F1 is achieved by PRO-STEP at 3B, compared with 26.22 / 33.51 for Search-R1 3B.PRO-STEP uses 5k seed questions, whereas Search-R1 uses 170k training questions.
  • Training configuration: Direct DPO on the 3B base model achieves only 12.40 EM, motivating SFT warmup before preference optimization.The successful 3B configuration also uses trajectories generated by the larger policy model.
  • Ablation evidence: 1.9 EM and 2.0 F1 are the average gains from the main VTS branching configuration over reduced branching.The comparison is reported in the VTS branching ablation.

H Intermediate-Flawed-But-Correct Analysis

The IFBC analysis measures correct answers supported by flawed intermediate steps and shows that PRM-guided preference selection reduces this failure mode while retaining more correct trajectories.

  • Definition: IFBC denotes a trajectory with at least one flawed intermediate step that still produces a correct final answer.Intermediate steps exclude the final-answer step.
  • Interpretation: Lower IFBC indicates that correct answers are less frequently supported by flawed intermediate trajectories.The analysis uses annotated trajectories from 500 paired questions each in HotpotQA, 2WikiMultiHopQA, and MuSiQue.
  • Results: 40.3% to 26.8% is the reduction in IFBC under PRM-guided selection, while correct trajectories increase from 447 to 474.The controlled comparison uses identical training trees and DPO hyperparameters, differing only in preference-pair selection.
  • Training signal: 19.8% to 10.7% is the reduction in flawed-but-successful trajectories promoted to the chosen side of DPO pairs.The comparison contrasts outcome-only selection with PRM-guided selection.
  • Baseline comparison: PRO-STEP has significantly lower IFBC rates than Search-R1 and StepSearch, while ReasonRAG’s numerically lower rate is not statistically significant.ReasonRAG also produces fewer correct trajectories, 430 versus 474, so IFBC is interpreted jointly with task performance.

I Detailed Recovery Analysis

The recovery analysis examines retrieval failures, recovered trajectories, and an attempted regeneration augmentation, showing benefits for PRO-STEP but a severe failure when training and inference formats diverge.

  • Recovery definitions: A retrieval succeeds when the gold passage appears among the top-3 documents, and recovery depth k is the first successful retrieval step.The initial retrieval is indexed as step 0, and reported subset sizes follow each model’s own search behavior.
  • Recovered trajectories: PRO-STEP leads on four of five datasets in token-level F1 over recovered trajectories, with the largest margin on 2WikiMulti at +7.8.Bamboogle is the exception, with only 25 questions in its recovered subset.
  • Recovery depth: Gains are consistent for recovery depths k = 1 through k = 3 on 2WikiMulti.The k ≥4 comparison favors PRO-STEP by a wide margin, but both subsets are small and unreliable.
  • No-recovery trajectories: PRO-STEP performs competitively or better on four of five datasets when initial retrieval already succeeds, indicating its recovered-trajectory advantage is not limited to an easier subset.This comparison uses the complementary no-recovery case.
  • Failed augmentation: 34.51 to 23.83 average EM and 44.09 to 34.27 average F1 occur after regeneration augmentation, despite stronger DPO training metrics.The augmented run’s loss roughly halves, preference margin grows by 60%, and training accuracy rises from 81% to 91%.
  • Failure diagnosis: The regeneration variant fails because feedback-conditioned trajectories differ from the deployed policy’s inference-time distribution.The outcome-margin filter checks answer-level quality but not stylistic drift in intermediate tokens.
Loading 2609.01658v1…