Technical Summary: SCPRM for Knowledge Graph Question Answering

Problem Statement

Large Language Models (LLMs) have demonstrated strong capabilities in complex reasoning via Chain-of-Thought (CoT) prompting. However, the linear structure of CoT makes LLMs vulnerable to error propagation, where a single intermediate mistake can derail the entire reasoning trajectory. While Outcome Reward Models (ORMs) verify final answers, they cannot distinguish between correct answers derived from logical reasoning versus those resulting from hallucination or compensatory errors.

Process Reward Models (PRMs) address this by evaluating intermediate steps. However, existing PRMs often suffer from a risk compensation effect: in tasks like mathematical reasoning, a risky or incorrect intermediate step can be offset by subsequent correct steps, leading the model to assign high rewards to flawed reasoning paths. This issue is critical in Knowledge Graph Question Answering (KGQA), particularly in risk-sensitive domains like medicine and law. In KGs, multiple paths may exist between entities, but only specific paths are logically valid. A risky step (e.g., traversing an irrelevant entity type) renders a path flawed, even if it coincidentally reaches the correct final answer. Furthermore, existing conditional reward models often introduce a length bias, where longer reasoning paths are disproportionately penalized due to the multiplicative nature of probability aggregation, hindering the exploration of necessary multi-hop paths.

Methodology: Schema-aware Cumulative Process Reward Model (SCPRM)

The authors propose SCPRM, a reward model designed to evaluate reasoning paths by conditioning on the reasoning prefix and incorporating schema distance. The approach is inspired by the A* search algorithm, decoupling evaluation into cumulative past rewards and future rewards.

1. Cumulative Past Reward Modeling

To prevent risk compensation, SCPRM models the safety of a reasoning trajectory step-by-step.

• Step-wise Safety Probability: For a state $s_k$ in a trajectory $\pi$, the model predicts a safety probability $p_k$ using a risk head ($\phi$) conditioned on the prompt, query, and the reasoning prefix up to $k$.
• Logarithmic Accumulation: Instead of multiplicative aggregation (which leads to vanishing rewards), the cumulative past reward $G(s_k)$ is calculated as the logarithmic sum of conditional safety probabilities:
  $$G(s_k) = \sum_{t=1}^{k} \log(1 - p_t)$$
  This formulation ensures that a single high-risk step (low $1-p_t$) imposes a persistent, multiplicative penalty on the total reward, effectively vetoing risky trajectories that vanilla PRMs might rescue.

2. Schema-aware Future Reward Modeling

To guide exploration without length bias, SCPRM estimates future success based on the alignment between the current reasoning path and the implicit constraints of the query.

• Schema Extraction: The query is parsed to extract a query schema (e.g., disease → gene → miRNA), representing the logical constraints. The current reasoning path is represented by a reasoning schema based on entity types and relations.
• Future Reward Estimation: A schema head ($\psi$) computes the similarity between the query embedding and the current path embedding. The future reward $H(s_k)$ is defined as the logarithmic probability of future success, optimized to minimize the distance between the predicted schema and the "ground truth" schema derived from the query constraints. This ensures logical consistency in multi-hop paths without penalizing path length.

3. Training and Optimization

The model is trained using a combination of pairwise ranking loss and future reward loss:

• Pairwise Loss: Positive paths (correct, risk-free) are ranked higher than negative paths (containing errors). The loss maximizes the reward difference between positive and negative pairs, including a specific penalty for the first erroneous step to enable early pruning.
• Future Loss: The predicted future success probability is optimized against a ground-truth future reward derived from the schema distance.
• Total Loss: $L = L_{pair} + \lambda \cdot L_{future}$.

4. Integration with MCTS

SCPRM is integrated into Monte Carlo Tree Search (MCTS) as the value function. The search process involves selection, expansion, evaluation, and backpropagation. SCPRM evaluates each potential action (triplet) based on the current state, guiding the search toward high-reward trajectories. The final answer is selected from the top-$k$ paths with the highest cumulative rewards.

Key Contributions

1. Non-Compensatory PRM Architecture: SCPRM introduces a cumulative reward mechanism using conditional probabilities that mathematically ensures risky steps impose persistent penalties, preventing the "risk compensation" effect common in additive PRMs.
2. Schema-Aware Future Estimation: The authors propose a learnable future reward module that utilizes logical schemas parsed from queries. By computing schema distance, the model learns query logic and overcomes the inherent length bias of traditional conditional reward models.
3. Strong Performance on Risk-Sensitive KGQA: The framework demonstrates superior performance in medical and legal domains, where logical safety is paramount, as well as on general multi-hop reasoning benchmarks.

Experimental Results

The authors evaluated SCPRM on three datasets: Medical KGQA, Legal KGQA, and Complex WebQuestions (CWQ).

• Pairwise Ranking (RQ1): SCPRM consistently outperformed baseline reward models (Vanilla PRM, PQM, CRM) across all datasets and base LLMs (Qwen2.5-1.5B, Qwen3-4B, Llama3.1-8B). For instance, on the Medical KGQA dataset with Llama3.1-8B, SCPRM achieved 95.34% pairwise accuracy, significantly higher than the next best baseline (91.83%).
• KG Reasoning Performance (RQ2 & RQ3): Integrated with MCTS, SCPRM-MCTS achieved an average improvement of 1.18% in Hits@k over strong baselines. Notably, SCPRM-MCTS using an 8B parameter model (Llama3.1-8B) outperformed prompting-based methods utilizing the larger, closed-source GPT-4o-mini on Medical and Legal KGQA tasks.
• Ablation Study: Removing the Cumulative Reward (w/o CR) caused severe degradation in step-level accuracy (dropping ~10%), confirming its role in penalizing historical errors. Removing the Future Reward (w/o FR) significantly impacted end-to-end Hits@k (dropping ~12% in Legal KGQA), highlighting its role in path selection and lookahead.
• Robustness: The model demonstrated graceful degradation under noisy schema conditions (up to 50% noise), retaining robust performance without catastrophic failure.

Significance and Claims

The paper claims that SCPRM addresses a critical gap in reasoning evaluation for risk-sensitive domains. By explicitly modeling the non-compensatory nature of errors and incorporating schema-aware lookahead, the model provides a more accurate and reliable evaluation of reasoning paths than existing PRMs. The authors emphasize that this approach allows smaller, open-source models to outperform larger, closed-source models in complex, multi-hop KG reasoning tasks, particularly where logical safety and adherence to domain constraints are essential. The work bridges heuristic search (A*) and modern learning systems, offering a principled framework for guiding LLMs in structured knowledge environments.