온라인 정책 증류 재검토: 경험적 실패 요인 및 간단한 해결책

Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes

온라인 정책 증류(OPD)는 교사의 고정된 행동 패턴 대신 학생이 생성한 결과에 대한 교사의 피드백을 평가하므로, 대규모 언어 모델(LLM)의 사후 훈련에 매력적인 방법입니다. 그러나 장기적인 환경 설정에서, 일반적으로 사용되는 샘플링된 토큰 방식은 불안정합니다. 이 방식은 분포 정렬을 단일 토큰 신호로 제한하며, 학생의 결과가 교사가 자주 방문하는 위치에서 멀어질수록 신뢰성이 떨어집니다. 본 연구에서는 추정 및 구현 측면에서 OPD를 재검토합니다. 이론적으로, 토큰 수준의 OPD는 시퀀스 수준의 역방향 KL 발산과 비교하여 편향이 있지만, 훨씬 더 엄격한 최악의 분산 경계를 가집니다. 우리의 간단한 연구 결과는 이론적 예측과 일치하며, 미래 보상의 강한 연관성은 더 높은 기울기 분산과 불안정한 학습을 초래합니다. 경험적으로, 샘플링된 토큰 OPD의 세 가지 실패 요인을 확인했습니다. 즉, 불균형한 단일 토큰 신호, 학생이 생성한 접두사에 대한 신뢰할 수 없는 교사의 안내, 그리고 토크나이저 또는 특수 토큰 불일치로 인한 왜곡입니다. 우리는 교사의 상위 K개 로컬 지원 매칭을 사용하여 이러한 문제를 해결하며, 이는 상위-p 로우트 샘플링 및 특수 토큰 마스킹을 사용한 잘린 역방향 KL로 구현됩니다. 단일 작업 수학 추론 및 다중 작업 에이전트 기반 수학 훈련에서, 이 목적 함수는 샘플링된 토큰 OPD보다 더 안정적인 최적화를 제공하며 더 나은 성능을 보입니다.

On-policy distillation (OPD) is appealing for large language model (LLM) post-training because it evaluates teacher feedback on student-generated rollouts rather than fixed teacher traces. In long-horizon settings, however, the common sampled-token variant is fragile: it reduces distribution matching to a one-token signal and becomes increasingly unreliable as rollouts drift away from prefixes the teacher commonly visits. We revisit OPD from the estimator and implementation sides. Theoretically, token-level OPD is biased relative to sequence-level reverse-KL, but it has a much tighter worst-case variance bound; our toy study shows the same tradeoff empirically, with stronger future-reward coupling producing higher gradient variance and less stable learning. Empirically, we identify three failure modes of sampled-token OPD: an imbalanced one-token signal, unreliable teacher guidance on student-generated prefixes, and distortions caused by tokenizer or special-token mismatch. We address these issues with teacher top-K local support matching, implemented as truncated reverse-KL with top-p rollout sampling and special-token masking. Across single-task math reasoning and multi-task agentic-plus-math training, this objective yields more stable optimization and better downstream performance than sampled-token OPD.