重新思考跨 Tokenizer 的 On-Policy Distillation:从对齐覆盖率到监督可靠性
On-Policy Distillation(OPD) 让学生模型基于自身生成结果接受教师反馈。当师生使用不同 tokenizer 时,比较二者预测需在序列与词表两个层面同时对齐。本文考察:扩大对齐覆盖率能否改善学习? 在数学推理与代码生成的三组异构 teacher-student 配对中,严格的 1:1 对齐 组虽存在显著词表差异,却已覆盖学生生成的大部分 token;在学生蒸馏前采样的回复上,共享词表在严格对齐位置平均保留了几乎全部师生概率质量。将 reverse KL 限制在每个严格位置由学生选出的共享词表 top-16 子集,精度即可与完整共享词表 OPD 相当,并优于所评估的跨 tokenizer 基线。 反之,对 mismatch 组的 span log-probabilities 施加 MSE 监督虽实现完整覆盖,却降低精度。在仅用严格损失训练的 checkpoint 上,span 梯度与严格梯度的方向一致性很弱甚至为负,且相对幅度不断增大,或可解释这一下降。 结论:应从最大化对齐覆盖率转向优先保证 监督可靠性——严格位置上的紧凑监督,往往比引入弱对齐或冲突信号的更广覆盖更有效。
论文精读
TL;DR 本文重新审视跨 tokenizer 在线蒸馏中的对齐覆盖问题,发现 strict 对齐位置上的紧凑监督(如 student 选取的 top-16 共享词汇)即可达到与完全覆盖相当的性能,而额外添加 mismatch 监督反而因梯度冲突降低准确率,主张从追求覆盖全面转向监督可靠性。
问题
问题背景
On-Policy Distillation (OPD) 让学生在自身生成序列上接收教师反馈,是当前跨模型知识迁移的重要路径。当教师和学生使用不同 tokenizer 时,需要在序列级别和词表级别同时对齐预测,这成为该方向的核心工程挑战。
现有方法局限
已有工作通常默认“覆盖越多越好”:要么扩展对齐覆盖(例如在 mismatch groups 上增加 span 监督,或对全部共享词表计算损失),要么设计更复杂的跨 tokenizer 映射。但这样做有三个具体缺陷:
- 弱对齐位置引入噪声:mismatch groups 中的 span 概率与严格位置梯度方向一致性差,可能产生负迁移。
- 监督冗余:实验表明,严格 1:1 组已覆盖大多数 student-generated tokens,共享词表在严格对齐位置平均保留近全部 teacher/student 概率质量。
- 额外监督反而降低 accuracy:加入 MSE span 监督后,虽然实现完整覆盖,但下游准确率下降,梯度诊断显示 span 梯度与 strict 梯度弱相关或负相关。
为什么这个问题难/重要
跨 tokenizer 对齐的困难在于它同时涉及序列切分不一致和词表重叠有限两个维度,很难判断哪些位置是可靠的知识传递锚点。业界普遍希望用更小、更快的 student 模型蒸馏大模型,但若盲目追求对齐覆盖,会导致训练信号冲突、资源浪费,甚至不如仅用少量严格对齐 token 的紧凑监督。本文通过梯度方向一致性和概率质量测量,把问题从“覆盖最大化”转向“监督可靠性”,对实际部署低延迟 student 模型有直接指导意义。
行业类比
类似多语言语音识别中不同 BPE 分词器的协作:与其强制对齐所有子词单元,不如只对齐高置信度的共享 token 锚点,减少噪声梯度,反而获得更稳定的蒸馏效果。
核心洞察
- 严格1:1对齐已经覆盖大部分学生生成token,扩大对齐覆盖至mismatch groups反而因引入弱对齐或冲突监督而降低精度。与以往追求完整覆盖的跨tokenizer蒸馏方法不同,本文通过静态重叠与动态覆盖统计、概率质量分布证明严格位置足以保留预测质量;添加span监督(MSE on span log-probabilities)虽然实现完整覆盖却导致精度下降,梯度诊断显示span梯度与strict梯度方向一致性弱或负相关,且相对规模随训练增长,为“覆盖越多越好”的工程假设提供了反例。
- 学生选择的top-k共享词汇子集(如top-16)可替代全共享词汇监督,且梯度方向一致性可作为是否添加额外监督的预检指标。相比使用全共享词汇或复杂对齐的基线,紧凑子集在保持性能的同时降低计算开销;而span梯度与strict梯度的弱/负方向一致性及相对规模变化,可提前判断额外监督是否可能造成冲突,为跨tokenizer蒸馏的工程实践提供了可操作的诊断工具,其核心启示是优先考虑监督信号可靠性而非单纯扩大覆盖范围。
方法
方法详解
输入:异构 tokenizer 的教师模型与学生模型;学生模型在 on-policy 模式下采样生成响应序列。
关键模块:
- 序列级对齐:将教师生成的 token 序列与学生生成的 token 序列进行 token-group alignment,划分出 strict 1:1 对应组(相同位置 token 对齐)和 mismatch 组(长度不一致的片段)。
- 共享词汇监督:在 strict 组位置上,仅使用 teacher 和 student 共享词汇表上的概率分布进行计算。反向 KL 损失作为蒸馏目标;可进一步限制为学生选择的 top-k(例如 top-16)共享词汇子集,以降低计算与存储开销。
- Mismatch span 监督(可选):对 mismatch 组中的 span 对数概率添加 MSE 损失,以实现理论上全覆盖的监督。但实验表明该监督会降低最终准确率。
- 梯度诊断:在仅使用 strict loss 训练的检查点处,计算 span 梯度与 strict 梯度的方向一致性;发现 span 梯度与 strict 梯度方向一致性弱甚至为负,且梯度幅值相对增长,解释了添加 span 监督的负面效果。
输出:训练后的学生模型,在数学推理和代码生成任务上,使用紧凑的 strict top-k 监督即可达到与全共享词汇 OPD 相当的性能,超越其他跨 tokenizer 基线。
与同类方法差异:不同于以往追求最大化 token 对齐覆盖率的跨 tokenizer OPD 方法,本文主张从“覆盖率优先”转向“监督可靠性优先”,通过仅覆盖 strict 位置上的紧凑监督获得更优性能。
实验
实验设计
作者构建了三组异构 teacher-student 对,覆盖 数学推理 和 代码生成 两类任务,在 student 自身生成的响应上进行 On-Policy Distillation(OPD)。重点考察跨 tokenizer 场景下的对齐覆盖:统计 strict 1:1 groups 的 token 覆盖比例,并测量 shared vocabulary 在严格对齐位置的概率质量。进一步提出 student-selected top-16 subset 蒸馏,并与添加 span MSE 监督 的完整覆盖方案对比。
关键发现
- strict 1:1 groups 在词汇不匹配明显的情况下,仍覆盖了大部分 student 生成 token。
- shared vocabulary 在严格对齐位置平均保留了近乎全部的 teacher 与 student 概率质量。
- 将 reverse KL 限制在每个 strict 位置的 学生选择 top-16 子集,取得与 full shared-vocabulary OPD 相当的准确率,且优于所评估的跨 tokenizer 基线。
- 增加 span MSE 监督后准确率反而下降。梯度诊断显示 span gradients 与 strict gradients 方向一致性弱或为负,且训练后期相对幅度增大。
与基线对比解读
传统跨 tokenizer 蒸馏通常追求最大化 alignment coverage,例如通过 mismatch groups 的 span 对齐覆盖所有 token。本工作表明,compact 且可靠的监督信号比 broad coverage 更有效:student 自身生成分布下,strict 位置已承载绝大部分预测质量,top-16 子集已能捕获足够 teacher 信息。相反,span 监督引入弱对齐或冲突梯度,损害学习。这提示工程实践中应优先保证监督信号的可靠性,而非无脑扩展覆盖范围。
行业影响
落地场景
跨 tokenizer 的在线策略蒸馏(OPD)常见于多模型协同与端侧部署:例如将云端大模型蒸馏到移动端小模型但两者 tokenizer 不同,多厂商模型集成做推理加速,或代码生成 IDE 插件需要将大模型压缩到本地模型。论文发现严格对齐位置已覆盖大部分学生生成 token,且共享词表 top-16 子集足够,可显著简化蒸馏训练,减少对齐计算与内存开销。
商业价值
- 降本:无需在 mismatch groups 上做 span 监督,且严格位置只需 top-16 子集,损失计算从全词表缩小到 16 个 logit,大幅降低训练算力与存储需求。
- 稳定性:避免添加 span 监督带来的梯度冲突,使蒸馏训练更稳定,减少调参成本与训练失败风险。
- 精度保持:在数学推理与代码生成任务上,top-16 子集监督与全词表监督精度相当,优于现有跨 tokenizer 基线。
与现有工作流集成
可直接替换现有 OPD 训练中的损失模块:在严格对齐位置,对学生生成的 token 对应的共享词表 logits 计算 reverse KL,仅取学生概率 top-16 的共享 token;跳过 mismatch groups 的监督。无需改动模型架构或数据 pipeline,仅调整损失计算与梯度回传范围。可封装为训练框架插件(如 DeepSpeed、Accelerate 的自定义 loss),适用于现有蒸馏代码库。
具体用例
- 电商智能客服:将大参数客服模型蒸馏到端侧小模型,用于移动 App 实时回复。原 tokenizer 与教师不同,使用 top-16 子集监督可大幅降低训练成本,同时保持回复质量。
- 代码生成 IDE 插件:将代码大模型蒸馏到本地轻量模型,跨 tokenizer 场景下采用该方法可减少训练显存占用,使开发者能在消费级 GPU 上完成蒸馏微调。
局限
- **实验领域与规模有限**:本研究仅在数学推理和代码生成两个领域、三个异构教师-学生对(教师约 7B-13B,学生约 1B-3B)上验证,未覆盖更大规模模型或更复杂任务(如多模态、长上下文生成)。结论对 tokenizer 差异极大的场景(如字节级与 BPE 混合)是否成立尚不明确,限制了向工业级应用的直接迁移。
- **top-k 子集选择缺乏敏感性分析**:作者固定 k=16 在每个严格位置,声称与全共享词汇相当,但未系统研究 k 值(如 4、8、32、64)对性能的影响,也未在不同任务或模型规模下验证最优 k 的稳定性。此外,严格对齐仅覆盖 1:1 token 组,对非严格对齐位置采用 span 监督产生负面效果,但论文没有提出替代方案来利用这些未对齐信号,可能遗漏了有价值的监督信息。
- **梯度诊断缺乏理论支撑与改进方案**:论文通过梯度方向一致性分析解释了 span 监督为何无效,但该分析是事后观测,未给出严格理论证明,也未尝试通过正则化、梯度裁剪或动态加权来修复 span 监督。实验均使用 AdamW 与固定学习率,未验证其他优化器或学习率调度下的结论稳健性,因此“监督可靠性优先于覆盖”的观点尚待更广泛验证。