不是所有分歧都可学习：On-Policy Distillation 中的 Token Teachability

TL;DR On-policy 蒸馏中，并非所有教师-学生分歧都可学；本文提出 token 可教性概念，并构建 TA-OPD，仅用 5% 高可教性 token 即超越全量 OPD，重构了选择性蒸馏的准则。

问题

问题背景

On-policy 蒸馏 (OPD) 是训练高效语言模型的核心范式：学生模型在自己的生成轨迹上，以教师模型的 token 级分布作为监督信号进行优化。由于不同 token 位置的学习价值差异巨大，近期工作转向 选择性 OPD，希望只保留“信息量大”的 token，以此降低训练成本并改善收敛。

现有方法局限

主流选择性策略依赖 启发式指标，如高熵 (high entropy) 或高 KL 散度 (high divergence)。这些指标隐含假设：教师与学生分歧越大的位置，越值得学生学习。但该假设存在根本缺陷：

• 混淆信号类型：KL 散度无法区分“可学习的分歧”(教师将概率质量集中到学生 top-K 候选项上) 与“不相容的分歧”(教师偏爱学生当前完全忽略的 token，仅因分布形状差异导致高 KL)。
• 固定预算下效果不佳：基于熵或散度的选择经常保留大量学生无法吸收的监督信号，浪费训练计算，甚至在极端情况下引入噪声，损害最终模型质量。

为什么这个问题难且重要

• 技术挑战：在 token 级别判断“可学习性”需要同时建模学生当前的局部支持集 (local support) 与教师信号的分布重叠，这是细粒度的分布兼容性问题，无法用单点标量散度捕捉。
• 工程意义：OPD 的成本与 token 数成正比，大规模语言模型训练中减少 95% 以上 token 可带来显著加速；但若选择标准粗糙，极度压缩预算时性能会急剧下降。因此，从“选择显著 token”转向“选择可学习 token”是提升 训练效率上限 的关键。

行业类比：这类似强化学习中的 经验回放 (experience replay) 优先级设计——不是所有高 TD-error 的转移都值得重放，不可学习的意外噪声反而会破坏策略；类似地，OPD 需要感知 token 的局部可塑性。

核心洞察

• **KL 分歧是粗粒度的学习价值代理，混淆了可学习与不可兼容的分歧**：原始 KL 值高并不保证对应的教师信号能被学生吸收。论文提出的 **token teachability** 将教师分布与学生的局部 top-K 支持集进行对齐度量，区分出“可学习分歧”（教师纠正性质量落在学生候选内）和“不可兼容分歧”（教师质量落在学生当前支持集之外）；前者才能可靠预测固定上下文下的 KL 下降，后者则成为噪声。这一视角揭示了传统基于熵或裸 KL 的选择策略的根本缺陷，为理解 OPD 中的信号效率提供了新的理论基础。
• **选择性 OPD 的重心应从“显著 token”转向“可学习 token”**：现有方法优先选择高熵或高分歧 token，但实验发现低熵高分歧区域存在大量不可教 token，而高熵低分歧区域则常被忽视。**TA-OPD** 通过轻量的可教性得分（仅依赖师生输出分布，无需奖励模型或验证器）进行 token 保留，在 Qwen2.5/Qwen3 多个模型对上仅用 5% 的 token 就超过了全 token OPD，并优于基于熵和分歧度的基线。这重新定义了选择性蒸馏的目标——不是挑出教师信号最强的 token，而是挑出学生最能从中获益的 token。

方法

输入与诊断框架

TA-OPD 的输入是学生模型在自身 rollouts 上产生的 token 序列，以及教师模型对每个 token 位置的 top-K 概率分布。该方法首先基于固定上下文诊断（fixed-context diagnostic）量化每个 token 位置的可学习性：将教师-学生间的 KL 分歧分解为两类——可学习分歧（教师概率质量集中于学生当前 top-K 候选内）与不可兼容分歧（教师质量主要落在学生支持集之外）。可学性被定义为token teachability，即教师分布与学生局部知识状态的兼容程度。

核心模块：可教性评分与选择性蒸馏

1. 可教性评分计算：对于每个 token 位置，计算一个连续分数，它反映教师分布与学生 top-K 候选的重叠度（例如，教师质量在学生 top-K 上的累积概率）。该分数无需奖励模型或验证器，仅依赖当前 logits。
2. 预算控制的选择策略：给定一个保留比例（如 5%），在训练过程中仅对可教性最高的 token 子集施加 OPD 损失（如 forward KL），其余位置不参与梯度更新。这相当于将教师信号过滤为“学生有能力吸收”的部分。
3. 轻量实现：所有计算均在线完成，不引入额外网络，也不依赖下游任务信号，保持了蒸馏流程的简洁性。

输出与优化目标

通过最小化所选高可教性 token 上的教师-学生 KL 散度，更新学生参数。目标是让学生仅从可学习的教师信号中获益，避开会引入噪声或冲突的不可兼容信号。

与同类方法的差异

传统选择性 OPD 方法（如基于熵或原始 KL 分歧）仅关注信号的显著性，而 TA-OPD 强调信号的可学习性：它区分了教师是否在学生的“最近发展区”内提供修正，从而在有限预算下更稳健地提升学生性能。

