Smaller Models, Smarter Rewards: A Two-Sided Approach to Process and Outcome Rewards

AI summary

Overview

• Research area: Reward modeling for large language models, specifically value estimation for Python code generation and reasoning-trace evaluation (process reward models vs. outcome reward models).
• Technical level: Intermediate. Readers need familiarity with decoder-only transformers, fine-tuning, reward modeling, and pass@k evaluation, but the paper explains its notation and training setup explicitly.
• One-sentence scope: The paper investigates whether small Phi-4 models (3.8B and 14B) can be turned into reward models that both score whole code rollouts and estimate the correctness of intermediate reasoning steps.

What This Paper Is About

Generating correct Python code remains difficult for LLMs, and improving reasoning models on coding tasks typically requires an intermediate reward model that can judge outputs or intermediate steps. The core problem is that outcome reward models (ORMs) give only a single end-of-generation score and suffer from delayed credit assignment, while process reward models (PRMs) give stepwise feedback but require costly fine-grained human labels. The authors ask whether small, open Phi-4 models can be fine-tuned into a single reward model that serves both roles, acting as a lightweight code critic and as an evaluator of partial reasoning traces.

Key Contributions

1. Architecture replacement: The last layer of the Phi-4 models is replaced by a single linear regression output with a sigmoid function, producing a decoder-only value-head model that estimates the success probability of a partial reasoning trace using only preceding tokens. The authors report these architectures fine-tune into reward models in just two episodes.
2. Critic capability: The trained reward model serves as a lightweight critic and intermediate-step evaluator, selecting correct rollouts among multiple candidates and improving the accuracy of the selected Python code by over 20% relative to the unassisted baseline.
3. Behavioral analysis: The paper extensively analyzes the model's scoring of full rollouts and intermediate reasoning steps, including confidence tracking across token positions and the amount of partial code required before predictions beat random guessing.
4. Size comparison: A direct comparison between the 3.8B Phi-4-mini and the 14B Phi-4 shows that model size is the primary limitation, with the smaller model falling significantly behind.

Main Findings

• Baseline generation performance: Without any reward model interference, generation with Phi-4-mini-instruct yields pass@1 of 45%, pass@3 of 65%, and pass@10 of 84%.
• Critic-driven selection gains: Applying the critic and selecting the best rollouts raises pass@1 to 50%–55% and pass@3 to 72%–78%. The paper frames this as a ceiling improvement of approximately 22.2% relative for pass@1 on best(3 choose 1) (55 vs. 45) and approximately 20% relative for pass@3 on best(10 choose 3) (78 vs. 65).
• Larger model wins: After training on the same datasets, the 14B Phi-4 substantially outperforms the 3.8B Phi-4-mini, with most prediction results reaching 60% or above. On imbalanced test data, the 14B model reaches 73.8% accuracy when trained on balanced data and 71.7% when trained on imbalanced data; the 3.8B model reaches 64.0% and 66.0% respectively.
• Balanced test data is harder: On the balanced test set, the 14B model reaches 65.8% (balanced training) and 60.5% (imbalanced training) accuracy, while the 3.8B model reaches 55.3% and 51.9%.
• Balancing helps but only moderately: Training accuracy on the imbalanced dataset is similar between balanced and imbalanced training. On the balanced test dataset, models trained on balanced data perform substantially better, but when the task is to rank or select the best rollouts, the two training regimes produce close accuracies. The authors conclude balancing has an impact, but a rather limited one.
• Process evaluation requires roughly half the trace: The model needs at least 50% of the partial code before it judges better than a random guess. The authors interpret this as a lower bound on the preceding code tokens needed for the model to have predictive power.
• Confidence dynamics: The model's confidence rises after seeing the generated output and aligns with ground truth, but the authors observe a warming-up stage before around 100 tokens, attributed to the Chain-of-Thought prompt that first elaborates the problem and only later writes code.
• Likely error detection without self-correction: Given that performance only improves around the 50% percentile, the authors suspect the model can catch errors in code but lacks the self-correction capability for already incorrect reasoning steps.
• Dataset composition: The APPS-derived dataset totals 4,449 problems, split into 465 test problems and 3,984 train problems. The unbalanced training dataset has 66,924 samples and the balanced version has 110,016.
• Small models can do reward modeling: The conclusion states that large-scale capacity is not a prerequisite for reward modeling, arguing Phi-4 family models (14B parameters) can serve as effective reward models, described as comparing favorably against state-of-the-art systems such as GPT-5 and Claude 4 Sonnet at vastly larger scales.

Methodology in Plain English

