Entrenamiento de Modelos de Lenguaje para Auto-Corrección mediante Aprendizaje por Refuerzo

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

Resumen

La capacidad de autocorrección es una característica altamente deseable de los grandes modelos de lenguaje (LLMs), sin embargo, se ha encontrado consistentemente que es poco efectiva en los LLMs modernos. Los enfoques existentes para entrenar la autocorrección requieren múltiples modelos o dependen de un modelo más capaz u otras formas de supervisión. Con este fin, desarrollamos un enfoque de aprendizaje por refuerzo en línea multi-turno, SCoRe, que mejora significativamente la capacidad de autocorrección de un LLM utilizando datos generados internamente. Para construir SCoRe, primero demostramos que las variantes del ajuste fino supervisado (SFT) en trazas de corrección generadas por el modelo fuera de línea son insuficientes para inculcar el comportamiento de autocorrección. En particular, observamos que el entrenamiento a través de SFT sufre de una falta de coincidencia de distribución entre los datos de entrenamiento y las respuestas propias del modelo o prefiere implícitamente solo un cierto modo de comportamiento de corrección que a menudo no es efectivo en el momento de la prueba. SCoRe aborda estos desafíos entrenando bajo la distribución propia del modelo de trazas de corrección generadas internamente y utilizando regularización apropiada para dirigir el proceso de aprendizaje hacia la adquisición de una estrategia de autocorrección efectiva en el momento de la prueba en lugar de simplemente ajustar respuestas de alto valor de recompensa para una indicación dada. Esta regularización prescribe ejecutar una primera fase de RL en un modelo base para generar una inicialización de política menos susceptible al colapso y luego utilizar un bono de recompensa para amplificar la autocorrección durante el entrenamiento. Cuando se aplica a los modelos Gemini 1.0 Pro y 1.5 Flash, encontramos que SCoRe logra un rendimiento de autocorrección de vanguardia, mejorando la autocorrección de los modelos base en un 15.6% y un 9.1% respectivamente en los benchmarks MATH y 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.