Preprint
Large Language Models

Training LLMs to Self-Correct via RL

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, D. Precup, Feryal M. P. Behbahani, Aleksandra Faust
September 19, 2024International Conference on Learning Representations448 citations
arXivPDFDOI

448

Citations

38

Influential Citations

International Conference on Learning Representations

Venue

2024

Year

Abstract

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.

Analysis

Why This Paper Matters

Self-correction is a critical capability for LLMs to produce reliable outputs, yet prior methods have been largely ineffective. This paper addresses a fundamental gap by showing that supervised fine-tuning on offline correction traces fails due to distribution mismatch and behavior collapse. SCoRe introduces a multi-turn online RL approach that trains the model on its own generated corrections, avoiding the need for external supervision or stronger models. This is significant because it makes self-correction training more accessible and scalable, potentially enabling LLMs to improve their outputs iteratively in real-world applications.

The paper also provides a deep analysis of why SFT fails, which is valuable for the community. By identifying behavior collapse—where the model learns a single correction mode that is not effective—the authors highlight a subtle issue that can undermine training. This insight can guide future research on training strategies for complex behaviors.

Technical Contributions

  • Multi-turn online RL framework: SCoRe trains the model to generate an initial response and then a corrected response, using RL to optimize the improvement from first to second attempt.
  • Two-stage training: The first stage uses multi-turn RL on a base model to create a policy initialization that is less prone to collapse. The second stage applies a reward bonus that amplifies self-correction, steering the policy toward effective correction behavior.
  • Self-generated data: The method relies entirely on the model's own outputs, eliminating the need for multiple models or external supervision.
  • Analysis of SFT failures: The paper provides a systematic study showing that SFT on offline traces suffers from distribution mismatch and behavior collapse, motivating the need for online RL.

Results

SCoRe achieves state-of-the-art self-correction performance on two benchmarks: MATH and HumanEval. On MATH, it improves the base Gemini 1.0 Pro model's self-correction by 15.6%, and on HumanEval, it improves Gemini 1.5 Flash by 9.1%. These gains are substantial and demonstrate the effectiveness of the method. The paper also compares against SFT baselines, showing that SCoRe outperforms them, likely due to its online training and regularization.

Significance

This work has broad implications for LLM training. By enabling self-correction without external supervision, it paves the way for more autonomous and reliable AI systems. The method could be applied to other tasks beyond math and code, such as reasoning, planning, and even multi-step decision making. The insights on behavior collapse and distribution mismatch are also relevant to other RL fine-tuning scenarios. Overall, SCoRe represents a significant step toward making LLMs self-improving, which is a key goal for AI research.

large-language-modelsself-correctionreinforcement-learningmulti-turn-rlonline-learningfine-tuning
Large Language ModelsFine TuningReinforcement LearningPromptingComputer Vision

Related Papers