LLM 강화 학습에서 능력 유지 및 보존을 위한 가소성을 고려한 KL 정규화

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

강화 학습(RL)은 대규모 언어 모델(LLM)의 사후 훈련에 중요한 패러다임이 되었지만, 새로운 목표를 향한 최적화는 기본 모델에 이미 존재하는 능력을 저하시킬 수 있습니다. KL 정규화는 종종 참조 모델을 기준으로 정책의 변화를 제한하여 이러한 망각 현상을 완화하는 데 사용됩니다. 그러나 표준적인 전체 정책 KL 정규화는 전체 응답 분포를 제약하며, 이는 불필요하게 탐색과 목표 작업 학습을 제한할 수 있습니다. 따라서 다음과 같은 질문이 제기됩니다. 보다 정확한 제약을 통해 기존 능력을 유지하면서 새로운 작업을 학습하는 데 미치는 간섭을 최소화할 수 있을까요? 이를 위해 우리는 '정확성 조건부 KL 정규화(CoKL)'라는 조건부 정규화 프레임워크를 제안합니다. CoKL은 보존 제약 조건을 전체 출력 분포에서 정확성-조건부 응답 분포로 좁힙니다. 우리는 순방향 KL 발산을 사용하여 CoKL을 구현하고, RL 기반 LLM 사후 훈련을 위한 실용적인 유한 그룹 학습 목표를 도출했습니다. 모집단 수준에서 CoKL은 올바른 응답에 할당된 전체 확률을 해당 정확성-조건부 분포와 분리하여, 참조 모델이 지원하는 올바른 응답 간의 상대적 확률 할당을 정규화합니다. 동시에 잘못된 출력을 직접 고정하거나 전체 정확도 척도를 고정하지 않습니다. 제어된 다중 해 환경 및 다양한 모델 크기에 대한 연속적인 사후 훈련 설정을 통해 수행한 실험 결과, CoKL은 기존 정규화 방법보다 목표 작업 개선과 이전 능력 유지 간의 더 나은 균형을 제공합니다. 저희 코드는 https://github.com/Lumina04/CoKL 에서 확인할 수 있습니다.

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.