Training von Sprachmodellen zur Selbstkorrektur mittels Reinforcement Learning

Training Language Models to Self-Correct via Reinforcement Learning

September 19, 2024

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

Zusammenfassung

Die Selbstkorrektur ist eine äußerst wünschenswerte Fähigkeit großer Sprachmodelle (LLMs), wurde jedoch in modernen LLMs bisher als weitgehend unwirksam erachtet. Bestehende Ansätze zur Schulung der Selbstkorrektur erfordern entweder mehrere Modelle oder setzen ein leistungsfähigeres Modell oder andere Formen der Aufsicht voraus. Zu diesem Zweck entwickeln wir einen Multi-Turn-Online-Verstärkungslernansatz (RL), SCoRe, der die Selbstkorrekturfähigkeit eines LLMs erheblich verbessert, indem ausschließlich selbstgenerierte Daten verwendet werden. Um SCoRe aufzubauen, zeigen wir zunächst, dass Varianten des überwachten Feintunings (SFT) anhand offline vom Modell generierter Korrekturtrajektorien nicht ausreichen, um ein Selbstkorrekturverhalten zu vermitteln. Insbesondere beobachten wir, dass das Training über SFT entweder unter einem Verteilungsfehler zwischen den Trainingsdaten und den eigenen Antworten des Modells leidet oder implizit nur eine bestimmte Art des Korrekturverhaltens bevorzugt, die oft nicht effektiv zur Testzeit ist. SCoRe begegnet diesen Herausforderungen, indem es unter der eigenen Verteilung des Modells von selbstgenerierten Korrekturtrajektorien trainiert und geeignete Regularisierung verwendet, um den Lernprozess in die Richtung einer Selbstkorrekturstrategie zu lenken, die zur Testzeit effektiv ist, anstatt einfach hochbelohnte Antworten für eine bestimmte Eingabe anzupassen. Diese Regularisierung schreibt vor, eine erste Phase des RL auf einem Basismodell durchzuführen, um eine Richtlinieninitialisierung zu generieren, die weniger anfällig für Zusammenbrüche ist, und dann einen Belohnungsbonus zu verwenden, um die Selbstkorrektur während des Trainings zu verstärken. Bei Anwendung auf die Modelle Gemini 1.0 Pro und 1.5 Flash stellen wir fest, dass SCoRe eine Spitzenleistung bei der Selbstkorrektur erzielt und die Selbstkorrektur der Basismodelle auf den MATH- und HumanEval-Benchmarks um jeweils 15,6 % bzw. 9,1 % verbessert.

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.