实验

实验设计

作者构建固定上下文诊断协议，测量同一上下文下教师-学生KL降低，解耦可学习不一致（教师信号与学生top-K候选有交集）与不相容不一致（教师支持与学生支持无重叠），形式化token teachability。在Qwen2.5和Qwen 3教师-学生对上，TA-OPD根据局部支持兼容性选择token位置，仅对高teachability位置施加OPD损失，无需额外模型或验证器。基线包括全token OPD、基于熵的选择和基于KL散度的选择。

关键发现

• 原始KL不一致作为学习价值指标粗糙，token teachability更精准预测固定上下文增益。
• TA-OPD仅在保留5% token的情况下，性能经常超过全token OPD，并全面优于熵基和散度基基线。
• 高熵–高不一致token区域非均匀可教，许多教师信号与学生当前分布不相容，盲目选择反而引入噪声。

基线对比深度解读

传统选择性OPD依赖显著性启发式（如高熵或高KL），这些方法忽略信号可学性。TA-OPD通过局部支持分解筛选可迁移信号，避免了无效梯度，将选择性蒸馏重构为“选可学信号”而非“选显著token”。这种设计不仅提升样本效率，还保证了蒸馏过程的稳定与方向正确，为未来的token级知识迁移提供了新范式。

行业影响

落地场景

TA-OPD 的核心思想——根据 token 的可教学性（teachability）选择性应用蒸馏损失——可直接嵌入任何使用 on-policy distillation 训练大语言模型的 pipeline。典型场景包括：

• 对话助手/客服模型的在线蒸馏：在部署阶段，用更大模型（教师）实时纠正较小模型（学生）的生成偏差，仅对高可教学性 token 施加损失，显著节省训练成本。
• 企业级模型微调：在私有数据上蒸馏自建小模型时，自动筛选出教师能真正“教会”学生的 token 位置，避免在不兼容分歧上浪费算力，提升收敛效率。
• 多智能体协作：当多个智能体互相切磋时，可基于 teachability 动态决定在哪些 token 上采纳对方的建议，实现更高效的集体学习。

商业价值

• 降本：原始方法通常要求对所有 token 计算 KL 散度，计算和通信开销巨大；TA-OPD 仅需 5% 的 token 即可达到甚至超越全量效果，直接降低 GPU 时间和电力成本。
• 增效：过滤掉不可教学的分歧后，梯度更新更精准，模型收敛更快，实验结果表明在多个 Qwen 系列教师-学生对中，低开销的 TA-OPD 优于基于熵或基于发散度的选择基线，意味着用更少资源达到更好质量。
• 体验提升：在对话式 AI 产品中，更高效的蒸馏可加快模型迭代速度，从而更快响应用户行为变化和需求演进。

与现有工作流的接口

TA-OPD 以轻量即插即用的方式集成：

• 只需在学生模型生成 rollout 时，额外计算一个 teachability score（基于教师分布中落在学生 top-K 候选内的概率质量），无需附加奖励模型或验证器，与现有的 Hugging Face 训练框架或 DeepSpeed 等分布式训练环境兼容。
• 可将 teachability selector 封装为自定义损失模块，替换或叠加到既有的蒸馏损失函数上，对现有代码侵入性小。
• 适用于任何基于 Transformer 的自回归模型，无论预训练还是微调阶段，都可作为选择性 OPD 的标准组件。

具体落地案例

1. 电商平台智能客服知识蒸馏：大型客服系统常采用教师-学生架构，教师模型（数千亿参数）负责离线生成高质答案，学生模型（百亿参数）在线部署。使用 TA-OPD 在线蒸馏时，系统仅针对那些学生有足够词汇覆盖但尚需纠偏的 token 进行学习，显著减少云端实时推理的额外开销，同时维持答复质量不降。
2. 低资源语言翻译模型快速迭代：在一个面向全球市场的翻译服务中，需要频繁用最新语料更新小模型。TA-OPD 可自动筛选出教师分布中能被学生“领悟”的词汇单元，跳过 token 级别的语言隔阂导致的不可学分歧，从而在有限训练预算下加速多语言模型的周期更新。

局限

• **方法依赖 Top-K 支持与超参选择**：Token Teachability 的计算基于教师分布对学生 Top-K 候选的局部兼容性，K 的取值直接影响可学习判断的粒度。论文虽在实验中固定 K=10 并做灵敏度分析，但该值可能对不同规模模型或不同生成场景不够鲁棒，额外超参引入会降低方法的即插即用性。
• **实验场景集中在语言模型，泛化性待验证**：TA-OPD 仅在 Qwen2.5 和 Qwen 3 系列的文本生成任务上评估，未涉及多模态模型或编码器-解码器架构。On-policy 蒸馏在其他领域（如语音、视觉）的应用可能面临不同的分布特性，Teachability 是否同样有效尚不清楚。
• **计算开销尚难完全忽略**：虽然选 token 的得分计算不依赖 reward model 或 verifier，但需实时获取学生 logits 和教师分布的对数几率排序，相比于简单的熵或散度选择器仍有额外运算。在大批量在线蒸馏中，该开销可能累积，论文未分析其吞吐量或延迟影响。