重新审视在策略蒸馏：经验性失败模式与简单修复方案

识别采样token蒸馏三大失败模式，提出top-K匹配+19.8%提升

前置知识

在策略蒸馏 (On-Policy Distillation, OPD)

在策略蒸馏是一种LLM后训练方法，核心思想是让学生模型在自己生成的轨迹（rollouts）上进行训练，同时用更强的教师模型对这些轨迹提供逐token的监督信号。与离策略蒸馏（如SFT在教师生成数据上训练）不同，OPD的训练数据来自学生自身的分布，这样可以避免训练时和推理时prefix分布不匹配的问题。标准实现中，每一步更新基于教师和学生在采样token上的对数概率比（log-ratio）。

OPD是本文研究的核心对象，理解其标准实现和动机是理解本文改进方案的前提。

反向KL散度 (Reverse KL Divergence)

反向KL散度 $D_{KL}(p \| q) = \sum_x p(x) \log \frac{p(x)}{q(x)}$ 衡量用分布q近似分布p时的信息损失。在OPD中，学生$p_\theta$向教师q靠近，使用反向KL意味着学生倾向于避免在教师低概率区域分配过多概率（mode-seeking行为），这与前向KL（mass-covering）形成对比。反向KL更适合生成任务，因为它倾向于产生更精确的输出而非过于保守的输出。

OPD的目标函数就是最小化学生到教师的反向KL散度，理解这一散度的性质才能理解为什么token-level近似是有偏差的但方差更低。

偏差-方差权衡 (Bias-Variance Tradeoff)

在梯度估计中，偏差指估计量的期望与真实梯度的系统性偏离，方差指估计量在不同采样中的波动程度。Token-level OPD去掉了未来奖励耦合项（future-reward coupling），因此是有偏的，但方差更小；序列级OPD无偏但方差更大。在长序列场景下，这种权衡至关重要：序列级估计的方差上界为$O(T^4)$，而token-level仅为$O(T^2)$。

本文的核心理论贡献就是阐明了这一权衡，并据此选择token-level路线，然后在其上做改进。

top-K / top-p 采样

Top-K采样是从概率最高的K个token中进行采样；top-p（nucleus）采样是选择累积概率达到阈值p的最小token集合并从中采样。本文中，top-K用于构建教师分布的支持集（support set），top-p用于控制rollout生成的稳定性。两者都是限制采样空间、避免低质量token的方法。

本文方法的核心就是用教师的top-K支持集来替代单token监督，理解top-K/top-p机制是理解方法的必要条件。

研究动机

在当前的LLM后训练流水线中（如Qwen3、MiMo-V2-Flash、GLM-5等），在策略蒸馏（OPD）已成为标准组件。标准实现采用采样token比较：在每一步训练中，更新仅由教师和学生在采样token上的对数概率差驱动。这种单token监督信号在长rollout场景下变得非常脆弱。具体表现为三个失败模式：第一，token级信号严重不平衡——大多数采样token获得负奖励（学生概率高于教师），优化被少数正奖励token主导，训练对高频filler和短续写敏感；第二，在学生生成的prefix上教师信号不可靠——当学生轨迹偏离教师分布后，教师仍可能在重复、自循环等无意义续写上给出高概率，造成token级教师一致性和轨迹级质量之间的目标错位；第三，tokenizer或特殊token不匹配——当师生使用不同tokenizer时（如学生将''拆分为'<','think','>'，教师期望''），相同的语义内容会被误判为不匹配。这些失败模式在数学推理任务中已经被反复观察到，也与Gu et al. (2024)报告的输出退化和Ko et al. (2026)报告的entropy collapse一致。

本文的目标是本文的具体目标是在保持token-level OPD方差优势（$O(T^2)$ vs $O(T^4)$）的前提下，改善其局部监督信号的质量，使其不再依赖单一采样token的log-ratio，而是在每个prefix上对教师支持的token子集进行分布级比较。具体来说，作者希望实现：(1) 更稳定的优化过程（更小的梯度范数、更少的梯度裁剪触发）；(2) 更好的下游性能，在数学推理和智能体任务上超过标准采样token OPD；(3) 减轻tokenizer不匹配带来的负面影响。

与已有工作不同的是，已有工作尝试用全词汇蒸馏（full-vocabulary distillation）来解决单token信号的信息丢失问题，但全词汇KL在内存上对长序列和大模型来说代价过高。Zhang & Ba (2026)提出了无偏的top-K KL估计器，但计算复杂度仍较高。本文的独特切入角度是：不追求无偏估计或全词汇匹配，而是在教师top-K支持集上计算截断反向KL，同时对师生分布做归一化——这是一个在信息量和计算成本之间的实用中间点。更重要的是，本文不仅提出方法，还系统地诊断了采样token OPD失败的具体机制，将理论分析（偏差-方差权衡）和经验分析（三个失败模式）结合，为方法设计提供了明确的动机。

核心方法

