The Many Faces of On-Policy Distillation: Pitfalls, Mechanisms, and Fixes¶

会议: NeurIPS 2026
arXiv: 2605.11182
领域: 模型压缩
关键词: 在策略蒸馏, 在策略自蒸馏, 逆向KL散度, 特权信息内化, 强化学习验证器

一句话总结¶

系统揭示在策略蒸馏（OPD）与自蒸馏（OPSD）在数学推理中因前缀失真、有偏 TopK 逆向 KL 梯度及特权信息边缘化导致的崩溃机理，提出 Stop-gradient TopK 目标、RLVR 适配教师与 SFT 预热的联合稳定化方案。

研究背景与动机¶

在策略蒸馏（On-Policy Distillation, OPD）与在策略自蒸馏（On-Policy Self-Distillation, OPSD）作为大语言模型后训练阶段的热门技术，通过在学生策略自身采样的轨迹上提供密集的逐 token 级教师监督，展现出整合强教师能力、消除离线分布偏移与内化上下文提示的巨大潜力。然而，现有文献对其真实效用评价褒贬不一：一方面有工作报道其成功内化系统提示并提升推理能力；另一方面，近期研究频繁暴露训练发散、生成长度爆炸、重复退化乃至评测准确率断崖式下滑等严重不稳定性。

这种现象的核心矛盾在于：学界普遍假设只要将教师模型（或由特权信息增强的学生自身）挂载在学生前缀后提供逆向 KL 监督，学生就能沿着自身状态空间稳步吸纳能力。但实际上，学生生成的低质前缀会强行将教师推入未曾经历的非稳态，造成教师监督发生语义冲突并诱导频繁分支跳转；与此同时，工程上为节省显存所普遍采纳的 TopK 截断逆向 KL 散度会悄然打破全词表梯度的零和常数消去条件，引入持续抑制教师优势候选的固有偏差梯度。更为关键的是，OPSD 在数学推理等具体实例强绑定的特权信息任务中，理论最优解仅仅是跨实例特权信息的几何平均无条件策略，导致学生根本无法习得推理机制。

本文的目标是全面理清 OPD 与 OPSD 何时生效、何时崩溃及其底层驱动机理，并提供兼具理论自洽性与工程可行性的修复手段。核心 idea：通过形式化解析揭示前缀诱导的教师状态失真、未归一化 TopK 逆向 KL 的常数项梯度偏差、以及实例特定特权信息的边际共识退化三大机理，进而构建由梯度截断（Stop-gradient）TopK 代理损失、可验证奖励强化学习（RLVR）适配教师和轨迹监督微调（SFT）预热组成的三位一体稳定化配方。

方法详解¶

整体框架¶

在策略蒸馏与自蒸馏的学习范式中，输入样本 $x$ 输入学生策略 $\pi_\theta$，模型按自身分布自回归采样生成完整轨迹 $y = (y_1, \dots, y_T) \sim \pi_\theta(\cdot \mid x)$。在每个生成时间步 $t$，教师模型 $\pi_T$（OPD 中为外部强模型，OPSD 中为挂载特权信息 $I$ 的学生镜像）基于相同的历史前缀 $(x, y_{<t})$ 评估下一词预测分布，并通过单步损失 $\ell_t$ 计算逐 token 监督信号以更新学生参数 $\theta$。

整个系统的运行与稳定化流程如下所示，涵盖数据采样、教师特权条件注入、梯度截断代理损失计算以及针对三大失效机理的防御干预：

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入样本 x 与特权信息 I<br/>区分实例特定与共享潜在规则"] --> B["学生轨迹生成<br/>SFT稳定化预热防乱码与长度爆炸"]
    B --> C["教师前缀条件化评估<br/>RLVR适配对齐学生分布减少状态失真"]
    C --> D["交集词表截断过滤<br/>S_K = S_tea ∩ S_stu 解决显存膨胀"]
    D --> E["Stop-gradient TopK 逆向KL优化<br/>截断 log π_S 梯度消除 +1 常数偏差"]
    E --> F["对齐收敛模型<br/>推理时无需特权信息高效自回归输出"]

关键设计¶

1. 消除未归一化 TopK 逆向 KL 梯度偏差的截断代理目标：彻底修复优化崩溃

