Entraînement de modèles de langage pour s'auto-corriger via l'apprentissage par renforcement.

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

Résumé

L'autocorrection est une capacité hautement souhaitable des grands modèles de langage (LLM), cependant, elle s'est avérée largement inefficace dans les LLM modernes. Les approches existantes pour entraîner l'autocorrection nécessitent soit plusieurs modèles, soit un modèle plus performant ou d'autres formes de supervision. Dans cette optique, nous développons une approche d'apprentissage par renforcement en ligne multi-tours, SCoRe, qui améliore significativement la capacité d'autocorrection d'un LLM en utilisant uniquement des données auto-générées. Pour construire SCoRe, nous montrons d'abord que les variantes du fine-tuning supervisé (SFT) sur des traces de correction générées hors ligne par le modèle sont insuffisantes pour inculquer un comportement d'autocorrection. En particulier, nous observons que l'entraînement via SFT souffre soit d'un désaccord de distribution entre les données d'entraînement et les propres réponses du modèle, soit préfère implicitement seulement un certain mode de comportement de correction qui n'est souvent pas efficace au moment du test. SCoRe relève ces défis en s'entraînant sous la distribution propre du modèle de traces d'autocorrection auto-générées et en utilisant une régularisation appropriée pour orienter le processus d'apprentissage vers l'acquisition d'une stratégie d'autocorrection efficace au moment du test, plutôt que de simplement ajuster des réponses à haute récompense pour une instruction donnée. Cette régularisation prescrit d'exécuter une première phase de RL sur un modèle de base pour générer une initialisation de politique moins sujette à l'effondrement, puis d'utiliser un bonus de récompense pour amplifier l'autocorrection pendant l'entraînement. Lorsqu'appliqué aux modèles Gemini 1.0 Pro et 1.5 Flash, nous constatons que SCoRe atteint des performances d'autocorrection de pointe, améliorant respectivement de 15,6% et 9,1% les capacités d'autocorrection des modèles de base sur les benchmarks MATH et HumanEval.

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.