SUGGERIMENTO: Importanza dei Token nella Distillazione On-Policy

TIP: Token Importance in On-Policy Distillation

April 15, 2026

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

cs.AI

Abstract

La distillazione della conoscenza on-policy (OPD) addestra uno studente sulle sue stesse traiettorie di esplorazione (rollout) sotto la supervisione a livello di token di un insegnante. Non tutte le posizioni dei token hanno la stessa importanza, ma le visioni esistenti sull'importanza dei token sono incomplete. Ci poniamo una domanda diretta: quali token veicolano il segnale di apprendimento più utile nell'OPD? La nostra risposta è che i token informativi provengono da due regioni: le posizioni con alta entropia dello studente e le posizioni con bassa entropia dello studente più alta divergenza insegnante-studente, dove lo studente è sovracconfidente e sbaglia. Empiricamente, l'entropia dello studente è un forte proxy del primo ordine: trattenere il 50% dei token con un campionamento basato sull'entropia eguaglia o supera l'addestramento su tutti i token, riducendo al contempo la memoria di picco fino al 47%. Ma l'entropia da sola trascura una seconda regione importante. Quando isoliamo i token a bassa entropia e alta divergenza, l'addestramento su meno del 10% di tutti i token si avvicina molto ai baseline con tutti i token, dimostrando che i token sovracconfidenti veicolano un segnale correttivo denso nonostante siano quasi invisibili alle regole basate solo sull'entropia. Organizziamo questi risultati con TIP (Importanza del Token nella distillazione on-Policy), una tassonomia a due assi su entropia dello studente e divergenza insegnante-studente, e forniamo una spiegazione teorica del perché l'entropia è utile ma strutturalmente incompleta. Questa visione motiva regole di selezione dei token "type-aware" che combinano incertezza e disaccordo. Convalidiamo questo quadro su tre coppie insegnante-studente che coinvolgono Qwen3, Llama e Qwen2.5 su MATH-500 e AIME 2024/2025, e sul benchmark DeepPlanning per la pianificazione agentica a lungo orizzonte, dove l'addestramento solo su Q3 con <20% dei token supera l'OPD con tutti i token. I nostri esperimenti sono implementati estendendo il repository OPD https://github.com/HJSang/OPSD_OnPolicyDistillation, che supporta la distillazione efficiente in memoria di modelli più grandi con budget GPU limitati.

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.