신뢰 영역 온-정책 증류

초록

온-정책 증류(On-Policy Distillation, OPD)는 대규모 언어 모델(LLM)의 효율적인 사후 학습을 위한 기본 기법으로, 에이전트 학습, 다중 작업 강화, 모델 압축 등에 폭넓게 적용된다. 그러나 OPD 학습은 교사와 학생 분포가 크게 다를 때 불안정해지는데, 학생이 생성한 토큰에 대한 교사의 감독이 신뢰할 수 없는 정책 기울기를 초래하고 최적화 실패로 이어질 수 있다. 본 연구는 신용 할당 전략을 통해 신뢰할 수 있는 온-정책 토큰 수준 감독 문제를 해결하며, 신뢰 영역 온-정책 증류(Trust Region On-Policy Distillation, TrOPD)를 제안한다. TrOPD는 다음과 같은 특징을 가진다: 1) 신뢰 영역 온-정책 학습: TrOPD는 교사가 신뢰할 수 있는 감독을 제공하는 영역에서만 OPD를 수행하여, 분포 불일치 하에서 K1 역방향-KL 추정기의 최적화 어려움을 완화한다. 2) 이상치 추정: 이상치 영역의 경우, 기울기 클리핑, 마스킹, 순방향-KL 추정을 탐색하여 신뢰할 수 없는 감독의 부정적 영향을 줄인다. 3) 오프-정책 가이던스: 학생은 교사 접두사로부터 생성을 계속하고 순방향 KL을 사용하여 오프-정책 가이던스를 모방함으로써, 신뢰할 수 있는 영역으로의 온-정책 탐색을 장려한다. 실험 결과, TrOPD는 수학적 추론, 코드 생성 및 일반 도메인 벤치마크에서 OPD, EOPD, REOPOLD 등 최신 OPD 기준선을 일관되게 능가함을 보여준다.

English

On-Policy Distillation (OPD) is a fundamental technique for efficient post-training of large language models (LLMs), with broad applications in agent learning, multi-task enhancement, and model compression. However, OPD training becomes unstable when the teacher and student distributions differ substantially, as teacher supervision on student-generated tokens may yield unreliable policy gradients and even cause optimization failure. This work addresses reliable on-policy token-level supervision through credit assignment strategies, and proposes Trust Region On-Policy Distillation, TrOPD. It features the following characteristics: 1) Trust-Region On-Policy Learning: TrOPD performs OPD only in regions where the teacher provides reliable supervision, mitigating the optimization difficulty of the K1 reverse-KL estimator under distribution mismatch. 2) Outlier Estimation: For outlier regions, we explore gradient clipping, masking, and forward-KL estimation to reduce the adverse effects of unreliable supervision. 3) Off-Policy Guidance: The student continues generation from teacher prefixes and uses forward KL to imitate off-policy guidance, encouraging on-policy exploration toward reliable regions. Experiments show that TrOPD consistently outperforms SoTA OPD baselines, including OPD, EOPD, and REOPOLD, across mathematical reasoning, code generation, and general-domain benchmarks.