
浙大蚂蚁团队提出OPRD把蒸馏从抄答案升级到抄思路——在隐藏状态空间监督学生绕过LM-head信息瓶颈实现零方差梯度、1.44倍训练加速、54%内存削减在AIME数学推理上首次让1.5B学生逼近教师水平。1. 蒸馏的困境只抄答案永远抄不像大模型蒸馏Distillation是老生常谈。让小模型学大模型的本事——这个思路听起来简单但做了十年瓶颈始终卡在同一个地方所有方法都在输出空间折腾。无论是最早的Hinton蒸馏soft targets还是最新的On-Policy DistillationOPD让学生自己采样答案然后对比教师的概率分布本质都一样比较学生和教师在next-token概率上的差异。浙大和蚂蚁团队的研究OPRD: On-Policy Representation Distillation指出这种输出空间-only范式有两大致命伤1.1 方差灾难后期训练信号被噪声淹没OPD的核心操作是让学生采样一个tokeny^tŷ_ty^t然后算logpt(y^t)−logqt(y^t)log p_t(ŷ_t) - log q_t(ŷ_t)logpt(y^t)−logqt(y^t)。这是单样本Monte Carlo估计KL divergence。问题当学生逐渐接近教师pt→qtp_t → q_tpt→qt信号趋近于零但方差不变。信噪比SNR在后期训练中崩溃导致精度plateau或振荡——无论你训练多久都无法突破那堵方差墙。更糟的是现代LLM词汇表巨大Qwen系列≈150K tokens方差问题被进一步放大。1.2 信息瓶颈教师只用了1%的脑容量输出空间蒸馏把教师当作黑盒概率oracle——只查询LM head之后的输出分布把整个中间层计算栈L层×d维隐藏状态当作垃圾扔掉。但这里有个数学陷阱输出分布任意接近的隐藏状态可能沿整个仿射子空间差异巨大。因为softmax对加性常数不变LM head的投影矩阵Whead∈R∣V∣×dW_head ∈ R^{|V|×d}Whead∈R∣V∣×d存在有效零空间null space——隐藏状态沿零空间方向的偏差完全不可被输出空间检测但可能代表完全不同的内部认知状态。换言之学生可能学会了鹦鹉学舌输出分布和教师一样但内部的思考过程完全不同——这在复杂推理任务数学、代码中是致命的。2. OPRD从抄答案到抄思路OPRD的核心创新极其简洁不要只比较输出概率直接比较学生和教师的中间层隐藏状态。2.1 损失函数MSE在隐藏状态空间LOPRDEx,y^[1∣Llayer∣∑l1∑mt∑tmt1d∥hθ,t(l)−sg(hT,t(l))∥22]\mathcal{L}_{\text{OPRD}} \mathbb{E}_{x, \hat{y}} \left[ \frac{1}{|L_{layer}|} \sum_{l} \frac{1}{\sum m_t} \sum_{t} m_t \frac{1}{d} \left\| h^{(l)}_{\theta,t} - \text{sg}\left(h^{(l)}_{T,t}\right) \right\|_2^2 \right]LOPRDEx,y^[∣Llayer∣1l∑∑mt1t∑mtd1hθ,t(l)−sg(hT,t(l))22]关键设计公式6组件含义典型设置L_layer蒸馏层集合全部28层P(ŷ)监督位置最后k2000个token答案收敛段m_t位置掩码1[t ∈ P(ŷ)]sg(·)stop-gradient教师冻结d隐藏维度1536为什么监督最后2000个token论文通过余弦相似度分析发现学生与教师的表示分歧集中在响应尾部chain-of-thought收敛到最终答案处首段几乎始终接近教师97%相似度。这是数据驱动的位置选择而非人工设计。2.2 与OPD的组合不是替代是互补LLOPDμ⋅LOPRD\mathcal{L} \mathcal{L}_{\text{OPD}} \mu \cdot \mathcal{L}_{\text{OPRD}}LLOPDμ⋅LOPRD两者共享同一on-policy rollout和单次教师前向传播——基础设施成本几乎为零。μ0是纯OPDμ1是平衡组合μ10是OPRD主导。μAIME24相对μ0提升042.3-147.75.41050.27.9单调提升验证了隐藏状态信号与输出空间信号的可加性——它们捕获的是不同的、互补的信息。3. 理论双杀零方差 信息瓶颈突破3.1 Theorem 1零方差梯度OPRD的梯度是确定性的——给定一个rollout损失是固定的MSE没有采样随机性。OPD的梯度是高方差的——即使给定rolloutlogpt(y^t)−logqt(y^t)log p_t(ŷ_t) - log q_t(ŷ_t)logpt(y^t)−logqt(y^t)的估计方差不随p→q消失因为score function项∇θlogpt(y^t)∇_θ log p_t(ŷ_t)∇θlogpt(y^t)始终引入噪声。后果OPD后期训练信噪比崩溃OPRD持续稳定优化。这解释了为什么Figure 3中OPD在中期plateau而OPRD单调上升至接近教师水平。3.2 Theorem 2LM-head信息瓶颈的量化设W_head的奇异值为σ1≥...≥σd0σ_1 ≥ ... ≥ σ_d 0σ1≥...≥σd0。核心结论输出空间不可检测的隐藏状态差异如果hθ−hT∈NWh_θ - h_T ∈ N_Whθ−hT∈NWLM head的零空间则输出损失ℓout0ℓ_out 0ℓout0即隐藏状态差异再大输出分布也完全一样。低奇异值方向的放大效应沿最小奇异值方向vdv_dvd隐藏状态范数与输出损失之比下界为条件数平方(σ1/σd)2(σ_1/σ_d)^2(σ1/σd)2。生产LLM中这个比值通常极大意味着隐藏状态可以偏差数个数量级而不影响输出损失。结论输出空间OPD对中间层隐藏状态没有任何约束能力。OPRD恰好惩罚这些不可检测的方向并监督任意子集的中间层。4. 实验又快、又省、又准4.1 模型与数据项目教师学生模型JustRL-Deepseek-1.5BDeepSeek-R1-Distill-Qwen-1.5B骨干Qwen2.5-1.5BQwen2.5-1.5B层数2828隐藏维度15361536词汇表≈151K≈151K训练数据DAPO-Math-17K数学推理prompts每prompt采样2个responses温度1.0最大长度16,384 tokens训练8×A100 (80G)FSDP500优化器步评估AIME 2024/2025、AIMOAMC 20222023Avg164.2 准确率逼近教师方法AIME24AIME25AIMOTeacher50.835.679.5Student (未修改)32.921.962.2OPD top-142.333.577.0OPD top-1647.134.076.5OPRD (ours)49.834.679.1关键发现OPRD与教师差距1.0 / 1.0 / 0.4点AIMO在评估噪声范围内视为effectively tiedOPD top-16严格信息超集于top-1也无法避免plateau证实Theorem 1——额外输出层信息无法抵消采样噪声训练动态OPD在中期达到plateauOPRD单调提升至接近教师4.3 效率Pareto三杀指标OPD top-1OPD top-16OPRD峰值GPU内存(GB)30.245.020.5500步训练时间(min)813812563AIME24准确率42.347.149.81.44×训练加速因为绕过LM head无需materialize B×T×|V| logits张量32-54%峰值内存削减OPD top-16需要构造top-k logits矩阵内存开销巨大同时达到更高准确率严格Pareto dominant4.4 响应更简洁方法收敛平均长度OPRD~5,700 tokensOPD~7,000 tokensOPRD在更高准确率同时产生更简洁的推理链进一步降低推理成本。5. Mechanistic分析Phase Transition假说5.1 Loss Spike现象所有OPDOPRD组合运行均出现loss spikeFigure 8推测为策略重组的phase transition。关键观察添加OPRD使spike提前到来μ1和μ10早于μ0spike后PG loss≈0但准确率差距持续存在5.4/7.9点这直接证实Theorem 2一旦策略梯度消失pt≈qtp_t ≈ q_tpt≈qt输出空间信号无法驱动进一步改进剩余差距存在于LM head的null space中——只有OPRD的表示级信号能继续优化。5.2 Top-16重叠率的Dip-Surge模式∣top−16(πθ)∩top−16(πT)∣/16|top-16(π_θ) ∩ top-16(π_T)| / 16∣top−16(πθ)∩top−16(πT)∣/16在OPRD运行中出现先dip后surgedip与PG-loss spike时间重合表示学生策略正在重组surge重组后超越纯OPD基线这支持phase transition假说——学生不是渐进式接近教师而是经历一次内部重组后跃迁到更高质量策略。6. 与相关工作的对比不是BERT蒸馏的翻版6.1 与特征蒸馏FitNets、TinyBERT、MiniLM的区别维度FitNets/TinyBERT/MiniLMOPRD监督数据固定预训练/下游语料学生生成的rollouts暴露偏置存在学生不生成自己的序列消除on-policy模型类型编码器BERT、CNN自回归解码器LLM表示特性一次性计算条件于整个采样前缀核心区别OPRD的隐藏状态对齐发生在学生自己的采样分布上每个ht(l)h_t^(l)ht(l)编码了在已生成前缀下对下一个token的预测信念。这是encoder蒸馏完全没有的on-policy对象。6.2 与输出空间蒸馏的对比维度OPD所有变体OPRD监督空间输出logits隐藏状态梯度方差高REINFORCE零确定性MSE教师信息利用仅最终分布全部中间层内存开销O(BT|V|)O(BTd)瓶颈突破无绕过LM head零空间7. 局限与未来7.1 当前局限同构假设实验要求教师和学生共享相同架构和维度无需投影器W。跨架构蒸馏如教师7B→学生1.5B不同维度需要额外验证。位置选择启发式last-k2000是基于cosine相似度分析的数据驱动选择但不同任务代码、多轮对话的最佳suffix长度可能不同。层数选择默认使用全部层但哪些层对蒸馏最关键是否可以只监督关键层来进一步加速推理模型特殊性实验基于数学推理的CoT数据在通用对话、创意生成等非结构化任务上是否同样有效7.2 未来方向跨架构蒸馏引入可学习投影器W实现不同维度/架构之间的表示对齐。动态层/位置选择基于训练进展自适应调整监督层数和位置而非固定配置。与强化学习的组合OPRD提供确定性表示信号与PPO/GRPO等强化学习结合可能实现更高效的推理能力迁移。多模态扩展在视觉-语言模型中隐藏状态空间是否包含视觉和语言模态的联合表示跨模态蒸馏是否适用8. 结论蒸馏的范式升级OPRD不是又一个蒸馏技巧而是对蒸馏范式本身的升级。它证明了三个核心命题输出空间不是唯一的监督通道——隐藏状态空间包含输出空间不可检测的结构性信息零方差梯度在LLM蒸馏中是可行的——确定性MSE损失替代了高方差的REINFORCE估计绕过LM head可以带来效率与质量的双重收益——1.44×加速、54%内存削减、同时逼近教师水平在LLM后训练post-training成为工业标配的今天OPRD为如何更高效地从强模型学习提供了一个全新的操作平面。对于那些正在用OPD蒸馏自家模型的团队——是时候看看教师的脑内活动了。“我们证明隐藏状态监督不仅是一个更丰富的信号源更是突破LM head信息瓶颈的必要条件。当输出空间的信号耗尽表示空间的优化才刚刚开始。”—— OPRD作者团队参考论文Yang, S., Zhu, G., Song, B., Wang, H., Xia, M., Zheng, X., Ma, Y., Chen, Z., Wang, W., Chen, G. (2026).OPRD: On-Policy Representation Distillation. Zhejiang University, Ant Group. arXiv:2606.06021.代码https://github.com/ShenzhiYang2000/OPRD#大模型蒸馏 #知识蒸馏 #隐藏状态监督 #推理能力迁移 #LLM后训练 #AIME #数学推理 #浙江大学 #蚂蚁集团 #On-Policy #表示学习 #信息瓶颈 #零方差梯度 #Transformer