Training Language Models to Self-Correct via Reinforcement Learning

📄 arXiv: 2409.12917v2 📥 PDF

作者: 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.LG

发布日期: 2024-09-19 (更新: 2024-10-04)

💡 一句话要点

提出SCoRe，通过强化学习显著提升大语言模型的自我纠错能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 自我纠错 强化学习 在线学习 奖励函数

📋 核心要点

  1. 现有大语言模型的自我纠错能力不足，且依赖额外资源或监督，限制了其应用。
  2. SCoRe方法通过在线强化学习，利用模型自生成数据进行训练，提升自我纠错能力。
  3. 实验表明，SCoRe在Gemini模型上显著提升了自我纠错性能，优于现有方法。

📝 摘要（中文）

自我纠错是大语言模型(LLM)一个非常理想的能力，但现有LLM在这方面表现不佳。目前的自我纠错训练方法通常依赖于多个模型、更高级的模型或额外的监督形式。为了解决这些不足，我们开发了一种多轮在线强化学习(RL)方法，SCoRe，它完全使用自生成数据来显著提高LLM的自我纠错能力。我们首先证明，在离线模型生成的纠错轨迹上进行监督微调(SFT)的变体通常不足以灌输自我纠错行为。特别是，我们观察到，通过SFT进行训练容易陷入数据收集策略所犯错误与模型自身响应之间的分布不匹配，或者陷入行为崩溃，即学习隐式地只偏好某种纠错行为模式，而这种模式通常不能有效地纠正测试问题。SCoRe通过在模型自身生成的纠错轨迹分布下进行训练，并使用适当的正则化来引导学习过程，使其学习一种在测试时有效的自我纠错行为，而不是拟合给定提示的高奖励响应，从而解决了这些挑战。通过Gemini 1.0 Pro和1.5 Flash模型，我们发现SCoRe实现了最先进的自我纠错性能，在MATH和HumanEval上分别将基础模型的自我纠错能力提高了15.6%和9.1%。

🔬 方法详解

问题定义：现有大语言模型的自我纠错能力较弱，并且依赖于多模型、更高级模型或额外监督。监督微调（SFT）在离线数据上的训练容易出现分布不匹配或行为崩溃问题，导致模型无法有效进行自我纠错。

核心思路：SCoRe的核心思路是利用强化学习，让模型在自身生成的纠错轨迹分布上进行训练，避免分布不匹配问题。通过正则化技术，引导模型学习有效的自我纠错行为，而不是简单地拟合高奖励的响应。

技术框架：SCoRe采用多轮在线强化学习框架。首先，在一个基础模型上进行多轮RL，生成一个策略初始化，使其不易崩溃。然后，使用奖励加成（reward bonus）来放大自我纠错行为。整个流程包括：1) 初始策略训练；2) 自我纠错数据生成；3) 强化学习训练；4) 模型评估。

关键创新：SCoRe的关键创新在于使用在线强化学习，直接在模型自身生成的纠错数据上进行训练，避免了离线数据带来的分布不匹配问题。此外，奖励加成机制能够有效引导模型学习有效的自我纠错策略。

关键设计：SCoRe的关键设计包括：1) 奖励函数的设计，用于衡量自我纠错的有效性；2) 正则化策略，防止模型陷入行为崩溃；3) 探索策略，鼓励模型探索不同的纠错路径；4) 奖励加成系数，用于调整自我纠错行为的权重。具体的参数设置和网络结构细节在论文中未明确给出，属于未知信息。

🖼️ 关键图片

img_0

📊 实验亮点

SCoRe在MATH和HumanEval数据集上取得了显著的性能提升。在Gemini 1.0 Pro模型上，自我纠错能力提升了15.6%；在Gemini 1.5 Flash模型上，提升了9.1%。这些结果表明，SCoRe是一种有效的自我纠错训练方法，优于现有的监督微调方法。

🎯 应用场景

该研究成果可应用于各种需要大语言模型进行自我纠错的场景，例如自动问答、代码生成、文本摘要等。通过提高模型的自我纠错能力，可以减少人工干预，提高工作效率，并提升用户体验。未来，该方法有望推广到更多的大语言模型和任务中。

📄 摘要（原文）

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. Current methods for training self-correction typically depend on either multiple models, a more advanced model, or additional forms of supervision. To address these shortcomings, 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 often insufficient for instilling self-correction behavior. In particular, we observe that training via SFT falls prey to either a distribution mismatch between mistakes made by the data-collection policy and the model's own responses, or to behavior collapse, where learning implicitly prefers only a certain mode of correction behavior that is often not effective at self-correction on test problems. 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 behavior that is effective at test time as opposed to fitting high-reward responses for a given prompt. This regularization process includes an initial phase of multi-turn RL on a base model to generate a policy initialization that is less susceptible to collapse, followed by using a reward bonus to amplify self-correction. With 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 MATH and HumanEval.