[论文解读] On Reinforcement Learning and Distribution Matching for Fine-Tuning Language Models with no Catastrophic Forgetting

Tomasz Korbak, Hady Elsahar+2 more|arXiv (Cornell University)|Jun 1, 2022

Topic Modeling被引用 6

一句话总结

该论文通过统一奖励最大化（RM）与分布匹配（DM），在微调语言模型时避免灾难性遗忘。通过将RM中的KL控制重新解释为一种DM形式，并在DM中引入基线，作者提升了可控文本生成任务（如情感、主题和性别控制）中的样本效率、稳定性和约束满足度。

ABSTRACT

The availability of large pre-trained models is changing the landscape of Machine Learning research and practice, moving from a training-from-scratch to a fine-tuning paradigm. While in some applications the goal is to "nudge" the pre-trained distribution towards preferred outputs, in others it is to steer it towards a different distribution over the sample space. Two main paradigms have emerged to tackle this challenge: Reward Maximization (RM) and, more recently, Distribution Matching (DM). RM applies standard Reinforcement Learning (RL) techniques, such as Policy Gradients, to gradually increase the reward signal. DM prescribes to first make explicit the target distribution that the model is fine-tuned to approximate. Here we explore the theoretical connections between the two paradigms, and show that methods such as KL-control developed for RM can also be construed as belonging to DM. We further observe that while DM differs from RM, it can suffer from similar training difficulties, such as high gradient variance. We leverage connections between the two paradigms to import the concept of baseline into DM methods. We empirically validate the benefits of adding a baseline on an array of controllable language generation tasks such as constraining topic, sentiment, and gender distributions in texts sampled from a language model. We observe superior performance in terms of constraint satisfaction, stability and sample efficiency.

研究动机与目标

• 解决语言模型微调中的灾难性遗忘问题，即模型在适应新偏好时会丧失原有能力。
• 统一两种主流范式——奖励最大化（RM）与分布匹配（DM）——以结合其各自优势。
• 通过借鉴强化学习中的基线概念，将基线引入DM方法，降低DM方法的高梯度方差，提升训练稳定性。
• 实现传统RM因序列级奖励不可行而难以实现的分布约束（如性别或主题平衡）。
• 在多个可控生成任务中，展示出在样本效率、约束遵循度与训练稳定性方面的优越性能。

提出的方法

• 将RM中的KL控制重新表述为一种分布匹配（DM）形式，理论证明惩罚性奖励最大化与最小化前向KL散度等价。
• 在分布匹配算法中引入基线，以降低梯度方差，灵感来源于强化学习中的策略梯度方法。
• 使用能量模型（EBMs）定义目标分布，以编码所需偏好，如情感、主题或性别平衡。
• 采用分布策略梯度（DPG）最小化目标分布与模型输出分布之间的前向KL散度。
• 通过指示函数与目标均值，结合多种特征（如50%提及女性且包含体育术语）实现混合约束。
• 使用基线减法的策略梯度更新组合方式训练模型，以在基于DM的优化中稳定学习。

实验结果

研究问题

• RQ1 RM中的KL控制能否被解释为一种分布匹配形式？
• RQ2 在语言模型微调背景下，RM与DM之间在理论与实践上存在何种联系？
• RQ3 在DM方法中引入基线是否能提升训练稳定性和样本效率？
• RQ4 基于DM的方法能否在文本生成中有效实施复杂分布约束（如性别或主题平衡）？
• RQ5 所提出的基线增强型DM方法相较于标准RM与DM基线，在约束满足度与稳定性方面表现如何？

主要发现

• RM中的KL控制在理论上等价于最小化预训练模型与目标分布之间的前向KL散度，因此可归入DM框架。
• 所提出的DM基线显著降低了梯度方差，相比标准DM方法，训练动态更加稳定。
• 该方法在约束满足方面表现优越——例如，能有效维持生成文本中50%的女性提及比例，同时保持流畅性与多样性。
• 在涉及混合约束的任务（如性别平衡+主题包含）中，基线增强型DM模型在样本效率与稳定性方面均优于标准DM与RM基线。
• 实证结果表明，该方法比标准强化学习微调更能保持原始模型的分布特性，避免灾难性遗忘。
• 该方法在多种可控生成任务中具有泛化能力，包括情感控制、主题控制与性别平衡生成，且在所有指标上均保持一致改进。