ヒント:オンポリシー蒸留におけるトークン重要度
TIP: Token Importance in On-Policy Distillation
April 15, 2026
著者: Yuanda Xu, Hejian Sang, Zhengze Zhou, Ran He, Zhipeng Wang, Alborz Geramifard
cs.AI
要旨
オン方針知識蒸留(OPD)は、トークンレベルの教師監督下で、学生モデル自身のロールアウトを用いて学習を行う。全てのトークン位置が同等に重要ではなく、既存のトークン重要性の見解は不完全である。我々は直接的な問いを立てる:OPDにおいて、どのトークンが最も有用な学習信号を持つのか?その答えは、有益なトークンが2つの領域から生じるというものである:学生モデルのエントロピーが高い位置と、学生モデルのエントロピーが低くかつ教師-学生間のダイバージェンスが高い(学生モデルが過信して誤っている)位置である。
実証的に、学生エントロピーは強力な一次代理指標である:エントロピーに基づくサンプリングで上位50%のトークンを残すことで、全トークン学習と同等かそれを上回る性能を達成しつつ、ピークメモリ使用量を最大47%削減する。しかし、エントロピーだけでは第二の重要な領域を見逃す。低エントロピー・高ダイバージェンストークンを単離すると、全トークンの10%未満での学習がほぼ全トークンベースラインに匹敵し、過信トークンがエントロピー単独の規則ではほぼ検出不能であるにも関わらず、高密度の修正信号を運ぶことを示す。
我々はこれらの知見を、学生エントロピーと教師-学生ダイバージェンスの2軸からなる分類法「TIP(オン方針蒸留におけるトークン重要性)」で体系化し、エントロピーが有用でありながら構造的に不完全である理論的説明を与える。この見解は、不確実性と不一致を組み合わせたタイプ認識型トークン選択規則を動機づける。我々はこの図式を、Qwen3、Llama、Qwen2.5に跨る3組の教師-学生ペアで、MATH-500およびAIME 2024/2025、さらに長期的なエージェンシック計画のためのDeepPlanningベンチマークにおいて検証する。後者では、全トークンの20%未満のQ3トークン(低エントロピー・高ダイバージェンス)のみでの学習が、全トークンOPDを凌駕する。本実験は、限られたGPU予算下での大規模モデルのメモリ効率的な蒸留をサポートするOPDリポジトリ(https://github.com/HJSang/OPSD_OnPolicyDistillation)を拡張して実装された。
English
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.