LLM 추론을 위한 강화 학습 인식 지식 증류

초록

강화 학습(RL) 사후 훈련은 최근 사고 연쇄 추론 대규모 언어 모델(LLM)에서 큰 성과를 거두었지만, 이러한 모델의 높은 추론 비용은 더 작은 학생 모델로의 지식 증류를 필요로 합니다. 기존 지식 증류(KD) 방법 대부분은 감독 미세 조정(SFT)을 위해 설계되어 고정된 교사 추적 또는 교사-학생 KL(Kullback-Leibler) 발산 기반 정규화에 의존합니다. 이러한 접근법을 RL과 결합할 경우, 분포 불일치와 목적 간섭 문제가 자주 발생합니다: 교사의 감독이 학생의 진화하는 롤아웃 분포와 일치하지 않을 수 있으며, KL 정규화기가 보상 최대화와 경쟁하고 신중한 손실 균형 조정을 필요로 할 수 있습니다. 이러한 문제를 해결하기 위해, 우리는 RL 중 선택적 모방을 수행하는 RL 인식 증류(RLAD)를 제안합니다. 이는 현재 정책 업데이트를 개선할 때만 학생을 교사 방향으로 유도합니다. 우리의 핵심 구성 요소인 신뢰 영역 비율 증류(TRRD)는 교사-학생 KL 정규화기를 PPO/GRPO 스타일의 가능도 비율 목적 함수로 대체하며, 이는 교사-이전 정책 혼합에 기반을 둡니다. 이를 통해 학생 롤아웃에 대한 이점 인식 및 신뢰 영역 경계 증류를 제공하고 탐색, 활용, 모방을 자연스럽게 균형 잡습니다. 다양한 논리 추론 및 수학 벤치마크에서 RLAD는 오프라인 증류, 표준 GRPO, KL 기반 온-정책 교사-학생 지식 증류를 일관되게 능가합니다.

English

Reinforcement learning (RL) post-training has recently driven major gains in long chain-of-thought reasoning large language models (LLMs), but the high inference cost of such models motivates distillation into smaller students. Most existing knowledge distillation (KD) methods are designed for supervised fine-tuning (SFT), relying on fixed teacher traces or teacher-student Kullback-Leibler (KL) divergence-based regularization. When combined with RL, these approaches often suffer from distribution mismatch and objective interference: teacher supervision may not align with the student's evolving rollout distribution, and the KL regularizer can compete with reward maximization and require careful loss balancing. To address these issues, we propose RL-aware distillation (RLAD), which performs selective imitation during RL -- guiding the student toward the teacher only when it improves the current policy update. Our core component, Trust Region Ratio Distillation (TRRD), replaces the teacher-student KL regularizer with a PPO/GRPO-style likelihood-ratio objective anchored to a teacher--old-policy mixture, yielding advantage-aware, trust-region-bounded distillation on student rollouts and naturally balancing exploration, exploitation, and imitation. Across diverse logic reasoning and math benchmarks, RLAD consistently outperforms offline distillation, standard GRPO, and KL-based on-policy teacher-student knowledge distillation.

LLM 추론을 위한 강화 학습 인식 지식 증류

Reinforcement-aware Knowledge Distillation for LLM Reasoning

초록

Support