PaperlayerView exact revision on arXiv

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

arXiv:2409.12917v2cs.LG

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.

    Source: §1 · pp. 1–2

  • 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.

    Source: §1 · p. 2, §4 · p. 6, §5 · p. 8

  • Results

    15.6% and 9.1% absolute self-correction gains are reported relative to base Gemini models on MATH and HumanEval, respectively.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    SCoRe attains significantly positive intrinsic self-correction performance without external feedback during inference.

    Source: §1 · pp. 2–3, §7 · p. 13

    SCoRe was trained for only one round of iterative self-correction, so later rounds may be less effective.

    Source: §7 · p. 14

Abstract

from arXiv · show

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.Source: §1 · p. 1
  • Existing approaches rely on prompting, multiple models, or teacher supervision, while SCoRe uses one model for answering and correction without oracle feedback.Source: §1 · p. 2
  • SFT-based approaches can suffer from distribution shift or behavior collapse, limiting transfer to self-correction under the learned model’s own mistakes.Source: §1 · p. 2
  • SCoRe is a multi-turn RL approach for teaching LLMs to correct their own mistakes.Source: §1 · pp. 2–3
  • 15.6% and 9.1% absolute gains in self-correction are reported on MATH and HumanEval, respectively, relative to base Gemini models.Source: §1 · pp. 2–3

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.Source: §2 · p. 3
  • Naive prompting can degrade intrinsic self-correction, partly because prior studies make different assumptions about available feedback or initial prompts.Source: §2 · p. 3
  • Fine-tuning methods commonly use revision demonstrations from human annotators or stronger models, whereas this work targets entirely self-generated training data.Source: §2 · p. 3
  • The paper distinguishes its focus from RL machinery research by treating self-correction itself as an RL problem.Source: §2 · p. 3
  • Figure 2 presents traces where SCoRe corrects either an arithmetic mistake or a reasoning error.Source: Figure 2 · p. 4

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.Source: §3 · p. 4
  • At test time, the model must detect and correct mistakes without access to the oracle reward used during training.Source: §3 · p. 4
  • The objective maximizes verifier-based correctness at the end of l + 1 turns, with this paper evaluating two turns.Source: §3 · p. 5
  • Unlike standard single-response training, the multi-turn objective trains attempts simultaneously through the summed intermediate rewards.Source: §3 · p. 5
  • Base RL uses REINFORCE with a KL-divergence penalty against a fixed reference policy.Source: §3 · p. 5
  • Self-correction is measured with first- and second-attempt accuracy and net improvement between attempts.Source: §3 · p. 5

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.Source: §4 · p. 6
  • 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.Source: §4 · p. 6
  • 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.Source: §4 · p. 8
  • 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.Source: Figure 5 · p. 7, §4 · p. 8
  • 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.Source: §4 · pp. 6, 8

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.Source: §5 · p. 8
  • 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.Source: §5 · pp. 8–9

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.Source: Figure 7 · p. 9, §5 · p. 9, §5.1 · pp. 9–10

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.Source: §5.1 · p. 10

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.Source: Figure 7 · p. 9, §5.2 · p. 10

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.Source: §5 · p. 8, §5.3 · p. 10

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.Source: §6 · p. 11

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.Source: §6.1 · p. 12
  • 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.Source: §6.1 · p. 12

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.Source: §6.2 · p. 12

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.Source: §6.3 · p. 13, Table 4 · p. 13
  • 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.Source: §6.3 · p. 13

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.Source: §7 · p. 13
  • 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.Source: §7 · p. 13

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.Source: §7 · p. 14

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.Source: §A.1 · p. 18
  • The base Gemini 1.5 Flash model is noisy across attempts and never surpasses its first-attempt performance.Source: §A.1 · p. 18
  • Pair-SFT does not improve beyond the second attempt.Source: §A.1 · p. 18
  • SCoRe’s performance plateaus as revision attempts increase, likely because the response distribution shifts quickly.Source: §A.1 · p. 18
  • γ=0.8 with α=1.0 does not prevent multi-turn RL from suffering the same non-correcting behavior collapse as standard multi-turn RL.Source: §A.2 · pp. 18–19

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.

Self-correction instructions

  • The MATH self-correction instruction asks the model to correct an earlier solution and output only the final solution.Source: C. Prompts · p. 19

Programming examples

  • The programming examples include largest-integer selection with a heap queue and testing whether an integer is non-prime.Source: C. Prompts · p. 20

Self-correction instructions

  • The MBPP/HumanEval instruction asks the model to correct possible code errors and output only the final Python program.Source: C. Prompts · p. 20

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.

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

Programming examples

Mathematical examples

Related papers

Loading 2409.12917v2…