Source-linked AI summary
Training Language Models to Self-Correct via Reinforcement Learning
Aviral Kumar, Vincent Zhuang, Rishabh Agarwal, Yi Su, John D Co-Reyes, Avi Singh, Kate Baumli, Shariq Iqbal, Colton Bishop, Rebecca Roelofs, Lei M Zhang, Kay McKinney, Disha Shrivastava, Cosmin Paduraru, George Tucker, Doina Precup, Feryal Behbahani, Aleksandra Faust
TL;DR
Modern LLMs rarely self-correct reliably, and existing methods often require multiple models, stronger models, or extra supervision. SCoRe addresses this with multi-turn online RL on self-generated traces, using staged regularization to improve intrinsic self-correction. It reports 15.6% and 9.1% gains over base Gemini models on MATH and HumanEval, respectively.
Problem
Modern LLMs have limited intrinsic self-correction, while existing training methods often depend on multiple models, stronger models, or additional supervision.
Method
SCoRe trains a single model with multi-turn online RL on self-generated correction traces, using a two-stage design and reward shaping to address distribution shift and behavior collapse.
Results
15.6% and 9.1% absolute self-correction gains are reported relative to base Gemini models on MATH and HumanEval, respectively.
Takeaways & Limitations
SCoRe attains significantly positive intrinsic self-correction performance without external feedback during inference.
Takeaways & Limitations
SCoRe was trained for only one round of iterative self-correction, so later rounds may be less effective.
Abstract
from arXiv · showhide
Self-correction is a highly desirable capability of large language models (LLMs), yet it has consistently been found to be largely ineffective in modern LLMs. Current methods for training self-correction typically depend on either multiple models, a more advanced model, or additional forms of supervision. To address these shortcomings, we develop a multi-turn online reinforcement learning (RL) approach, SCoRe, that significantly improves an LLM's self-correction ability using entirely self-generated data. To build SCoRe, we first show that variants of supervised fine-tuning (SFT) on offline model-generated correction traces are often insufficient for instilling self-correction behavior. In particular, we observe that training via SFT falls prey to either a distribution mismatch between mistakes made by the data-collection policy and the model's own responses, or to behavior collapse, where learning implicitly prefers only a certain mode of correction behavior that is often not effective at self-correction on test problems. SCoRe addresses these challenges by training under the model's own distribution of self-generated correction traces and using appropriate regularization to steer the learning process into learning a self-correction behavior that is effective at test time as opposed to fitting high-reward responses for a given prompt. This regularization process includes an initial phase of multi-turn RL on a base model to generate a policy initialization that is less susceptible to collapse, followed by using a reward bonus to amplify self-correction. With Gemini 1.0 Pro and 1.5 Flash models, we find that SCoRe achieves state-of-the-art self-correction performance, improving the base models' self-correction by 15.6% and 9.1% respectively on MATH and HumanEval.
1. Introduction
Modern LLMs often lack reliable intrinsic self-correction, despite possessing useful reasoning knowledge. SCoRe trains a single model with multi-turn reinforcement learning and achieves positive self-correction gains on MATH and HumanEval.
- Intrinsic self-correction remains severely lacking in current LLMs, especially without external input.
- Existing approaches rely on prompting, multiple models, or teacher supervision, while SCoRe uses one model for answering and correction without oracle feedback.
- SFT-based approaches can suffer from distribution shift or behavior collapse, limiting transfer to self-correction under the learned model’s own mistakes.
- SCoRe is a multi-turn RL approach for teaching LLMs to correct their own mistakes.
- 15.6% and 9.1% absolute gains in self-correction are reported on MATH and HumanEval, respectively, relative to base Gemini models.
2. Related Work
Prior intrinsic self-correction work uses prompting, fine-tuning, or additional supervision, while this paper studies self-generated training without larger models or humans. The paper illustrates self-correction through arithmetic and reasoning-error traces.
- Intrinsic self-correction without external input is substantially harder than settings with environmental feedback.
- Naive prompting can degrade intrinsic self-correction, partly because prior studies make different assumptions about available feedback or initial prompts.
- Fine-tuning methods commonly use revision demonstrations from human annotators or stronger models, whereas this work targets entirely self-generated training data.
- The paper distinguishes its focus from RL machinery research by treating self-correction itself as an RL problem.
- Figure 2 presents traces where SCoRe corrects either an arithmetic mistake or a reasoning error.
3. Preliminaries and Problem Setup
The paper formulates intrinsic self-correction as a multi-turn policy problem: a model receives the problem, prior attempts, and auxiliary instructions, then produces a final response. Training uses verifier rewards across attempts.
- The policy receives the problem, previous model attempts, and auxiliary instructions, then produces the most correct response possible.
- At test time, the model must detect and correct mistakes without access to the oracle reward used during training.
- The objective maximizes verifier-based correctness at the end of l + 1 turns, with this paper evaluating two turns.
- Unlike standard single-response training, the multi-turn objective trains attempts simultaneously through the summed intermediate rewards.
- Base RL uses REINFORCE with a KL-divergence penalty against a fixed reference policy.
- Self-correction is measured with first- and second-attempt accuracy and net improvement between attempts.
4. SFT on Self-Generated Data is Insufficient for Self-Correction
SFT-based self-correction methods improve some outcomes but remain insufficient because they suffer from distribution shift, conservative or collapsed correction behavior, and limited positive self-correction gains.
- 1.8% gain in self-correction is achieved by Pair-SFT despite a substantially higher Δ(t1, t2) than the base model.The improvement mainly comes from reducing correct-to-incorrect changes rather than correcting incorrect first attempts.
- STaR does not clearly learn when to modify responses, while adding correct-to-correct data yields only a 0.4% change in Δ(t1, t2).The added data helps STaR substantially but produces only a small change in the self-correction gap.
- Standard multi-turn RL improves both attempts but leaves Δ(t1, t2) unchanged because the responses become tightly coupled.The resulting policy becomes biased against changing its answer, reducing exploration and learning progress.
- Offline SFT improves correction on fixed first attempts but degrades on self-generated attempts, revealing distribution shift.The learner optimizes correction accuracy on static data while its self-correction accuracy falls on its own responses.
- SFT-based methods suffer from distribution shift and behavior collapse when learning self-correction.Training on on-policy data can address distribution shift but not behavior collapse.
5. SCoRe: Self-Correction via Multi-Turn Reinforcement Learning
SCoRe addresses distribution shift and behavior collapse by training self-correction with on-policy, multi-turn RL, a two-stage initialization, and reward shaping. The method explicitly favors changing an initial response when appropriate rather than optimizing both attempts toward non-correcting behavior.
- Motivation: Standard multi-turn RL improves both attempts but leaves their difference unchanged, converging to behavior that is overly biased against changing responses.This produces no self-correction ability and resembles the collapse observed with STaR.
- Motivation: RL can learn either genuine improvement between attempts or the equally optimal training-set strategy of producing the best first response followed by no correction.Only the first strategy generalizes to new problems, making behavior collapse a central challenge.
- Method overview: SCoRe uses Stage I to initialize a policy that improves second attempts while preserving base-model behavior on first attempts, then jointly optimizes both attempts in Stage II.This two-stage design decouples the attempts before multi-turn RL and reduces susceptibility to behavior collapse.
- Stage I: Training an Initialization that Decouples Attempts: Stage I is critical: compared with standard multi-turn RL, it more effectively decouples responses and leads to better Stage II performance.A strict first-attempt KL penalty helps avoid shifting first-turn responses.
- Stage II: Multi-Turn RL with Reward Shaping: Stage II applies shaped rewards that favor progress toward self-correction rather than merely producing the best first response followed by superficial edits.The reward shaping biases optimization toward a self-correcting solution when ground-truth rewards alone permit coupled responses.
- Takeaways and Implications: SCoRe’s core insight is to use self-generated online data and regularization to make nuanced self-correction more attractive than degenerate behavior modes.This directly targets both distribution shift and behavior collapse.
6. Experimental Evaluation
SCoRe is evaluated on math and coding tasks with sequential self-correction attempts, alongside prompting and fine-tuning baselines and component ablations. It improves intrinsic self-correction across benchmarks, benefits from sequential self-correction at inference time, and depends on its multi-stage, shaped-reward, on-policy design.
- Experimental Setup: SCoRe is evaluated on MATH, MBPP, and HumanEval using two sequential attempts, with MBPP-R additionally measuring offline repair.The experiments compare self-correction accuracy and, for code generation, offline correction performance.
- Benchmark Results: SCoRe improves MATH self-correction over the base Gemini 1.5 Flash model by 15.6% and raises Accuracy@t2 by 23.0%.Its intrinsic self-correction gain is 4.4%, while it fixes 14.5% of initially incorrect answers versus 9.5% for the base model.
- Benchmark Results: SCoRe improves MBPP-R from 47.3% to 60.6% and achieves a 12.2% intrinsic self-correction delta on HumanEval, 9% higher than the base model.Pair-SFT performs similarly on static repair but degrades the base model in the self-correction setting.
- Inference-Compute Scaling with Self-Correction: With a 32-solution budget, parallel sampling gains 7.4% accuracy, while sequential sampling combined with self-correction gains 10.5%.The comparison supports sampling K solutions followed by one self-correction round rather than sampling 2K solutions in parallel.
- Ablation Studies: Single-turn training improves first-attempt performance but produces negative self-correction deltas, while removing Stage I lowers Δ(t1,t2) by 2% and Accuracy@t2 by 3%.Removing reward shaping also hurts performance.
- Ablation Studies: Replacing REINFORCE with STaR in Stage II yields significantly lower absolute performance without visible self-improvement gains.The results suggest on-policy samples are especially important for this multi-turn self-correction problem.
7. Discussion, Limitations, and Conclusion
The paper concludes that SCoRe uses two-stage RL and reward shaping to obtain positive intrinsic self-correction while avoiding distribution shift and behavior collapse. Its main scope limitation is that it was trained for only one round of iterative self-correction.
- Conclusion: SCoRe is presented as one of the first methods to attain significantly positive intrinsic self-correction through multi-turn RL.The method analyzes failure modes including non-correcting behavior and distribution shift.
- Conclusion: The two-stage design and reward shaping help prevent collapse into ineffective self-corrective behavior.The conclusion identifies these mechanisms as central to training a self-correcting strategy.
- Limitations: SCoRe was not trained beyond one round of iterative self-correction, so subsequent rounds may be less effective than the first.The authors identify training with more than two attempts as future work.
A.1. Scaling to Multiple Attempts
SCoRe is evaluated over repeated self-correction attempts beyond its two-turn training setup. Unlike the base model and Pair-SFT, SCoRe improves slightly after two turns, though gains plateau.
- SCoRe increases slightly beyond two self-correction turns despite training over only two attempts.
- The base Gemini 1.5 Flash model is noisy across attempts and never surpasses its first-attempt performance.
- Pair-SFT does not improve beyond the second attempt.
- SCoRe’s performance plateaus as revision attempts increase, likely because the response distribution shifts quickly.
- γ=0.8 with α=1.0 does not prevent multi-turn RL from suffering the same non-correcting behavior collapse as standard multi-turn RL.
B. Additional Experiment Details
This section records SCoRe’s MATH and MBPP hyperparameters and includes the self-correction instructions and programming examples used in the experiments.
- Table 5 lists SCoRe hyperparameters for MATH and MBPP.
- Training uses Gemini 1.0 Pro with Adam, a 1e-5 learning rate, 1500 steps, batch size 128, and sampling temperature 1.0.
- Self-correction instructions: The MATH self-correction instruction asks the model to correct an earlier solution and output only the final solution.
- Programming examples: The programming examples include largest-integer selection with a heap queue and testing whether an integer is non-prime.
- Self-correction instructions: The MBPP/HumanEval instruction asks the model to correct possible code errors and output only the final Python program.
D. Qualitative Analysis of SCoRe
Qualitative examples show SCoRe refining responses through complete rewrites or targeted revisions, especially for computational mistakes, and sometimes correcting within a turn.
- SCoRe revises responses either by rewriting the entire solution or by preserving correct parts while changing incorrect ones.
- SCoRe is especially adept at revising computational mistakes.
- The model sometimes shows more steps in computations and manipulations to increase its probability of producing a correct answer.
- SCoRe occasionally performs self-correction within a single turn, including in a MATH example.
E. Examples of SCoRe Self-Correction
The examples cover mathematical and programming tasks in which SCoRe’s second-turn response corrects an initial response or code solution.
- Example coverage: The examples include cyclic-pattern checking, even–odd digit counting, fraction simplification, polynomial roots, and minimization problems.
- Programming examples: In the even–odd digit example, the second-turn code changes iteration from str(num) to str(abs(num)).
- Mathematical examples: In the fraction example, SCoRe changes the final answer from 4/21 on turn one to 8/21 on turn two.
- Mathematical examples: In the minimization example, the second-turn solution rejects the incorrect value 6 and concludes that the minimum is 9.
- Mathematical examples: In the circle-intersection example, the first-turn probability 1/12 is corrected to 1/3 on turn two.