Imitating Language via Scalable Inverse Reinforcement Learning
作者: Markus Wulfmeier, Michael Bloesch, Nino Vieillard, Arun Ahuja, Jorg Bornschein, Sandy Huang, Artem Sokolov, Matt Barnes, Guillaume Desjardins, Alex Bewley, Sarah Maria Elisabeth Bechtle, Jost Tobias Springenberg, Nikola Momchev, Olivier Bachem, Matthieu Geist, Martin Riedmiller
分类: cs.LG, cs.AI, cs.CL, stat.ML
发布日期: 2024-09-02 (更新: 2024-12-09)
备注: Published at NeurIPS 2024
💡 一句话要点
提出基于逆强化学习的语言模型微调方法，提升生成质量和多样性。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 逆强化学习 模仿学习 语言模型微调 序列生成 奖励函数学习
📋 核心要点
- 现有语言模型微调主要依赖最大似然估计，忽略了序列生成的内在结构。
- 论文提出基于逆强化学习的模仿学习方法，提取奖励函数并直接优化序列生成。
- 实验表明，该方法在保持生成多样性的同时，提升了任务性能，优于传统监督微调。
📝 摘要（中文）
当前，语言模型训练主要依赖于模仿学习，包括预训练、监督微调以及强化学习中人类反馈的起始条件。最大似然估计（MLE）因其简单性和可扩展性而成为主流范式。然而，更广义的模仿学习可以更有效地利用自回归生成的序列结构。本文着重研究模仿学习的逆强化学习（IRL）视角，提取奖励并直接优化序列，而非单个token的似然。本文从新的角度出发，将逆软Q学习重新表述为MLE的时间差分正则化扩展，从而在MLE和IRL之间建立起有原则的联系，并允许在监督微调（SFT）环境中权衡复杂性与生成性能和多样性。研究发现，基于IRL的模仿学习具有明显的优势，尤其是在最大化任务性能的同时保持多样性，使得IRL成为固定SFT数据集上的强大替代方案，即使没有在线数据生成。对IRL提取的奖励函数的分析进一步表明，通过更紧密地集成监督和基于偏好的LLM后训练，可以获得更鲁棒的奖励函数。
🔬 方法详解
问题定义：现有语言模型微调方法，如监督微调（SFT），主要采用最大似然估计（MLE）来预测下一个token。这种方法将序列生成问题简化为独立的token预测，忽略了序列的整体结构和长期依赖关系。此外，MLE可能导致模型过度拟合训练数据，降低生成的多样性。
核心思路：论文的核心思路是将模仿学习问题转化为逆强化学习（IRL）问题。通过从专家数据（例如，SFT数据集）中学习奖励函数，模型可以更好地理解专家行为的内在目标，并生成更符合目标且更具多样性的序列。这种方法不再仅仅关注单个token的预测，而是关注整个序列的优化。
技术框架：该方法的核心是将逆软Q学习重新表述为MLE的时间差分正则化扩展。具体来说，首先使用专家数据训练一个初始语言模型。然后，使用逆强化学习算法从该模型生成的序列中提取奖励函数。该奖励函数反映了专家行为的内在目标。最后，使用该奖励函数来微调语言模型，使其生成的序列更符合专家行为。整体框架可以看作是在传统MLE的基础上，增加了一个奖励函数学习和序列优化步骤。
关键创新：最重要的技术创新点在于将逆软Q学习与MLE联系起来，提出了时间差分正则化扩展。这种方法在MLE的基础上引入了时间差分学习的思想，使得模型能够更好地利用序列的结构信息。与传统的MLE方法相比，该方法能够生成更符合目标且更具多样性的序列。
关键设计：关键的设计包括：1) 使用逆软Q学习来提取奖励函数，该方法能够处理连续状态空间和动作空间；2) 将逆软Q学习重新表述为MLE的时间差分正则化扩展，从而在MLE和IRL之间建立起有原则的联系；3) 使用提取的奖励函数来微调语言模型，使其生成的序列更符合专家行为。损失函数的设计需要平衡MLE损失和时间差分正则化项，以避免模型过度拟合或欠拟合。
🖼️ 关键图片
📊 实验亮点
实验结果表明，基于IRL的模仿学习方法在保持生成多样性的同时，提升了任务性能。具体而言，该方法在固定SFT数据集上取得了优于传统监督微调方法的结果，并且在没有在线数据生成的情况下也表现出色。对IRL提取的奖励函数的分析表明，该方法能够学习到更鲁棒的奖励函数，从而提升语言模型的泛化能力。
🎯 应用场景
该研究成果可应用于各种需要高质量和多样性文本生成的场景，例如对话系统、文本摘要、机器翻译和代码生成。通过学习更鲁棒的奖励函数，可以提升语言模型在复杂任务中的表现，并减少对大量人工标注数据的依赖。未来，该方法有望应用于更广泛的序列生成任务，例如视频生成和机器人控制。
📄 摘要（原文）
The majority of language model training builds on imitation learning. It covers pretraining, supervised fine-tuning, and affects the starting conditions for reinforcement learning from human feedback (RLHF). The simplicity and scalability of maximum likelihood estimation (MLE) for next token prediction led to its role as predominant paradigm. However, the broader field of imitation learning can more effectively utilize the sequential structure underlying autoregressive generation. We focus on investigating the inverse reinforcement learning (IRL) perspective to imitation, extracting rewards and directly optimizing sequences instead of individual token likelihoods and evaluate its benefits for fine-tuning large language models. We provide a new angle, reformulating inverse soft-Q-learning as a temporal difference regularized extension of MLE. This creates a principled connection between MLE and IRL and allows trading off added complexity with increased performance and diversity of generations in the supervised fine-tuning (SFT) setting. We find clear advantages for IRL-based imitation, in particular for retaining diversity while maximizing task performance, rendering IRL a strong alternative on fixed SFT datasets even without online data generation. Our analysis of IRL-extracted reward functions further indicates benefits for more robust reward functions via tighter integration of supervised and preference-based LLM post-training.