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
Published: 2024-09-19 17:16:21

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