BiToK-SD用双层优化学习Top-K token选择，在数学推理自蒸馏中三个规模均取得最佳平均成绩
核心概要
该工作提出BiToK-SD，把在线策略自蒸馏中的Top-K token选择建模为双层优化问题：下层用可微的阈值松弛把离散Top-K选择变成随学生策略演化的软权重，上层只在被选位置上做知识蒸馏；在Qwen3-1.7B/4B/8B上、以OpenThoughts数学推理子集约3万条数据训练200步，于AIME24、AIME25、HMMT25的Avg@64评测中取得所有对比方法里最高的平均成绩，并比最强基线OPSD分别高出0.63、0.58、0.21个平均分。
Figure 1: Exploration dynamics during training. (a) Student entropy over training steps. (b) Student–teacher entropy gap over all tokens. (c) Student–teacher entropy gap on selected tokens. BiToK-SD maintains higher student entropy and larger student–teacher entropy gaps than the baselines, suggesting stronger policy exploration.
arXiv深度剖析
提出把自蒸馏中的token选择写成双层优化：下层以可微阈值松弛近似Top-K选择，上层在选中的token位置上最小化加权自蒸馏损失，从而同时学习“在哪里蒸馏”和“以多强力度蒸馏”。 此前的选择性蒸馏要么对所有token均匀蒸馏，要么用熵、师生散度等固定启发式分数做二值Top-K选择，且被选中的位置权重相同；BiToK-SD让选择阈值随学生策略在线更新，并给出连续而非二值的token权重。 方法在正文第3节完整给出，包括效用分数、松弛阈值目标、软掩码与加权损失；消融显示掩码密度稳定在目标稀疏度附近，阈值在训练中持续变化（1.7B上从0.061升至0.063、范围0.044–0.105；4B从0.030到0.102；8B从0.069到0.159）。
在三个模型规模的数学推理基准上，BiToK-SD取得所有对比方法中最高的平均成绩，并在九个“规模×基准”单元中七次最优。 相对最强基线OPSD，平均分在1.7B、4B、8B上分别提升0.63、0.58、0.21；相对使用相同Soft-OR分数但硬Top-K掩码与反向KL的TIP，在全部九个单元中均领先。 评测为AIME24、AIME25、HMMT25的thinking模式Avg@64（每题64次独立采样），所有方法训练200步并在最终检查点评估，且刻意不做基于测试集的检查点选择；各方法从同一检查点、同一数据初始化。
学习到的选择器在保持目标稀疏预算的同时，把更大权重分配给与反思相关的token位置，并维持更强的策略探索。 OPSD对所有有效位置均匀加权，TIP使用二值Top-K掩码，二者对反思token与其他被选token的相对权重在构造上就是均匀的；BiToK-SD的连续权重随效用分数相对自适应阈值而增大。 以十个认知标记词（wait、hmm、perhaps、maybe、actually、alternatively、seems、might、likely、check）作为自反思行为的词法指标，BiToK-SD每次补全生成的标记多于OPSD与TIP，差距约在第60步出现并扩大；反思相关位置获得的平均软掩码权重高于其他有效位置；同时学生熵与师生熵差更大，且AIME24上生成token数比OPSD少3.7%、1.6%、4.8%，重复4-gram比例与截断率更低。
该双层选择机制可迁移到离线策略蒸馏：在仅有截断教师分布可用的设定下，仍优于全位置蒸馏与基于学生熵的选择方法。 离线设定中因存储完整教师分布代价过高，只能缓存top-64概率，Soft-OR所需的完整分布KL无法忠实计算，因此改用仅依赖学生预测的学生熵分数，选择比例固定为10%。 按Tavor等人的协议，在FineWeb的8000万token子集上把Qwen3-8B蒸馏进Qwen3-1.7B，训练一个epoch、约2.4K次优化更新；在ARC-Easy、GSM8K、HellaSwag、PIQA、LAMBADA五任务平均准确率、IFEval严格指令遵循分数与LAMBADA困惑度三项汇总指标上均优于Full-KD、SE-KD与SE-KD3X。
启示与展望
该结果面向资源受限场景下压缩推理模型的自蒸馏：在Qwen3-1.7B/4B/8B上、以OpenThoughts数学推理子集约3万条数据训练200步，并在AIME24、AIME25、HMMT25的thinking模式Avg@64下评估，适用于希望在不依赖外部大教师模型的前提下提升小模型数学推理的实践者。方法本身不绑定特定效用分数，可替换为学生熵、教师熵、前向KL或反向KL；离线实验进一步表明，在只能缓存教师top-64概率、无法计算完整分布KL时，改用学生熵分数仍可在通用语料蒸馏中带来增益。
效用分数与选择比例的选择仍会影响相对均匀加权的增益：消融中Soft-OR与前向KL优于均匀加权，而基于熵的分数与反向KL未超过均匀加权，各分数间Avg@64差异约在两点以内；Top-K比例呈倒U形并在50%处最优，说明收益依赖中间稀疏度。TIP与BiToK-SD同时存在选择规则与散度两处差异，正文也指出该差距不能单独归因于任一因素。离线设定中只能使用截断教师分布，Soft-OR无法忠实计算，因此该场景的结论建立在学生熵分数之上。此外，反思行为以十个认知标记词作为词法代理，属于间接指标；主实验限于数学推理基准与Qwen3系列，其他推理领域与模型族的适用性仍是开放问题。