TIP: Token Importance in On-Policy Distillation

📄 arXiv: 2604.14084v1 📥 PDF

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

分类: cs.LG, cs.AI

发布日期: 2026-04-15

🔗 代码/项目: GITHUB

💡 一句话要点

提出TIP：基于Token重要性的On-Policy蒸馏方法，提升训练效率并降低内存占用

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 知识蒸馏 On-Policy学习 Token重要性 学生熵 教师-学生差异 模型压缩 长序列建模

📋 核心要点

  1. 现有On-Policy知识蒸馏方法对所有token同等对待，忽略了不同token对学生模型学习贡献的差异性，导致训练效率低下。
  2. 论文提出TIP框架，通过学生熵和教师-学生差异性双轴分类，识别携带关键学习信号的token，并设计类型感知token选择规则。
  3. 实验结果表明，仅使用少量关键token进行训练，即可达到甚至超过全token训练的效果，显著降低内存占用，提升训练效率。

📝 摘要（中文）

On-policy知识蒸馏(OPD)通过教师模型对学生模型在自身rollout上的token级监督进行训练。并非所有token位置都同等重要，但现有的token重要性观点并不完整。本文直接提出问题：在OPD中，哪些token携带最有用的学习信号？答案是：信息丰富的token来自两个区域：具有高学生熵的位置，以及具有低学生熵和高教师-学生差异的位置，此时学生模型过度自信且错误。实验表明，学生熵是一个强大的初步代理：保留50%的基于熵采样的token，其性能与全token训练相匹配或超过，同时减少高达47%的峰值内存。但仅靠熵会遗漏第二个重要区域。当隔离低熵、高差异的token时，在不到10%的token上进行训练几乎与全token基线相匹配，表明过度自信的token携带密集的校正信号，尽管熵规则几乎无法发现它们。本文使用TIP（On-Policy蒸馏中的Token重要性）组织这些发现，这是一个基于学生熵和教师-学生差异的双轴分类法，并给出了熵有用但结构上不完整的理论解释。这种观点激发了结合不确定性和分歧的类型感知token选择规则。在MATH-500和AIME 2024/2025上，以及在用于长时程智能体规划的DeepPlanning基准上，本文在Qwen3、Llama和Qwen2.5三个教师-学生模型对上验证了这一观点，其中仅使用Q3训练<20%的token就超过了全token OPD。实验通过扩展OPD存储库https://github.com/HJSang/OPSD_OnPolicyDistillation来实现，该存储库支持在有限的GPU预算下对更大的模型进行内存高效的蒸馏。

🔬 方法详解

问题定义：论文旨在解决On-Policy知识蒸馏(OPD)中token重要性评估不准确的问题。现有OPD方法平等对待所有token，忽略了不同token对学生模型学习的贡献差异，导致计算资源浪费和训练效率低下。尤其是在处理长序列时，这种问题更加突出。

核心思路：论文的核心思路是识别并利用携带关键学习信号的token进行训练。作者认为，高学生熵的token代表学生模型的不确定性，而低学生熵但高教师-学生差异的token代表学生模型的过度自信和错误。通过关注这两类token，可以更有效地指导学生模型的学习。

技术框架：TIP框架主要包含以下几个阶段：1) 使用教师模型和学生模型进行rollout，生成token序列。2) 计算每个token的学生熵和教师-学生差异。3) 基于学生熵和教师-学生差异，将token划分为不同的类型。4) 根据预设的类型感知token选择规则，选择一部分token用于训练。5) 使用选择的token计算损失函数，并更新学生模型的参数。

关键创新：论文最重要的技术创新在于提出了基于学生熵和教师-学生差异的双轴token重要性评估方法。这种方法能够更准确地识别携带关键学习信号的token，从而提高训练效率和模型性能。与现有方法相比，TIP框架能够更好地利用有限的计算资源，实现更高效的知识蒸馏。

关键设计：在关键设计方面，论文提出了类型感知的token选择规则，例如，优先选择高学生熵的token，或者选择低学生熵但高教师-学生差异的token。此外，论文还探索了不同的教师-学生差异度量方法，例如KL散度。具体的损失函数设计取决于具体的蒸馏任务，但通常会包含教师模型的输出和学生模型的输出之间的差异项。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在MATH-500和AIME数据集上，使用TIP框架仅保留50%的基于熵采样的token，其性能与全token训练相匹配或超过，同时减少高达47%的峰值内存。更重要的是，在DeepPlanning基准测试中，仅使用Qwen3训练<20%的token就超过了全token OPD，验证了TIP框架的有效性。

🎯 应用场景

该研究成果可广泛应用于自然语言处理、强化学习等领域，尤其是在需要进行长序列建模和知识蒸馏的场景下。例如，可以用于训练更高效的对话系统、智能客服、代码生成模型等。通过减少训练所需的计算资源，该方法有助于推动更大规模模型的训练和部署。

📄 摘要（原文）

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.