온폴리시 지식 증류 재고: 경험적 실패 양상과 단순한 해결책

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

March 26, 2026

저자: Yuqian Fu, Haohuan Huang, Kaiwen Jiang, Yuanheng Zhu, Dongbin Zhao

cs.AI

초록

온-폴리시 디스틸레이션(OPD)은 고정된 교사 추적이 아닌 학생 생성 롤아웃에 대한 교사 피드백을 평가하기 때문에 대규모 언어 모델(LLM) 사후 학습에 매력적입니다. 그러나 장기간 설정에서 일반적인 샘플링된 토큰 변형은 취약합니다: 이는 분포 매칭을 단일 토큰 신호로 축소하며, 롤아웃이 교사가 일반적으로 방문하는 접두사에서 점점 더 멀어질수록 신뢰성이 낮아집니다. 우리는 추정기와 구현 측면에서 OPD를 재검토합니다. 이론적으로 토큰 수준 OPD는 시퀀스 수준 역 KL에 비해 편향되지만 훨씬 더 엄격한 최악의 경우 분산 한계를 가집니다. 우리의 토이 연구는 동일한 트레이드오프를 경험적으로 보여주며, 더 강한 미래 보상 결합은 더 높은 그래디언트 분산과 덜 안정적인 학습을 초래합니다. 경험적으로 우리는 샘플링된 토큰 OPD의 세 가지 실패 모드를 확인했습니다: 불균형한 단일 토큰 신호, 학생 생성 접두사에 대한 신뢰할 수 없는 교사 지도, 그리고 토크나이저 또는 특수 토큰 불일치로 인한 왜곡. 우리는 이러한 문제를 교사 상위 K 지역 지원 매칭으로 해결하며, 이는 상위 p 롤아웃 샘플링과 특수 토큰 마스킹을 사용한 절단 역 KL로 구현됩니다. 단일 작업 수학 추론 및 다중 작업 에이전트+수학 학습 전반에 걸쳐, 이 목적 함수는 샘플링된 토큰 OPD보다 더 안정적인 최적화와 더 나은 다운스트림 성능을 제공합니다.

English

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.