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

• URL: http://arxiv.org/abs/2608.01743v1
• Date: Mon, 03 Aug 2026 06:10:33 GMT
• Title: Toward Plasticity-Preserving KL Regularization for Capability Retention in LLM Reinforcement Learning
• Abstract summary: We propose a conditional regularization framework that narrows the preservation constraint from the full output distribution to correctness-conditioned response distributions.<n>CoKL achieves a more favorable balance between target-task improvement and prior-capability retention than existing regularization methods.
• Score: 43.981833556858135
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Beyond the Stability-Exploration Dilemma: Environmental Regularization for LLM Policy Optimization [27.121081058414376]
  Policy-KL constrains response behavior and consumes the action-side exploration budget.<n>We argue for an alternative that breaks the dilemma by moving regularization to the input side.<n>ERPO replaces the standard Policy-KL regularizer while achieving effective control over query distribution drift.
  arXiv  Detail & Related papers  (2026-08-24T14:30:59Z)
• Post-Training at the Edge of Detectability: A Game-Theoretic Approach to Fine-Tuning [58.088509347353465]
  Reinforcement learning (RL) fine-tuning is widely used in language model training.<n>We propose a game-theoretic framework that gives this trade-off an explicit statistical interpretation.
  arXiv  Detail & Related papers  (2026-07-29T00:19:09Z)
• RL Forgets! Towards Continual Policy Optimization [57.4336338996653]
  Continual post-training is becoming a central paradigm for adapting vision-language models to evolving tasks.<n>Recent work has increasingly favored reinforcement learning over supervised fine-tuning, driven by the belief that reinforcement learning is inherently less prone to forgetting.<n>We propose Continual Policy Optimization (CPO), a replay-free framework grounded in a prior-task behavioral KL objective.
  arXiv  Detail & Related papers  (2026-07-05T15:46:11Z)
• Online KL-Regularized Reinforcement Learning with Function Approximation under Misspecification [70.9534986000242]
  We study KL-regularized contextual bandits and episodic reinforcement learning (RL) under general function approximation with model misspecification.<n>Existing guarantees rely on realizability and therefore do not extend to misspecified models, where classical regret bounds may fail.<n>This work introduces KL misspecification formulations for contextual bandits and episodic RL and analyzes regression-based algorithms with Gibbs policy updates.
  arXiv  Detail & Related papers  (2026-06-04T11:54:23Z)
• Improving Search Agent with One Line of Code [68.58667107354253]
  Tool-based Agentic Reinforcement Learning (TARL) has emerged as a promising paradigm for training search agents to interact with external tools.<n>We propose textbfSearch textbfAgent textbfPolicy textbfOptimization (textbfSAPO), which stabilizes training via a conditional token-level KL constraint.
  arXiv  Detail & Related papers  (2026-03-10T04:07:39Z)
• A Theoretical Lens for RL-Tuned Language Models via Energy-Based Models [9.324642081509756]
  Large language models (LLMs) trained via KL-regularized reinforcement learning demonstrate strong instruction following, self-correction, and reasoning abilities.<n>We exploit the closed-form energy-based model (EBM) structure of the optimal KL-regularized policy to provide a unified variational analysis of LLMs.
  arXiv  Detail & Related papers  (2025-12-21T13:28:58Z)
• Sharp Analysis for KL-Regularized Contextual Bandits and RLHF [52.519416266840814]
  Reverse-Kullback-Leibler (KL) regularization has emerged to be a predominant technique used to enhance policy optimization in reinforcement learning.<n>We show that a simple two-stage mixed sampling strategy can achieve a sample complexity with only an additive dependence on the coverage coefficient.<n>Our results provide a comprehensive understanding of the roles of KL-regularization and data coverage in RLHF, shedding light on the design of more efficient RLHF algorithms.
  arXiv  Detail & Related papers  (2024-11-07T11:22:46Z)
• Rethinking Kullback-Leibler Divergence in Knowledge Distillation for Large Language Models [18.870276152694245]
  Kullback-Leiber divergence has been widely used in Knowledge Distillation (KD) to compress Large Language Models (LLMs)<n>Contrary to prior assertions, reverse Kullback-Leibler (RKL) divergence is mode-seeking and thus preferable over the mean-seeking forward Kullback-Leibler (FKL) divergence.<n>We propose a simple yet effective Adaptive Kullback-Leiber (AKL) divergence method, which adaptively allocates weights to combine FKL and RKL.
  arXiv  Detail & Related papers  (2024-04-03T11:40:17Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.