The researchers started with the APPS coding challenge benchmark, which mixes problems that read and write via standard input/output with problems that ask for a completed function, and wrote two different prompting templates to match (both reproduce a full Chain-of-Thought request before code). To produce training data, they generated one main rollout with Phi-4-mini, then identified six branching positions—the token positions where the model assigned the lowest probability, since these are the points where the reasoning path is most likely to change. From each of those six positions they generated six additional rollouts, giving 6 × 6 = 36 rollouts per problem. Every rollout was labeled by executing the APPS unit tests in a sandbox, yielding a binary correctness verdict.

The state value they wanted to predict is the expected correctness of a completion generated from a given token prefix. To estimate it, they took the Phi-4-mini (3.8B) and Phi-4 (14B) models, replaced the output classification layer with a linear regression layer, applied a sigmoid so the output is a probability, and trained with binary cross-entropy loss. Only the last layers were fine-tuned: layers 236–248 for the 4B model and layers 180–192 for the 14B model. Training used the Adam optimizer at a learning rate of 1e-4. Reported batch sizes are 64 and 24 for the two models, although Table 1 and the appendix assign these numbers to opposite models. Training ran for two episodes. Because the raw rollouts are class-imbalanced per problem—risking a model that learns problem difficulty rather than reasoning correctness—the authors also built a balanced dataset by oversampling the smaller class within each problem, for both train and test splits, and trained and evaluated on both versions. Compute ran on an Azure Standard_NC96ads_A100_v4 instance with 96 AMD EPYC 7V13 CPU cores, 880 GB RAM, 256 GB disk, and 4 Nvidia A100 GPUs with 320 GB GPU memory, totaling around 95 GPU hours.

Why This Matters

The work suggests that reward modeling for code does not require frontier-scale models, which matters for teams with limited compute who still want to apply reinforcement learning or best-of-n selection to their own code models. It also provides evidence about when a reward model becomes useful during a generation—roughly halfway through the token stream—which has direct implications for early stopping and rollout guidance, and it highlights a limitation: the model appears to detect errors in code but not to correct faulty reasoning that already occurred.

Real-world applications:

• Code generation assistants: Using a compact critic to rank multiple candidate solutions and surface the most likely correct one to a developer.
• Training-data filtering: Serving as a lightweight data quality filter to identify correct and incorrect generations at scale.
• Inference-time compute scaling: Applying best-of-n selection in resource-constrained deployments where a frontier judge model is too expensive.
• Reinforcement learning for code: Supplying process and outcome reward signals for reasoning models without expensive intermediate-step human annotation.

Industry relevance: The paper is a collaboration involving SAP Lab and notes collaboration with the Stanford Human-Centered AI Institution, indicating direct interest from enterprise software organizations in efficient, self-hosted code evaluation models.

Future Directions

• Reducing the compute cost of data generation: The current setup of 36 rollouts per problem takes over 48 hours for one pass on a 4×A100 Azure instance, mainly throttled by the output token throughput of Phi-4-mini, so tuning branching hyperparameters was left to future work.
• Solving the cold-start problem: The method depends on a ground-truth correct rollout ("tree trunk") from which branching occurs, and samples where the Phi model produced no correct rollout were excluded.
• Finding the point of diminishing returns for branching: Preliminary experiments with smaller branching factors such as 6x6 indicate a positive impact from increased branching, but broader branching was not computationally tractable, so the saturation point is unknown.
• Post-training dynamics: The setup assumes the predicted correctness probability distribution does not shift between the base policy and the policy after post-training; the authors leave exploration of this shift to future work.
• APPS distribution alignment: Because the APPS train and test splits (each 5,000 samples) differ substantially in problem format and complexity, the authors performed their own train-test split and defer a deeper investigation of task-formulation effects to future work.

Target Audience

Researchers and engineers working on reward modeling, reinforcement learning from human feedback, and inference-time search for code generation will benefit most, particularly those interested in whether small open models can replace large proprietary judges. Practitioners building code assistants or data-filtering pipelines with tight compute budgets will also find the size comparison and confidence-tracking analysis directly actionable. Readers unfamiliar with reward models, process versus outcome supervision, or pass@k metrics should expect to consult the related-work section for background.

Authors’ abstract

Generating high-quality code remains a challenge for Large Language Models (LLMs). For the evolution of reasoning models on this task, reward models are a necessary intermediate step. These models judge outcomes or intermediate steps. Decoder-only transformer models can be turned into reward models by introducing a regression layer and supervised fine-tuning. While it is known that reflection capabilities generally increase with the size of a model, we want to investigate whether state-of-the-art small language models like the Phi-4 family can be turned into usable reward models blending the consideration of process rewards and outcome rewards. Targeting this goal, we construct a dataset of code samples with correctness labels derived from the APPS coding challenge benchmark. We then train a value-head model to estimate the success probability of intermediate outputs. Our evaluation shows that small LLMs are capable of serving as effective reward models or code evaluation critics, successfully identifying correct solutions among multiple candidates. Using this critic, we achieve over a 20% improvement in the search capability of the most accurate code out of multiple generations.