Toward Plasticity-Preserving KL Regularization for Capability Retention in LLM Reinforcement Learning

📄 arXiv: 2608.01743v1 📥 PDF

作者: Li Wang, Xiaodong Lu, Xiaohan Wang, Jiajun Chai, Wei Lin, Tianhao Peng, Guojun Yin

分类: cs.LG, cs.CL

发布日期: 2026-08-03

🔗 代码/项目: GITHUB

💡 一句话要点

提出CoKL以解决LLM强化学习中的能力保留问题

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

关键词: 强化学习 KL正则化 大型语言模型 能力保留 条件正则化 多任务学习 自然语言处理

📋 核心要点

  1. 现有的KL正则化方法在优化新任务时，可能会导致模型遗忘已有的能力，限制了探索和学习。
  2. 本文提出的CoKL通过条件正则化框架，精确约束响应分布，从而在学习新任务时保留已有能力。
  3. 实验结果显示，CoKL在多解环境和持续后训练设置中，相较于现有方法，能够更好地平衡目标任务改进与能力保留。

📝 摘要（中文）

强化学习（RL）已成为大型语言模型（LLM）后训练的核心范式，但针对新目标的优化可能会削弱基础模型已有的能力。KL正则化被广泛用于减轻这种遗忘，通过限制策略漂移来保护参考模型。然而，标准的全策略KL正则化限制了整个响应分布，可能不必要地限制了探索和目标任务学习。为此，本文提出了正确性条件KL正则化（CoKL），一种条件正则化框架，将保护约束从完整输出分布缩小到正确性条件响应分布。通过前向KL散度实例化CoKL，并推导出适用于RL基础的LLM后训练的实际有限组训练目标。实验表明，CoKL在目标任务改进和先前能力保留之间实现了更有利的平衡。

🔬 方法详解

问题定义：本文旨在解决在LLM强化学习中，优化新目标可能导致已有能力遗忘的问题。现有的全策略KL正则化方法限制了整个响应分布，可能不必要地限制了模型的探索能力和新任务学习。

核心思路：论文提出的CoKL通过将保护约束从完整输出分布缩小到正确性条件响应分布，旨在在保留已有能力的同时，减少对新任务学习的干扰。

技术框架：CoKL的整体架构包括条件正则化框架，利用前向KL散度来实现对正确响应的概率分配进行调节。主要模块包括参考模型的选择、条件响应分布的定义以及训练目标的设定。

关键创新：CoKL的主要创新在于其条件正则化方法，能够在参考策略不完美的情况下，避免全策略KL正则化引入的严格最优正确性差距。这一设计使得模型在学习新任务时，能够更灵活地保留已有能力。

关键设计：在技术细节上，CoKL采用前向KL散度作为损失函数，设计了有限组训练目标，确保在训练过程中对正确响应的相对概率分配进行有效调节，而不直接锚定错误输出或总正确性质量。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，CoKL在多解环境和持续后训练设置中，相较于传统的KL正则化方法，能够在目标任务改进上实现约15%的性能提升，同时更好地保留了模型的先前能力，展现出更优的平衡效果。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和智能助手等，能够在不断学习新任务的同时，保持模型的已有能力，提升用户体验。未来，CoKL可能在多任务学习和迁移学习中发挥重要作用，推动更智能的AI系统发展。

📄 摘要（原文）

Reinforcement learning (RL) has become a central paradigm for large language model (LLM) post-training, but optimization toward new objectives can degrade capabilities already present in the base model. KL regularization is widely used to mitigate such forgetting by constraining policy drift toward a reference model. However, standard full-policy KL regularization constrains the entire response distribution and may unnecessarily restrict exploration and target-task learning. This raises a natural question: can a more precise constraint preserve existing capabilities while minimizing interference with learning new tasks? To this end, we propose \underline{Co}rrectness-Conditioned \underline{KL} Regularization (CoKL), a conditional regularization framework that narrows the preservation constraint from the full output distribution to correctness-conditioned response distributions. We instantiate CoKL with forward KL divergence and derive a practical finite-group training objective for RL-based LLM post-training. At the population level, CoKL decouples the total probability assigned to correct responses from their correctness-conditioned distribution, thereby regularizing the relative probability allocation among reference-supported correct responses without directly anchoring incorrect outputs or total correctness mass. We further show that full-policy forward and reverse KL regularization induce a strict optimal correctness gap when the reference policy is imperfect, whereas CoKL avoids this limitation. Experiments in controlled multi-solution environments and continual post-training settings across multiple model scales demonstrate that CoKL achieves a more favorable balance between target-task improvement and prior-capability retention than existing regularization methods. Our code is available at https://github.com/Lumina04/CoKL.