本文的方法可以用一个类比来理解：标准采样token OPD就像只通过一个采样点来评判一幅画——如果恰好采到画的背景区域，就会给出过于负面的评价；如果采到高光区域，又可能给出过于正面的评价。本文的做法是让教师先圈出画中最重要的K个区域（top-K支持集），然后在这K个区域上比较师生的'评价分布'，这样即使单个区域有偏差，整体信号也更稳定可靠。技术路线分为三步：(1) 在每个prefix位置，取教师概率最高的K个token作为支持集；(2) 将师生分布在这个支持集上重新归一化；(3) 计算截断反向KL作为局部损失，再对所有rollout位置取平均。此外，配合top-p rollout采样和特殊token masking，形成完整的训练方案。

核心创新点在于将单token监督替换为教师top-K局部支持匹配（Local Support Matching, LSM）。与已有方法的本质区别体现在三个层面：第一，相对于采样token OPD，LSM不再将整个更新压在一个token的log-ratio上，而是在K个teacher支持的token上做分布级比较，消除了单点估计的脆弱性；第二，相对于全词汇KL，LSM只在top-K子集上计算，将计算成本从$O(|V|)$降低到$O(K)$（K通常为32），同时保留了分布级比较的优势；第三，师生分布在截断支持上重新归一化（各自除以在支持集上的概率之和），这一步至关重要——如果不做归一化，优化会快速崩溃（如消融实验所示）。此外，top-p rollout采样确保学生生成的轨迹保持在典型区域，使教师信号更可靠；特殊token masking则从实现层面消除tokenizer不匹配的虚假负反馈。

方法步骤详情

方法的完整流程如下：第一步，对于每个prompt $x$，学生推理策略$\pi_{infer}$采样一组输出$\{o_i\}_{i=1}^G$（group size $G=8$），使用top-p采样（$p=0.9$）和温度$T=1$。第二步，对于每个输出$o_i$的每个位置$t$，构建prefix $c_{i,t} = (x, y_{i,<t})$。第三步，计算教师在该prefix上的top-K支持集$S(c_{i,t}) = \text{TopK}_q(c_{i,t})$，即教师概率最高的K个token（$K=32$）。第四步，在支持集$S(c_{i,t})$上对师生分布分别做归一化：$\hat{\pi}_\theta(v|c_{i,t}) = \frac{\pi_\theta(v|c_{i,t})}{\sum_{u \in S} \pi_\theta(u|c_{i,t})}$，$\hat{q}(v|c_{i,t}) = \frac{q(v|c_{i,t})}{\sum_{u \in S} q(u|c_{i,t})}$。第五步，计算截断反向KL损失：$L_{LSM} = \mathbb{E} \left[ \frac{1}{\sum_i |o_i|} \sum_i \sum_t \sum_{v \in S} \hat{\pi}_\theta \log \frac{\hat{\pi}_\theta}{\hat{q}} \right]$。第六步，在训练中配合特殊token masking（移除EOS等特殊token的比较），使用AdamW优化器，学习率$2 \times 10^{-6}$，batch size 128，mini-batch size 64。

技术新颖性

本文的技术新颖性主要体现在四个方面。首先，它首次系统地将token-level OPD的偏差-方差权衡做了严格的理论推导——证明了token-level估计方差上界为$O(T^2)$而序列为$O(T^4)$——并据此提出'保持token-level更新以控制方差，改进局部信号以弥补偏差'的设计原则。其次，教师top-K局部支持匹配是一种新的蒸馏目标函数构造方式：它既不是全词汇KL（太贵），也不是单token log-ratio（太脆），而是在教师'认为重要'的token子集上做分布级比较，这个子集既高效又有信息量。第三，支持集归一化这一步看似简单，但在消融实验中被证明是必不可少的——没有它训练会快速崩溃——这揭示了截断KL目标的一个关键实现细节。第四，本文将三个独立的工程实践（top-K支持、top-p rollout、特殊token masking）组合为一个连贯的方案，并通过消融实验证明每个组件都有贡献，整体效果优于任何单一修改。

实验结果

本文在三个实验设置下验证了方法的有效性。在单任务数学推理中（Table 1），使用Qwen2.5-7B-Instruct作为学生、OpenThinker3-7B作为教师，标准采样token OPD将学生从28.2提升到36.4平均分，加入特殊token masking后进一步到40.7，而本文方法（不带masking）达到41.7，带masking达到41.5。这表明LSM的增益不仅来自masking，还来自更强的分布级蒸馏信号。在多任务交替训练中（Table 2，ALFWorld + 数学推理），LSM不带masking将数学平均分从34.8提升到41.7（+19.8%），同时ALFWorld成功率保持在95.3%；带masking的版本ALFWorld达到97.7%，但数学提升幅度有所下降。这说明LSM对推理任务特别有效，而masking则在任务间移了权衡点。在WebShop上的额外实验（Table A2），使用Qwen2.5-1.5B-Instruct作为学生，LSM将成功率从50.0%提升到57.8%，进一步验证了方法的泛化性。消融实验（Table 3）表明：单独使用teacher top-K（无top-p rollout）只得到17.7 avg@32（低于采样token OPD的20.4），但加上top-p后跃升到23.6，说明rollout策略对方法效果至关重要；去除归一化会导致训练快速崩溃。支持集大小K在16-48范围内不敏感，但太小会不稳定。训练动态分析（Figure A6）显示LSM产生更小的梯度范数、更少的梯度裁剪触发、更短的响应长度，以及师生对数概率差更接近零。

