Training Language Models to Self-Correct via Reinforcement Learning
Authors: name: Aviral Kumar, name: Vincent Zhuang, name: Rishabh Agarwal, name: Yi Su, name: John D Co-Reyes, name: Avi Singh, name: Kate Baumli, name: Shariq Iqbal, name: Colton Bishop, name: Rebecca Roelofs, name: Lei M Zhang, name: Kay McKinney, name: Disha Shrivastava, name: Cosmin Paduraru, name: George Tucker, name: Doina Precup, name: Feryal Behbahani, name: Aleksandra Faust

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. Existing approaches for training self-correction either require multiple
models or rely on a more capable model or other forms of supervision. To this
end, 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
insufficient for instilling self-correction behavior. In particular, we observe
that training via SFT either suffers from a distribution mismatch between the
training data and the model's own responses or implicitly prefers only a
certain mode of correction behavior that is often not effective at test time.
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 strategy that is
effective at test time as opposed to simply fitting high-reward responses for a
given prompt. This regularization prescribes running a first phase of RL on a
base model to generate a policy initialization that is less susceptible to
collapse and then using a reward bonus to amplify self-correction during
training. When applied to 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 the MATH and
HumanEval benchmarks.