TIP: オンポリシー・ディスティレーションにおけるトークン重要度

要旨: オンポリシー知識蒸留（OPD）は、教師からのトークン単位の監督のもとで、生徒自身のロールアウトにより生徒を訓練します。すべてのトークン位置が等しく重要というわけではありませんが、トークン重要度に関する既存の見解は不十分です。私たちは直接的な問いを立てます。OPDにおいて最も有用な学習シグナルを運ぶのはどのトークンか？ 結論として、有益なトークンは2つの領域から来ることが分かります。すなわち、生徒のエントロピーが高い位置、そして、生徒のエントロピーが低いにもかかわらず教師—生徒間のダイバージェンスが高い位置です。後者では、生徒は過信しており誤っています。
実験的には、生徒のエントロピーは強力な一次の代理指標です。エントロピーに基づくサンプリングでトークンの50\%を保持することで、ピークメモリを最大47\%削減しながら、全トークン訓練と同等、あるいはそれを上回る性能が得られます。しかし、エントロピーだけでは2つ目の重要な領域を見落とします。低エントロピーかつ高ダイバージェンスのトークンだけを切り出すと、全トークンの10\%未満での訓練でも、全トークンのベースラインにほぼ匹敵します。これは、エントロピーのみの規則ではほとんど見えないにもかかわらず、過信トークンが密な修正シグナルを運んでいることを示しています。
私たちはこれらの知見をTIP（Token Importance in on-Policy distillation）として整理します。TIPは、生徒のエントロピーと教師—生徒間ダイバージェンスの2軸による分類体系です。また、エントロピーが有用である一方で、構造的に不完全である理由を理論的に説明します。この見方は、不確実性と不一致（disagreement）を組み合わせる、タイプに応じたトークン選択ルールを動機づけます。私たちは、この理解を、Qwen3、Llama、Qwen2.5の3つの教師—生徒ペアにわたり、MATH-500およびAIME 2024/2025、さらに長期ホライズンのエージェント計画のためのDeepPlanningベンチマークで検証します。そこでは、トークンの$<$$20\%$に対するQ3のみの訓練が、全トークンOPDを上回ります。私たちの実験は、OPDリポジトリ https://github.com/HJSang/OPSD_OnPolicyDistillation を拡張することで実装しており、限られたGPU予算のもとでより大きなモデルをメモリ効率よく蒸留することをサポートしています。