Table of Contents
Fetching ...

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

Jan Niklas Groeneveld, Xi Qin, Alexander Schaefer, Yaad Oren

TL;DR

This work demonstrates that compact decoder-only Phi-4 models can serve as effective dual-role reward models for code generation, handling both process and outcome evaluations via a regression-based value head. By constructing a carefully sampled dataset of 36 rollouts per APPS problem and training with a sigmoid-valued head, the authors show that small models can reliably distinguish correct vs. incorrect solutions and substantially improve rollout selection when used as a critic. They further show that the reward model can also assess intermediate reasoning steps, providing early reliability signals after about half of the generation. Despite promising results, the study notes significant computational demands, data-balance effects, and distributional considerations that guide future work on scaling and generalization.

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.

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

TL;DR

This work demonstrates that compact decoder-only Phi-4 models can serve as effective dual-role reward models for code generation, handling both process and outcome evaluations via a regression-based value head. By constructing a carefully sampled dataset of 36 rollouts per APPS problem and training with a sigmoid-valued head, the authors show that small models can reliably distinguish correct vs. incorrect solutions and substantially improve rollout selection when used as a critic. They further show that the reward model can also assess intermediate reasoning steps, providing early reliability signals after about half of the generation. Despite promising results, the study notes significant computational demands, data-balance effects, and distributional considerations that guide future work on scaling and generalization.

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.
Paper Structure (20 sections, 2 equations, 6 figures, 2 tables)

This paper contains 20 sections, 2 equations, 6 figures, 2 tables.

Figures (6)

  • Figure 1: The CDF plot on the left is generated through kernel density estimation with Gaussian kernels. Then this model is used to generate the success probability estimations on the right.
  • Figure 2: Accuracy of percentiles in rollout generation from four combinations of the 14B model trained and tested on balanced/imbalanced data.
  • Figure 3: The prompt we use for problems that require I/O operations. We use the chat template which the model was trained with, laying out both system and user prompts. In the system prompt, we give instructions and give an piece of example code. The example code is for a straightforward string problem that is neither in the test nor train dataset. The model is supposed to continue with its thoughts and its code after the assistant tag.
  • Figure 4: The prompt we use for problems that require a complete function as the solution. In the assistant part of the message, we directly give the function signature to force the model to produce a valid function, and also not deviate from the given signature. Moreover, we force the model to start with a docstring. Using this prompt, we significantly reduced the number of malformed generations.
  • Figure 5: Example of one piece of training sample.
  • ...and 1 more figures