Treinando Modelos de Linguagem para se Auto-Corrigirem via Aprendizado por Reforço

Training Language Models to Self-Correct via Reinforcement Learning

September 19, 2024

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

Resumo

A capacidade de autocorreção é uma característica altamente desejável em grandes modelos de linguagem (LLMs), no entanto, tem sido consistentemente considerada amplamente ineficaz nos LLMs modernos. As abordagens existentes para treinar a autocorreção exigem múltiplos modelos ou dependem de um modelo mais capaz ou de outras formas de supervisão. Para isso, desenvolvemos uma abordagem de aprendizado por reforço online de múltiplas etapas, SCoRe, que melhora significativamente a capacidade de autocorreção de um LLM usando dados totalmente autogerados. Para construir o SCoRe, primeiro demonstramos que variantes de ajuste fino supervisionado (SFT) em traços de correção gerados pelo modelo offline são insuficientes para instilar o comportamento de autocorreção. Em particular, observamos que o treinamento via SFT sofre de uma disparidade de distribuição entre os dados de treinamento e as próprias respostas do modelo ou prefere implicitamente apenas um certo modo de comportamento de correção que frequentemente não é eficaz no momento do teste. O SCoRe aborda esses desafios treinando sob a distribuição do próprio modelo de traços de correção autogerados e utilizando regularização apropriada para direcionar o processo de aprendizado para aprender uma estratégia de autocorreção que seja eficaz no momento do teste, ao invés de simplesmente ajustar respostas de alta recompensa para uma determinada solicitação. Essa regularização prescreve a execução de uma primeira fase de RL em um modelo base para gerar uma inicialização de política menos suscetível a colapsos e, em seguida, usar um bônus de recompensa para amplificar a autocorreção durante o treinamento. Quando aplicado aos modelos Gemini 1.0 Pro e 1.5 Flash, descobrimos que o SCoRe alcança um desempenho de autocorreção de última geração, melhorando a autocorreção dos modelos base em 15,6% e 9,1%, respectivamente, nos benchmarks MATH e 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.