ASTUCE : Importance des tokens dans la distillation en politique active

TIP: Token Importance in On-Policy Distillation

April 15, 2026

Auteurs: Yuanda Xu, Hejian Sang, Zhengze Zhou, Ran He, Zhipeng Wang, Alborz Geramifard

cs.AI

Résumé

La distillation de connaissances en mode « on-policy » (OPD) entraîne un étudiant sur ses propres déroulements (« rollouts ») sous la supervision au niveau des tokens d’un enseignant. Toutes les positions de token n’ont pas la même importance, mais les conceptions actuelles de l’importance des tokens sont incomplètes. Nous posons une question directe : quels tokens portent le signal d’apprentissage le plus utile dans l’OPD ? Notre réponse est que les tokens informatifs proviennent de deux régions : les positions à forte entropie de l’étudiant, et les positions à faible entropie de l’étudiant mais avec une forte divergence enseignant–étudiant, où l’étudiant est surconfiant et se trompe. Empiriquement, l’entropie de l’étudiant est un bon indicateur de premier ordre : conserver 50 % des tokens par un échantillonnage basé sur l’entropie égal ou dépasse l’entraînement sur tous les tokens tout en réduisant la mémoire maximale jusqu’à 47 %. Mais l’entropie seule néglige une deuxième région importante. Lorsque nous isolons les tokens à faible entropie et haute divergence, l’entraînement sur moins de 10 % de tous les tokens atteint presque les performances de base avec tous les tokens, montrant que les tokens surconfiants portent un signal correctif dense bien que quasi invisibles pour les règles basées uniquement sur l’entropie. Nous organisons ces résultats avec TIP (« Token Importance in on-Policy distillation »), une taxonomie à deux axes basée sur l’entropie de l’étudiant et la divergence enseignant–étudiant, et fournissons une explication théorique expliquant pourquoi l’entropie est utile mais structurellement incomplète. Cette perspective motive des règles de sélection de tokens « conscientes du type » qui combinent incertitude et désaccord. Nous validons cette vision sur trois paires enseignant–étudiant couvrant Qwen3, Llama et Qwen2.5 sur MATH-500 et AIME 2024/2025, ainsi que sur le benchmark DeepPlanning pour la planification agentique à long terme, où un entraînement Q3 uniquement sur <20 % des tokens surpasse l’OPD avec tous les tokens. Nos expériences sont implémentées en étendant le dépôt OPD https://github.com/HJSang/OPSD_OnPolicyDistillation, qui prend en charge la distillation efficace en mémoire de modèles plus grands avec des budgets GPU limités.

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.