TIP: 온-폴리시 증류에서의 토큰 중요도

TIP: Token Importance in On-Policy Distillation

온-폴리시 지식 증류(OPD)는 교사 모델의 토큰 단위 감독 하에 학생 모델이 자체적으로 생성한 데이터에 대해 학습하는 방식입니다. 모든 토큰 위치가 동일한 중요도를 갖는 것은 아니지만, 기존의 토큰 중요도에 대한 관점은 아직 불완전합니다. 본 연구에서는 다음과 같은 직접적인 질문을 던집니다: OPD에서 가장 유용한 학습 신호를 전달하는 토큰은 무엇인가? 그 답은 정보적인 토큰이 두 가지 영역에서 나타난다는 것입니다. 첫째는 학생 모델의 엔트로피가 높은 위치이고, 둘째는 학생 모델의 엔트로피가 낮으면서 교사-학생 모델 간의 불일치가 높은 위치, 즉 학생 모델이 과신하고 잘못된 위치입니다. 실험적으로, 학생 모델의 엔트로피는 강력한 지표로 작용합니다. 엔트로피 기반 샘플링을 통해 전체 토큰의 50%를 유지하면 모든 토큰을 사용하는 학습 방식과 동일하거나 그 이상의 성능을 달성하면서, 최대 47%까지 메모리 사용량을 줄일 수 있습니다. 하지만 엔트로피만으로는 중요한 정보를 놓칠 수 있습니다. 엔트로피가 낮고 불일치가 높은 토큰을 분리하여 학습했을 때, 전체 토큰의 10% 미만을 사용해도 전체 토큰을 사용하는 기준 성능에 거의 근접하는 결과를 얻었습니다. 이는 과신된 토큰이 엔트로피 기반 규칙으로는 거의 감지할 수 없지만, 강력한 교정 신호를 담고 있음을 보여줍니다. 본 연구에서는 이러한 결과를 TIP (Token Importance in on-Policy distillation)이라는 두 축(학생 모델의 엔트로피와 교사-학생 모델 간의 불일치)을 기반으로 하는 분류 체계를 통해 정리하고, 엔트로피가 유용하지만 구조적으로 불완전한 이유에 대한 이론적 설명을 제공합니다. 이러한 관점은 불확실성과 불일치를 결합한 토큰 선택 규칙을 개발하는 데 영감을 줍니다. 본 연구에서는 Qwen3, Llama, Qwen2.5 모델을 사용하여 MATH-500 및 AIME 2024/2025 데이터셋, 그리고 장기적인 에이전트 기반 계획을 위한 DeepPlanning 벤치마크에서 세 가지 교사-학생 모델 조합을 사용하여 이 가설을 검증했습니다. 실험은 메모리 효율적인 대규모 모델 증류를 지원하는 https://github.com/HJSang/OPSD_OnPolicyDistillation 레포지토리의 확장 기능을 통해 구현되었으며, 이는 제한된 GPU 환경에서 유용하게 활용될 수 있습니다.

On-policy knowledge distillation (OPD) trains a student on its own rollouts under token-level supervision from a teacher. Not all token positions matter equally, but existing views of token importance are incomplete. We ask a direct question: which tokens carry the most useful learning signal in OPD? Our answer is that informative tokens come from two regions: positions with high student entropy, and positions with low student entropy plus high teacher--student divergence, where the student is overconfident and wrong. Empirically, student entropy is a strong first-order proxy: retaining $50\%$ of tokens with entropy-based sampling matches or exceeds all-token training while reducing peak memory by up to $47\%$. But entropy alone misses a second important region. When we isolate low-entropy, high-divergence tokens, training on fewer than $10\%$ of all tokens nearly matches full-token baselines, showing that overconfident tokens carry dense corrective signal despite being nearly invisible to entropy-only rules. We organize these findings with TIP (Token Importance in on-Policy distillation), a two-axis taxonomy over student entropy and teacher--student divergence, and give a theoretical explanation for why entropy is useful yet structurally incomplete. This view motivates type-aware token selection rules that combine uncertainty and disagreement. We validate this picture across three teacher--student pairs spanning Qwen3, Llama, and Qwen2.5 on MATH-500 and AIME 2024/2025, and on the DeepPlanning benchmark for long-horizon agentic planning, where Q3-only training on $<$$20\%$ of tokens surpasses full-token OPD. Our experiments are implemented by extending the OPD repository https://github.com/HJSang/OPSD_OnPolicyDistillation, which supports memory-efficient distillation of larger models under limited GPU budgets.