找出强化学习中最有用的注意力位置,提升训练效率。
TIP: Token Importance in On-Policy Distillation

- 按学生模型不确定性和师生差异双重标准筛选关键token
- 仅保留50%高熵token即可达全量训练效果,内存降低47%
- 超自信错误位置蕴含强纠错信号,<10%令牌训练接近全量表现
在策略知识蒸馏(OPD)中,学生模型基于自身采样结果接受教师指导,但并非所有标记位置都同等重要。本文提出直接问题:哪些标记承载最有效的学习信号?研究发现,有效信号来自两类区域:高学生熵的位置,以及低熵但师生差异大的位置——此时学生过于自信且出错。实验证明,仅基于熵选择50%的标记,性能与全量训练相当,峰值内存减少47%;而单独提取低熵高差异标记,使用不足10%的标记训练即接近全量基线,说明超自信错误位置蕴含密集纠正信号。为此构建TIP(Token Importance in on-Policy distillation)双轴分类框架,涵盖学生熵与师生差异,并给出理论解释。该视角推动了结合不确定性与分歧度的类型感知选择策略。在覆盖Qwen3、Llama、Qwen2.5的三组师生对上,于MATH-500、AIME 2024/2025及DeepPlanning长程智能体规划任务中均验证有效性。其中,仅用<20%标记的Q3训练超越全量OPD。实验基于扩展的OPD仓库实现,支持有限显存下大模型高效蒸馏。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。