在全词表逆向 KL 散度中，梯度包含由于全概率求和为 1 从而内积为 0 的常数项： $$\nabla_{\theta} D_{\mathrm{KL}}(p_{\theta} \parallel p_{T}) = \sum_{v \in \mathcal{V}} p_{\theta}(v) \left[ \log \frac{p_{\theta}(v)}{p_{T}(v)} + 1 \right] \nabla_{\theta} \log p_{\theta}(v)$$ 由于 $\sum_{v} p_{\theta}(v) \nabla_{\theta} \log p_{\theta}(v) = \nabla_{\theta} \sum_{v} p_{\theta}(v) = 0$，常量 $+1$ 自然消去。然而，当显存限制迫使工程实现仅截取局部候选集 $S_K(y_{<t}) \subset \mathcal{V}$ 时，$\sum_{v \in S_K} p_{\theta}(v) < 1$，导致 $\nabla_{\theta} \sum_{v \in S_K} p_{\theta}(v) \neq 0$。残存的 $+1$ 项彻底改变了参数更新法则，仅当教师概率严格满足 $\pi_T(v) > e \cdot \pi_S(v)$ 时对应 token 才被提升；教师仅略占优势的合理 token 反遭抑制，最终诱发模型退化至全篇重复输出 "maybe"、"wait" 的死锁状态。为此，设计了 Stop-gradient TopK 目标，直接对加权对数项停止求导： $$\mathcal{L}_{\mathrm{SG\text{-}Top}K}(t) = -\sum_{v \in S_K(y_{<t})} \pi_S(v \mid x, y_{<t}) \left[ \log \pi_T(v \mid x, y_{<t}, I) - \operatorname{stopgrad}(\log \pi_S(v \mid x, y_{<t})) \right]$$ 该目标将对数比率严格视作策略梯度的标量优势度（Advantage），不仅数学上自然规避了 $+1$ 偏置项，同时保留了局部重要性加权。

2. 基于可验证奖励强化学习（RLVR）的教师域适配：化解前缀状态失真

在策略蒸馏中，外部强教师模型被迫在学生生成的中间前缀上续写。实验揭示强教师独立推理准确率虽高，但在学生截断前缀下准确率断崖式下跌 16.16 个百分点（从 62.12% 暴跌至 45.96%），教师频繁在高熵状态下倾向于生成否定与转折词汇进行支路纠偏，反而为学生注入混乱监督。针对该痛点，单纯选用更大参数量教师并不奏效；关键在于使教师具备处理学生前缀的局部相容性。采用 RLVR（如基于 DAPO 框架对 1.7B 模型执行针对性强化学习）对教师在目标任务域先验微调，生成具备鲁棒解题能力且分布特征贴近学生策略的特化教师（Qwen3-1.7B-GRPO）。该教师在 Top20 词表分布上与学生具有更高重合度，大幅减轻局部语义冲突。

3. 特权信息结构性判定准则：划定 OPSD 的适用边界

OPSD 的理论最优策略满足各特权信息条件下教师分布的几何均值归一化形式： $$p_S^\star(y \mid x) = \frac{\exp\left(\mathbb{E}_{I \sim \mathcal{D}(\cdot \mid x)} \log p_T(y \mid x, I)\right)}{\sum_{y'} \exp\left(\mathbb{E}_{I \sim \mathcal{D}(\cdot \mid x)} \log p_T(y' \mid x, I)\right)}$$ 这表明 OPSD 的本质是在无特权输入下学习所有可能特权条件下教师策略的“共识”。当特权信息 $I$ 属于实例特定（Instance-specific）信息（例如数学题的标准答案或完整推理链）时，不同样本的特权信息异质性极高，边缘化几何平均无法还原确定性的推理因果链，反而让学生退化为低置信度的无用妥协；反之，当特权信息属于共享潜在规则（Shared Latent Rule）（例如全局系统提示词、安全对齐规范、或 CharacterBench 中的特定风格偏好）时，所有样本在特权信息引导下展现相同的先验模式，学生能够极其高效地将外在提示压缩为内在的参数归纳偏置。

4. 教师轨迹预热与双端 TopK 交集显存约束：夯实系统级工程实现

针对小参数量基座模型在未对齐初始状态下容易在 OPD 初期产生不可读乱码（无序 Unicode 符号）导致教师难以介入评分的工程硬伤，引入教师轨迹 SFT 预热阶段。采用非思考模式教师生成正确规范的解答样本对学生进行轻量监督微调，显著压低前缀困惑度（PPL 从 1.896 降至 1.397），确保 rollouts 始终运行在语义良构区间。同时，针对现代推理引擎（如 SGLang）不支持位置依赖异构 TopK 提取而引发张量维度爆炸的问题，设计了双端 TopK 交集动态截断 $S_K(y_{<t}) = S_{\mathrm{tea}, K}(y_{<t}) \cap S_{\mathrm{stu}, K}(y_{<t})$，在兼顾显存开销线性化的同时保持稳定的监督保真度。

实验关键数据¶

主实验¶

论文在数学推理（OpenThoughts 训练，Math500 / AIME24 / AIME25 测试）、风格对齐（CharacterBench / EmotionBench）和系统提示词内化（DAPO-Math 推理压缩、Wildguardmix 安全对齐）多维度上全面验证。

任务场景	评测基准 / 模型设置	本文方案 (OPD/OPSD+Fixes)	对比基线 (Vanilla/GRPO/Direct)	性能表现与增益
教师前缀敏感度测试	GPQA-Diamond (Qwen3-14B Teacher)	62.12% (独立教师基线)	45.96% (学生前缀续写)	强教师因前缀失真暴跌 -16.16%
风格对齐收敛加速	CharacterBench (Qwen3-4B 学生)	OPSD (Reward=1.0 快速收敛)	GRPO / PPO (缓慢爬升)	OPSD 样本效率显著领先 RL 基线
推理长度压缩	DAPO-Math-17k (Qwen3-8B 思考模式)	OPSD (保持 Math500 精度无损)	GRPO + 长度惩罚	OPSD 实现大幅度词数压缩且样本文本更少
SFT 预热稳定性	Qwen3-4B $\to$ 1.7B-Base OPD 训练	SFT + OPD (稳定收敛, 准确率上升)	Direct OPD (输出崩溃, 精度归零)	根除乱码退化，重构稳定学习通道

