Het trainen van taalmodellen om zichzelf te corrigeren via versterkend leren.

Training Language Models to Self-Correct via Reinforcement Learning

September 19, 2024

Auteurs: 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

cs.AI

Samenvatting

Zelfcorrectie is een zeer gewenste capaciteit van grote taalmodellen (LLM's), maar het is consequent gebleken dat het grotendeels ineffectief is in moderne LLM's. Bestaande benaderingen voor het trainen van zelfcorrectie vereisen meerdere modellen of vertrouwen op een capabeler model of andere vormen van toezicht. Om dit te bereiken, ontwikkelen we een multi-turn online reinforcement learning (RL) benadering, SCoRe, die aanzienlijk de zelfcorrectiemogelijkheid van een LLM verbetert met volledig zelf gegenereerde gegevens. Om SCoRe te bouwen, tonen we eerst aan dat varianten van supervised fine-tuning (SFT) op offline model-genererende correctietraces onvoldoende zijn om zelfcorrectiegedrag in te prenten. In het bijzonder merken we op dat training via SFT ofwel lijdt aan een distributieverschil tussen de trainingsgegevens en de eigen reacties van het model, ofwel impliciet de voorkeur geeft aan slechts een bepaalde modus van correctiegedrag die vaak niet effectief is bij testtijd. SCoRe pakt deze uitdagingen aan door te trainen onder de eigen distributie van zelf gegenereerde correctietraces van het model en door passende regularisatie te gebruiken om het leerproces te sturen naar het aanleren van een zelfcorrectiestrategie die effectief is bij testtijd in plaats van simpelweg het passen van hoog-reward reacties voor een gegeven prompt. Deze regularisatie schrijft voor om een eerste fase van RL uit te voeren op een basismodel om een beleidsinitialisatie te genereren die minder vatbaar is voor ineenstorting en vervolgens een beloningsbonus te gebruiken om zelfcorrectie te versterken tijdens de training. Wanneer toegepast op de Gemini 1.0 Pro en 1.5 Flash modellen, vinden we dat SCoRe state-of-the-art zelfcorrectieprestaties behaalt, waarbij de zelfcorrectie van de basismodellen respectievelijk met 15,6% en 9,1% verbetert op de MATH en HumanEval benchmarks.

English

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.