Reward models for small language models

An unpublished paper I co-wrote for my Machine Learning class at Lawrence Tech, finished in December 2025. It asks whether a small, self-hosted language model can be steered toward more reliable step-by-step reasoning by reward models trained on a stronger model’s judgments.

Read the paper (PDF) All projects

Type Class paper, unpublished
Class Machine Learning, Lawrence Tech
When Fall 2025, finished December
Authors Two, including me
Models 3B generator · two 1.5B reward models
Teacher ChatGPT
Stack Python · Hugging Face Transformers · Reward models · Chain of thought

Two steps of a sampled solution to the wolf, goat and cabbage puzzle, each in a card with its process reward model score: a long, rambling step that restates and muddles the answer, outlined in red with a score of 0.08, and the clean step “Take the goat across the river,” outlined in green with a score of 0.99

From the paper’s Figure 2: the process model’s score for each step of one sampled solution. The step that rambles and muddles the answer scores 0.08; the clean one after it, 0.99.

The question

Asking a model to think step by step helps it on multi-step problems, but the steps can read well and still be wrong. Earlier work showed that a process reward model, which judges each step, can pick out good reasoning better than an outcome reward model, which judges only the final answer. That work used very large models and very large labelled datasets.

We wanted to know how much of it carries over to a small open model with a few billion parameters, the kind you can host and fine-tune yourself, using labels from a stronger model as the teacher.

Pipeline

How it works

Generate. A 3B-parameter instruction-tuned model writes several chain-of-thought solutions to each math or logic question, one step per line, ending with an answer.
Label. ChatGPT splits each solution into steps and marks every step, and the final answer, as correct or incorrect.
Train. Two 1.5B-parameter models get a classification head and are fine-tuned on those labels with Hugging Face Transformers: the process model on single steps, the outcome model on whole solutions.
Select. At answer time the small model samples several solutions, and the reward models pick one. The process model’s score for a solution is the average of its step scores.

We compared six ways of answering: a direct answer with no reasoning, a single chain of thought, a majority vote over the sampled answers, picking by the outcome model, picking by the process model, and a critic mode that asks the small model to revise its weakest step. None of them changes the small model itself; everything happens when it answers.

The questions were 20 seed problems in arithmetic, algebra, probability and basic logic. With 8 sampled solutions each, that gave 2,216 labelled steps and 160 labelled solutions, split 90% for training and 10% for validation.

Outcome

What we found

The process model judged steps well on its validation set. When a solution went wrong partway through, its scores usually dropped sharply at the first bad step and stayed low.

Reward models on their validation sets
Model	Accuracy	Precision	Recall	F1	ROC-AUC
Process model (steps)	0.98	0.98	0.99	0.99	0.99
Outcome model (whole solutions)	1.00	1.00	1.00	1.00	–

From the paper’s Table I. The outcome model’s perfect scores are inflated: its dataset was small and mostly correct solutions, so we treat them as illustrative.

Picking by the process model helped most where majority vote fails. When the small model made the same subtle algebra mistake in most of its samples, majority vote chose the wrong answer with confidence, while the process model scored those solutions low and chose one whose every step held up. On our small held-out set, the direct answer did worst, a single chain of thought did a little better, majority vote better again, and picking by the process model was usually best or tied for best.

Limits

What it does not show

A small dataset. 20 course-level questions is enough to test the pipeline, not to measure it, so we compared the strategies by the errors they made rather than by statistics.
One teacher. Every label came from ChatGPT, with no human check.
Confident about the wrong things. The process model sometimes scored well-formed steps highly even when they had nothing to do with the question.
The small model never learned. All of the gains come from choosing better among what it already writes.

The next steps in the paper are to test on established benchmarks such as GSM8K and MATH, to compare other small models and reward model sizes, and to use the two reward models as the reward for reinforcement learning with PPO, so the small model itself learns to reason better rather than only being filtered.

The full paper is here as a PDF (7 pages).