Sam Babichenko

Reading notes / Reinforcement learning for reasoning

Let's Verify Step by Step

Lightman et al., 2024

Search and scoring
NNsolutions sampled per problem in best-of-NN
m;m; pjp_jnumber of steps in a solution; PRM probability that step jj is correct
∏jpj\prod_{j} p_jPRM score of a solution
Models
PRMlarge\mathrm{PRM}_{\rm large}the PRM trained on PRM800K, used as the labeler at small scale

Context. A reward model that scores solutions chooses among samples at test time, and in reinforcement learning: training a model to raise a reward on outputs it samples itself it is what the the language model being trained, viewed as a distribution over responses to each prompt learns to please, so its reliability limits both uses. For multi-step math it can be supervised in two ways. An a model trained to predict, from a whole solution, whether its final answer is correct needs only final answers, which a program can check; a a model trained on human judgments of each reasoning step, predicting for every step whether it is correct needs a person to judge every step. Uesato et al. (2022) found the two about equal on grade-school math. Lightman et al. (2024) redo the comparison with a stronger model (a large language model, used here as pretrained, before any tuning from human feedback), far more human labels and the harder a benchmark of 12,500 competition mathematics problems (7,500 training, 5,000 test), each with a final answer that can be checked automatically benchmark, and judge a reward model by how often the solution it ranks highest among many samples is correct.

Telling a reward model where a solution first goes wrong teaches it more than telling it only that the solution is wrong. An ORM has to work out for itself which step spoiled a wrong solution, and on hard problems, where most samples are wrong somewhere, a bare "wrong" carries little information; step labels give it the location. The advantage grows with the number of samples searched, since a larger pool holds more convincing wrong solutions for a weak verifier to pick. Process labels also cost less than they seem to, since checking a solution without a known answer means finding its first mistake anyway, and in a small-scale simulation, choosing to label the solutions the current PRM wrongly trusts was about 2.6 times as data-efficient as labeling uniform samples.

Setting. All solutions come from one fixed the model that writes the candidate solutions; it is never trained with RL in this paper, GPT-4 further trained on about 1.5B math tokens and then for one epoch on its own correct solutions to MATH training problems, only so that it writes one step per line. "Supervision" here always means supervision of the reward model. Labelers mark each step of a sampled solution positive (correct and reasonable), negative (incorrect or unreasonable) or neutral (ambiguous), and why: it keeps the comparison fair; on a correct solution both kinds of supervision say every step is right, and on a wrong one process labels add only the location of the first mistake. The resulting dataset, the released set of about 800,000 step labels on 75,000 solutions to 12,000 problems, collected from this generator, trains a PRM that predicts the label as a token after each step, so one forward pass scores every step. Write the PRM's probability that step jj of a solution is correct, with neutral counted as correct for its output; a solution with mm steps is scored by the probability that every step is correct, PRM score=∏j=1mpj.\text{PRM score}=\prod_{j=1}^{m}p_j. The ORM predicts, at every token, whether the final answer is right, and its score is read at the last token. A reward model is evaluated by sampling NN solutions to a problem, keeping the one the reward model scores highest, and grading only that one by its final answer on 500 MATH test problems drawn uniformly at random (the other 4,500 test problems were used for training), against answering with the final answer that occurs most often among the NN samples, with no reward model as the baseline.

Main findings (all empirical). With N=1860N=1860 samples per problem the PRM picks a correct solution for 78.2% of the test problems, against 72.4% for the ORM and 69.6% for majority voting, and it leads at every NN by a margin that widens as NN grows. That comparison is not controlled, since the two models train on different data, so a small-scale experiment with the large PRM standing in for the labelers isolates the kind of label. On identical solutions, process labels beat outcome labels, whether those come from final-answer checks or from the large PRM, at every dataset size. Choosing which solutions to label by spending labels on the examples the current model is most likely to get wrong, here solutions the PRM scores highly whose final answer is wrong makes process labels about 2.6 times as data-efficient. On 234 new AP and AMC exam problems, released after the pretraining data was collected, the PRM again leads at best-of-100, 72.9% against 63.8% for the ORM and 61.3% for voting.

Consequences and limits. For reranking sampled solutions in math, process labels give a more reliable verifier than final-answer labels, and per label, a small amount of process supervision matches a large amount of outcome supervision, and process supervision keeps improving as it is scaled, which the authors offer as the reason Uesato et al. found a tie at smaller scale. The paper never uses the PRM as an RL reward, where a learned verifier is pushed hardest. DeepSeek-R1 (Guo et al., 2025) later used only rule-based answer checks for its reasoning RL, having found neural reward models, outcome- or process-based, open to a policy raising a learned reward by exploiting the reward model's errors while the quality the reward stands for stalls or falls in large-scale RL. PRM800K has been widely used to train and evaluate PRMs since, and later work replaced human step labels with automatic ones, labeling a step good when sampled completions from it tend to reach the right answer (Math-Shepherd, Wang et al., 2024). The authors argue that step-level rewards favor reasoning a person can endorse, so that here a safer method also performs better, an argument they make without measuring it. The scope is narrow. It covers one domain with checkable answers and one model family, and the test set may have appeared in pretraining (the authors argue contamination would affect all methods alike, and the fresh exam problems agree).

Step 1