消融实验¶

针对损失函数构造方式与 TopK 梯度的核心消融结果如下：

配置	优化行为与稳定性	说明
Unnormalized Top-5 Reverse KL (Eq. 6)	崩溃退化	梯度中保留 $+1$ 偏置项，第 700 步长度爆炸，第 1000 步退化为单一词 "maybe"
Renormalized Top-5 Reverse KL (Eq. 15)	稳定收敛	候选集内部重归一化消去 $+1$ 偏置项，但忽略了集外概率质量
Stop-gradient Top-5 Reverse KL (Eq. 12)	稳定收敛 (最优)	显式截断 $\log \pi_S$ 梯度，完全消除偏差，且与策略梯度优势视角自洽
Sampled-token in Policy Gradient (Eq. 4)	稳定运行	仅依赖抽样 token 标量差，无 TopK 偏差，但无法捕捉全词表分布
Sampled-token as Single Loss (TopK=1)	崩溃退化	单 token 逆向 KL 丧失归一化约束，陷入局部发散

关键发现¶

教师与学生能力差距的悖论：单纯依赖绝对能力更强的教师（Qwen3-8B）进行数学 OPD，效果反而显著落后于经 RLVR 适配后的同规模教师（Qwen3-1.7B-GRPO），证实教师与学生词表分布的局部对齐度比教师基准天花板更具决定性。
特权信息的双刃剑效应：在数学推理中，提供完整解答（full-response）作为特权信息导致的蒸馏损失与退化程度，比仅提供最终答案（final-answer）更为严重；更丰富的实例特定特权信息反而急剧放大了分布漂移。
逆向 KL 梯度的数学陷阱：简单地将截断候选词数 $K$ 从 5 扩大到 20 根本无法挽救未归一化目标的坍塌命运，只有通过重归一化或截断求导彻底抹除常数偏置项，才能确保收敛。

亮点与洞察¶

透彻解构了 TopK 逆向 KL 的优化病态：指出在连续域和全离散域成立的概率和守恒常数消除定理，在局部截断集合中失效，从数学底层解释了语言模型在策略蒸馏中长期困扰学术界的“重复词死锁”现象。
确立了特权信息内化的拓扑判别准则：提出“实例特定变量”与“共享潜在流形”的二元分类标准，明晰了自蒸馏方法在复杂推理任务中的内在理论缺陷，避免了后续研究盲目在推理题中堆砌提示的试错成本。
揭示了前缀条件化对强模型的降智效应：实证证明将大模型作为下游监督信号时，学生的前缀起到了“对抗性干扰上下文”的作用，为后续构建解耦前缀纠偏或回溯蒸馏架构提供了关键启示。

局限与展望¶

作者指出的局限：实验主要覆盖 Qwen 系列架构与 1.7B 至 14B 参数区间，在更大规模开源模型或不同分词器体系下的普遍性仍待系统评估。
自身发现的局限：对于前缀失真问题，当前方案主要依赖 RLVR 适配和 SFT 初始化进行被动防御，尚未提出在蒸馏损失内自适应衰减低置信度前缀权重的主动动态机制。
未来方向：探索 SFT 预热、强化学习探索与在策略蒸馏三者交替迭代的闭环自演进流水线（Iterative Self-Improvement Pipeline），使模型在自身生成策略稳步提升的同时吸纳高质量稠密梯度。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 深入数学机理剖析了 TopK 逆向 KL 梯度偏差与 OPSD 边缘化极限，视角极其犀利深刻。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖数学、推理压缩、风格对齐、安全等多元场景，理论推导与失败轨迹可视化兼备。
写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严谨自洽，对实验现象的归因论证层层递进，洞察力极强。
价值: ⭐⭐⭐⭐⭐ 为大语言模型后训练时代的在策略蒸馏实践排除了重大技术暗坑，具有不可替代的工程指导价值。

配置	优化行为与稳定性	说明
Unnormalized Top-5 Reverse KL (Eq. 6)	崩溃退化	梯度中保留 \(+1\) 偏置项，第 700 步长度爆炸，第 1000 步退化为单一词 "maybe"
Renormalized Top-5 Reverse KL (Eq. 15)	稳定收敛	候选集内部重归一化消去 \(+1\) 偏置项，但忽略了集外概率质量
Stop-gradient Top-5 Reverse KL (Eq. 12)	稳定收敛 (最优)	显式截断 \(\log \pi_S\) 梯度，完全消除偏差，且与策略梯度优势视角自洽
Sampled-token in Policy Gradient (Eq. 4)	稳定运行	仅依赖抽样 token 标量差，无 TopK 偏差，但无法捕捉全词表分布
Sampled-token as Single Loss (TopK=1)	崩溃退化	单 token 逆向 KL 丧失归一化约束，陷入局部发散