Technical Summary: Outcome-Supervised Process Reward Modeling via Learnable Credit Assignment (LCA)

1. Problem Statement

Process Reward Models (PRMs) enhance the reasoning capabilities of Large Language Models (LLMs) by providing fine-grained, step-level feedback. However, training PRMs typically requires expensive human annotations for every reasoning step. To address this, Outcome-Supervised PRMs attempt to learn from final-answer correctness alone, avoiding step-level labels.

The core challenge in this setting is credit assignment: attributing the final outcome (correct/incorrect) to the specific reasoning steps that caused it. Existing approaches suffer from fundamental deficiencies:

• Uniform Assignment: Assumes all steps contribute equally to the final outcome (e.g., via additive aggregation). This fails to distinguish between correct precursor steps and the specific step where an error occurred.
• Causal Assignment: Propagates influence backward from the final answer using Monte Carlo rollouts (e.g., estimating the probability of reaching a correct answer from a given step). This conflates a step's logical correctness with the expected success of downstream completions, making it difficult to localize the actual first error.

Both methods fail to anchor credit in the actual correctness of individual steps, hindering the model's ability to identify process errors accurately.

2. Methodology: LCA Framework

The authors propose Outcome-Supervised Process Reward Modeling via Learnable Credit Assignment (LCA), a framework that jointly learns credit assignment and reward modeling without step-level supervision.

2.1 Theoretical Formulation: Weakest Link Assignment

LCA is grounded in the Weakest Link Assignment principle: a reasoning chain is as strong as its weakest link. A trajectory is correct if and only if every step is correct; it is incorrect if any step is wrong.

• This principle transforms the outcome-supervised PRM problem into a Multiple Instance Learning (MIL) problem.
• Bag: The full reasoning trajectory.
• Instances: The reasoning prefixes (steps) within the trajectory.
• Label: The trajectory-level correctness (positive if the trajectory contains an error, negative otherwise).
• Assumption: The Standard Multi-Instance Assumption holds: a bag is positive if and only if it contains at least one positive instance (an erroneous step).

2.2 The Challenge of Dependence

Standard MIL methods often assume instances are independent or weakly dependent. However, reasoning trajectories exhibit autoregressive, nested, and monotonic properties. In a trajectory, the final prefix alone determines the bag label, causing standard pooling techniques (like Max Pooling or Attention Pooling) to degenerate. The model tends to ignore intermediate steps, effectively becoming an Outcome Reward Model (ORM) that only predicts the final state.

2.3 Softmax-Weighted-Sum (SWS) Pooling

To resolve the mutual dependence between credit assignment and reward modeling, LCA introduces Softmax-Weighted-Sum (SWS) pooling.

• Mechanism: Given instance-level error probabilities $p = (p_1, \dots, p_T)$, the bag-level prediction is a weighted sum:
  $$\hat{p} = \sum_{t=1}^T w_t p_t, \quad \text{where } w_t = \frac{e^{p_t/\alpha}}{\sum_{j=1}^T e^{p_j/\alpha}}$$
• Role of Temperature ($\alpha$): The parameter $\alpha$ controls the entropy of the weight distribution.
  • As $\alpha \to 0$, SWS approaches Max Pooling (strict Weakest Link).
  • As $\alpha \to \infty$, SWS approaches Average Pooling (Uniform Assignment).
• Advantages:
  1. Non-negligible Credit: Unlike Max Pooling, SWS ensures every instance receives a strictly positive weight, preventing the model from ignoring intermediate steps.
  2. Structural Prior: The softmax weighting naturally encodes the Standard Multi-Instance Assumption, where instances with higher predicted error probabilities (positive instances) dominate the bag prediction.
  3. Bayes Consistency: The authors prove that under mild assumptions, minimizing the expected loss with SWS pooling leads to a Bayes-optimal instance classifier, even without instance-level labels.

2.4 Training Objective

The PRM is trained end-to-end using a trajectory-level cross-entropy loss. The model outputs per-prefix error probabilities, which are aggregated via SWS pooling to match the trajectory-level ground truth (final answer correctness). A stop-gradient operation is applied to the softmax weights to stabilize training.

3. Key Contributions

1. Insightful Analysis: The paper identifies that existing outcome-supervised PRMs (Uniform and Causal) fail to anchor credit in step correctness. It proposes Weakest Link Assignment as a principled alternative that aligns with the "first-error" criterion used in PRM evaluation.
2. Technical Innovation (LCA):
   • Formalizes outcome-supervised PRM as an MIL problem under the Weakest Link principle.
   • Introduces SWS pooling, an MIL technique tailored for the strong dependence and redundancy of reasoning states, embedding a smoothed Weakest Link assignment as a structural prior.
   • Provides a theoretical proof of Bayes consistency for the learning objective at the instance level.
3. Empirical Performance: Demonstrates that LCA consistently outperforms state-of-the-art outcome-supervised PRMs across multiple tasks (error identification, best-of-N scaling) and backbone models (Llama, Qwen).

4. Experimental Results

The authors evaluated LCA on mathematical reasoning tasks using the Math-Shepherd dataset for training and ProcessBench for evaluation.

• Error Identification (RQ1): On ProcessBench, LCA achieved a new state-of-the-art.
  • With a 7B backbone, LCA surpassed the second-best method by 5.5% in F1 score.
  • With a 3B backbone, it outperformed all competing 3B methods by a substantial margin.
  • Notably, LCA achieved these results without the heavy computational cost of Monte Carlo sampling required by baselines like MathShepherd, SCAN, or OmegaPRM.
• Test-Time Scaling (RQ2): In Best-of-N sampling experiments on MATH-500, LCA consistently achieved the highest average scores across different generator models (Mistral, Llama, Qwen) and sample sizes ($N=8, 32, 64$).
• Ablation Studies:
  • Temperature ($\alpha$): Performance varies with $\alpha$. Lower $\alpha$ (closer to Max Pooling) benefits simple datasets (GSM8K), while higher $\alpha$ helps complex datasets with many process errors (OlympiadBench), highlighting a trade-off between theoretical consistency and empirical effectiveness.
  • Pooling Techniques: SWS pooling with $\alpha=0.2$ significantly outperformed other MIL pooling techniques (Instance, Embedding, Attention, Conjunctive) and approached the performance of fully supervised PRMs trained on human-annotated step labels.

5. Significance and Claims

The paper claims that LCA offers a principled alternative for fine-grained process supervision in complex reasoning tasks. By resolving the "chicken-and-egg" dilemma between credit assignment and reward modeling through the MIL framework and SWS pooling, LCA produces PRMs that accurately attribute credit to the steps that determine the outcome.

The authors emphasize that their method:

• Closes the training-test gap: By aligning the training objective (Weakest Link) with the evaluation metric (First Error Detection).
• Reduces computational overhead: It avoids the expensive Monte Carlo rollouts required by other outcome-supervised methods.
• Enables generalization: The learned PRMs show strong generalization across different backbone models and reasoning difficulties.

The paper concludes modestly, acknowledging that while LCA outperforms other outcome-supervised methods, it still lags behind PRMs trained on human-annotated step-level labels. They identify learning with noisy labels (where correct answers are reached via flawed reasoning) and integrating PRMs into Reinforcement Learning (RL) as important directions for future work.