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

2026-08-03

AI Summary

This work addresses the challenge of catastrophic forgetting in large language models during reinforcement learning fine-tuning, where conventional KL regularization methods hinder new task acquisition by constraining the entire output distribution. To mitigate this limitation, the authors propose Correctness-conditioned KL Regularization (CoKL), which, for the first time, restricts KL divergence only to the conditional distribution over correct responses. This approach decouples the total probability mass assigned to correct answers from their internal token-level allocation, thereby preserving pre-existing capabilities while minimizing interference with exploration on new tasks. Leveraging forward KL divergence, a bounded grouping training objective, and multi-solution environment settings, CoKL consistently outperforms existing methods across diverse model scales and continual post-training scenarios, achieving superior target-task performance alongside more effective retention of original abilities.

Abstract

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.

Problem

• capability retention
• KL regularization
• reinforcement learning
• large language models
• catastrophic forgetting

Innovation

• Correctness-Conditioned KL Regularization
• Capability Retention
• Large Language Model Reinforcement Learning
• Policy Regularization
• Forward KL Divergence