📄 Enhancing Mathematical Reasoning in LLMs by Stepwise Correction

Summary

Best-of-N decoding samples several independent solutions to a math problem and keeps the highest-scoring one. Because the samples get no feedback, the model often makes the same mistake in all of them, and an early error tends to carry through every later step. The authors propose Stepwise Correction (STEPCO), a prompting framework that repairs one reasoning path instead of sampling many.

The LLM first writes a tagged, step-by-step solution. A process-supervised verifier (PSV), fine-tuned from Llama-3-8B, then estimates for each step the chance that it leads to the correct answer. The first step below a threshold is reported back to the LLM with its score. The LLM keeps the earlier steps and rewrites that step and the ones after it. This cycle repeats for a bounded number of rounds.

The PSV's training data is labelled automatically. Each step's score is the share of continuations through that step that reach the gold answer.

With GPT-4o, STEPCO beats Best-of-N on eight math benchmarks while using far fewer tokens. Targeted, verifier-guided revision is therefore both more accurate and cheaper than repeated sampling.

Key points

• With GPT-4o, STEPCO averages 94.1 accuracy over eight math datasets. That is +2.4 over Best-of-10, with 77.8% fewer tokens (about 612 vs. 2,760 tokens per question).
• A single verify-then-revise round already reaches 92.3, which is +0.6 over Best-of-10. Most questions are fixed after one revision.
• The PSV training data is built with two demonstrations at each step. D⁺ steers the model toward correct steps, and D⁻ toward wrong or alternative ones, so the dataset stays balanced.
• The method was tested with GPT-3.5-Turbo, GPT-4o and Llama-3-8B. It was also extended to open-domain QA and commonsense reasoning, where it again beat the baselines.
• The code is public at wzy6642.github.io/stepco.github.io.