查看结构化数据

任务	指标	本文	基线	提升
数学推理 (Avg. 5 benchmarks)	pass@1 平均分	41.7 (w/o mask)	36.4 (sampled-token OPD)	+14.6%
AIME24	pass@1	23.3%	10.0% (sampled-token OPD)	+133%
AIME25	pass@1	26.7%	16.7% (sampled-token OPD)	+60%
多任务数学推理 (Avg.)	pass@1 平均分	41.7 (w/o mask)	34.8 (sampled-token OPD)	+19.8%
多任务 ALFWorld	成功率	97.7% (w/ mask)	90.6% (sampled-token OPD)	+7.8%
WebShop 成功率	成功率	57.8%	50.0% (sampled-token OPD)	+15.6%

局限与改进

作者在论文中坦诚地指出了多个局限性。首先，教师匹配仍然是任务成功的不完美代理——即使OPD目标定义良好，教师偏好的局部续写仍可能与成功行为偏离（reward hacking案例在附录中有详细展示）。实验中与教师的差距仍然显著（教师56.0 vs 本文41.7），说明更好的局部监督只是蒸馏问题的一部分。其次，支持集KL是一个截断目标——只在受限token子集上计算KL，相对于全词汇反向KL引入了偏差，且支持集外的token不参与梯度更新。第三，训练-推理不匹配：prefix在top-p采样的rollout策略下生成，但训练引擎不对此采样过程做修正。第四，多任务设置下的支持集定义对结果影响很大，teacher top-K、student top-K、teacher top-K + sampled token三种变体在多任务下的表现差异很大，目前缺乏理论解释。第五，本文的实验规模相对有限——主要在7B模型上验证，未在更大模型上测试；训练步数也较少（400步）。

独立分析的弱点

从独立分析的角度，本文存在几个值得关注的弱点。第一，方法的效果高度依赖top-p rollout采样——消融实验显示没有top-p时teacher top-K反而比采样token OPD更差（17.7 vs 20.4），这说明方法的鲁棒性不足，不能独立于rollout策略而稳定工作。改进方向：可以研究自适应的top-p策略，根据师生分布差异动态调整采样参数。第二，支持集大小K=32是固定的，对不同任务和序列位置可能不是最优的——在师生分布接近时K可以小一些，在分歧大时需要更大K。改进方向：用top-p截断代替top-K作为支持集选择策略（作者在Future Directions中也提到了这一点）。第三，特殊token masking是一个针对特定tokenizer的工程补丁，不具备普遍性——如果师生使用完全不同的tokenizer（不只是特殊token不同），masking无法解决问题。改进方向：研究tokenizer-agnostic的对齐方法，如在token embedding空间而非token空间做匹配。第四，实验仅在7B和1.5B模型上验证，未测试更大规模模型（如70B、405B），而这些才是OPD最有价值的应用场景。

未来方向

作者在附录B中提出了多个未来方向。首先，top-p截断作为自适应支持——用累积概率阈值代替固定K，使支持集大小自动适应分布的集中程度。其次，OPD vs RL在多任务设置下的直接对比——OPD的优化目标通过教师过滤，而RL直接读取环境奖励，理解两者在知识转移上的区别对于选择后训练策略至关重要。第三，持续学习场景——OPD可以作为记忆保持机制，但也会面临分布偏移、教师陈旧化等本文揭示的问题。此外，基于本文成果可以延伸：(1) 将LSM与EMA-anchor稳定化、扰动修正等方法结合，构建更完整的蒸馏工具箱；(2) 研究教师不确定性的作用——在师生分歧大的区域，教师的低置信度信号可能本身就不应该被信任；(3) 将LSM与结果可验证奖励（outcome-verifiable rewards）结合，弥补教师匹配与任务成功之间的gap。

复现评估

本文的复现条件较好。代码已开源（GitHub链接在论文中提供），基于verl-agent框架实现。实验在8×NVIDIA H100 GPU节点上进行，使用Qwen2.5-Instruct模型系列（7B和1.5B），这些模型都是公开可用的。训练数据使用DAPO-Math-17K的英文部分和ALFWorld环境，均为公开数据集。超参数在Table A1中完整列出：group size 8、top-p 0.9、温度1、教师top-K 32、学习率$2 \times 10^{-6}$、batch size 128、训练400步。方法本身实现简单——核心修改就是在每个prefix上取教师top-K、归一化、计算KL，不涉及复杂的架构改变。复现难度中等：主要挑战在于需要同时加载师生两个模型（7B级别），以及需要teacher logits（白盒设置）。对于无法访问teacher logits的场景，本文方法不适用。

重新审视在策略蒸馏：经验性失败模式与简单修复方案 Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes