Ignore the KL Penalty! Boosting Exploration on Critical Tokens to Enhance RL Fine-Tuning

📄 arXiv: 2502.06533v1 📥 PDF

作者: Jean Vassoyan, Nathanaël Beau, Roman Plaud

分类: cs.CL, cs.LG

发布日期: 2025-02-10

备注: 11 pages, 6 figures, 5 tables. Accepted for publication in the Findings of the North American Chapter of the Association for Computational Linguistics (NAACL) 2025

💡 一句话要点

通过强化关键token探索，忽略KL惩罚提升RL微调效果

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

关键词: 强化学习 语言模型微调 KL惩罚 探索策略 关键Token 策略梯度 奖励函数

📋 核心要点

  1. 现有RL微调LLM方法难以平衡探索新方案和保持模型基本能力，KL惩罚是常用控制手段。
  2. 论文提出修改KL惩罚，鼓励对影响最终结果的关键token进行探索，提升微调效率。
  3. 实验表明，该方法在算术任务上有效，通过关注关键token显著提升了RL微调的性能。

📝 摘要（中文）

在大型语言模型(LLM)的发展中，实现长期目标是一项关键挑战。为了解决这个问题，可以使用强化学习(RL)对预训练的LLM进行微调，以探索优化给定目标的解决方案。然而，LLM的探索是困难的，因为必须在新方案发现和保持接近预训练模型之间取得平衡，以避免降低基本能力。这通常通过Kullback-Leibler (KL)惩罚来控制。本文研究了一个小型语言模型在简单算术任务上的探索动态，展示了不同程度的预训练如何影响探索，并强调了对最终结果有显著影响的“关键token”的重要性。因此，我们对KL惩罚进行了一个简单的修改，该修改有利于对关键token的探索，从而提高了RL微调阶段的效率。

🔬 方法详解

问题定义：现有基于强化学习的语言模型微调方法，为了防止模型偏离预训练分布过远，通常会引入KL散度惩罚项。然而，过强的KL惩罚会限制模型的探索能力，导致无法找到更优的策略。因此，如何在探索和稳定性之间取得平衡是一个关键问题。现有方法通常采用固定的KL惩罚系数，无法根据不同token的重要性进行自适应调整。

核心思路：论文的核心思路是，并非所有token对最终结果的贡献都是相同的，某些“关键token”对结果的影响更大。因此，应该鼓励模型在这些关键token上进行更多的探索，而对非关键token则可以适当增加KL惩罚。通过这种方式，可以在保证模型稳定性的同时，提升探索效率。

技术框架：整体框架仍然是标准的强化学习微调流程，包括：1) 使用语言模型生成文本；2) 根据环境反馈计算奖励；3) 使用策略梯度方法更新模型参数。关键的修改在于KL惩罚项的计算方式。具体来说，对于每个token，计算其对最终奖励的敏感度，作为该token的关键性指标。然后，根据该指标调整KL惩罚系数，关键性越高的token，KL惩罚越小。

关键创新：最重要的创新点在于提出了“关键token”的概念，并将其与KL惩罚相结合。与现有方法不同，该方法不再对所有token施加相同的KL惩罚，而是根据token的重要性进行自适应调整。这种方法能够更有效地利用有限的探索资源，提升RL微调的效率。

关键设计：关键的设计包括：1) 如何定义和计算token的关键性。论文采用的方法是计算每个token的输出概率对最终奖励的梯度，梯度越大，表示该token对结果的影响越大。2) 如何根据token的关键性调整KL惩罚系数。论文采用了一种简单的线性缩放方法，即KL惩罚系数与token关键性成反比。3) 实验中使用的语言模型是一个小型Transformer模型，任务是解决简单的算术问题。奖励函数根据模型输出的答案是否正确进行设计。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文在简单的算术任务上验证了所提出方法的有效性。实验结果表明，通过关注关键token的探索，可以显著提升RL微调的性能。具体来说，相比于使用固定KL惩罚的方法，该方法能够更快地找到更优的策略，并获得更高的奖励。

🎯 应用场景

该研究成果可应用于各种需要通过强化学习微调大型语言模型的场景，例如对话生成、文本摘要、机器翻译等。通过关注关键token的探索，可以提升模型在特定任务上的性能，并降低微调所需的计算资源。未来，该方法可以进一步推广到更复杂的任务和更大的模型上。

📄 摘要（原文）

The ability to achieve long-term goals is a key challenge in the current development of large language models (LLMs). To address this, pre-trained LLMs can be fine-tuned with reinforcement learning (RL) to explore solutions that optimize a given goal. However, exploration with LLMs is difficult, as a balance has to be struck between discovering new solutions and staying close enough to the pre-trained model, so as not to degrade basic capabilities. This is typically controlled with a Kullback-Leibler (KL) penalty. In this paper, we investigate the exploration dynamics of a small language model on a simple arithmetic task. We show how varying degrees of pre-training influence exploration and demonstrate the importance of "critical tokens" which have a dramatic impact on the final outcome. Consequently, we introduce a simple modification to the KL penalty that favors exploration on critical tokens, increasing the efficiency of the RL fine-tuning stage.