At scale the PRM is the better verifier, and the gap widens with NN. The large ORM trains on 100 uniform samples per problem graded by final answer, an order of magnitude more solutions than PRM800K and none in common with it; training it on PRM800K instead would be skewed, since active learning made that set mostly wrong-answer solutions, and adding uniform samples to PRM800K did not help the ORM. At every NN above 1, up to 1860, the PRM is ahead, and the gap is widest at the largest NN. By difficulty, measured as the generator's own solve rate, the gap shows in every fifth of the test set. On the easiest fifth the ORM's accuracy even falls slightly as NN grows, because a larger pool contains wrong solutions that fool it, while the PRM's holds; on the hardest fifth, sampling more helps most, since a correct and convincing solution may be rare. Scoring details matter little. Taking the minimum over steps instead of the product, or counting neutral steps as wrong, moves best-of-1860 accuracy only between 77.4% and 78.2%, and combining PRM scores with voting did not help noticeably. What this step cannot rule out is that the gap comes from the different training sets. The figure below uses a simple model of this trade-off, not the paper's data.

Step 2

With the solutions held fixed, process labels still win, and final-answer labels are noisier. Human labels are too expensive for ablations, so the large PRM, the PRM trained on PRM800K in Step 1, used here as the labeler for smaller models, labels samples from a small generator whose base model had about 200 times less pretraining compute. From the same 1 to 200 samples per problem it produces three kinds of supervision: process labels, which mark steps up to the first one it judges wrong (a step counts as wrong when its probability of being negative exceeds 20%, since PRMlarge\mathrm{PRM}_{\rm large} slightly favors positive labels); outcome labels from the same judgments, a solution being correct when every step passes; and outcome labels from final-answer checks. Judged by best-of-500, the process-supervised reward models beat both outcome-supervised series at every dataset size. The two outcome series also differ. PRMlarge\mathrm{PRM}_{\rm large}'s outcome labels beat final-answer checks, because a final-answer check rewards solutions that reach the right answer by wrong reasoning. The authors take PRMlarge\mathrm{PRM}_{\rm large}'s outcome labels as the fairer baseline, since the false positives that hurt final-answer checks are, they argue, more common on MATH than in other domains. This rules out the training-set difference of Step 1 and isolates the information in the label; whether the same holds for human labels at matched scale is not tested.

Step 3

Labeling the solutions the current PRM wrongly trusts makes labeling about 2.6 times as data-efficient. A a solution that the current PRM scores highly but whose final answer is wrong, so the PRM must be wrong about at least one of its steps is the most informative thing to label, since the PRM is known to be mistaken about it. In the small-scale experiment, a selector PRM trained on one sample per problem scores 1000 samples per problem; each training set takes 80% of its samples from the most convincing wrong-answer ones and 20% from the most convincing of the rest, so that it is not made only of wrong answers, and PRMlarge\mathrm{PRM}_{\rm large} labels them. Comparing the slopes of the fitted lines of accuracy against the number of samples labeled, with and without this selection, gives the factor of 2.6. The largest selected set (200 per problem) falls slightly below the trend, which the authors put down to drawing 200 from a pool of only 1000; retraining the selector during collection was unstable and gave no gain. The human data were collected in ten rounds that surfaced the current PRM's highest-scoring wrong-answer solutions, retraining the PRM between rounds, and ended up with only 13.2% of solutions correct in that phase while 74.1% of the labeled steps were correct, so the PRM still sees many correct steps.

Step 4

The ranking holds on problems the models cannot have seen. To check that the results are not memorization of MATH problems available online, the large models are evaluated at best-of-100 on recent AP Calculus, Chemistry and Physics and AMC 10/12 exams, released after the pretraining data was collected. The PRM leads on every exam (86.7%, 80.0%, 86.7% and 53.2%, against 68.9%, 68.9%, 77.8% and 49.1% for the ORM), while majority voting beats the ORM on the three AP exams and does worst on AMC, at 32.8%. The generator's low solve rates on many MATH problems the PRM handles, often a few percent, are a further sign that it has not memorized them.

Terms

ORM
outcome-supervised reward model: predicts from the whole solution whether the final answer is correct
PRM
process-supervised reward model: predicts for each step whether it is correct, from human step labels
PRM score
∏jpj\prod_j p_j, the probability that every step is correct; the minimum over steps does almost as well
Best-of-NN
keep the highest-scoring of NN samples and grade its final answer; the paper's measure of a reward model
Majority voting
the most common final answer among NN samples; the baseline with no reward model
Labels stop at the first error
makes process and outcome labels carry the same information except the location of the first mistake
False positive
a solution with the right final answer and wrong reasoning, which final-answer checks reward
Convincing wrong-answer solution
scored highly by the current PRM but with a wrong final answer; the PRM must be wrong somewhere in it
Active learning gain
about 2.6 times the data efficiency of uniform labeling in the small-scale experiment, from the slopes of fitted lines
PRMlarge\mathrm{PRM}_{\rm large} as labeler
a large PRM supplies process or outcome labels for small models, so the label type can be varied alone
Headline numbers
best-of-1860 on 500 MATH problems: PRM 78.2%, ORM 72.4%, majority voting 69.6%

Quiz

  1. The PRM's lead over the ORM widens as the number of samples NN grows. Which explanation does the note support?
  2. The active-learning training sets take 20% of their samples from the most convincing solutions regardless of answer. What is this share for?
  3. What does this paper leave untested about using a PRM in practice?

Sources

  • Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's Verify Step by Step. ICLR, 2024.