OPRD：蒸馏不只学答案，还要偷看老师的“脑内活动“
浙大+蚂蚁团队提出OPRD：把蒸馏从"抄答案"升级到"抄思路"——在隐藏状态空间监督学生，绕过LM-head信息瓶颈，实现零方差梯度、1.44倍训练加速、54%内存削减，在AIME数学推理上首次让1.5B学生逼近教师水平。
1. 蒸馏的困境：只抄答案，永远抄不像
大模型蒸馏（Distillation）是老生常谈。让小模型学大模型的本事——这个思路听起来简单，但做了十年，瓶颈始终卡在同一个地方：
所有方法都在输出空间折腾。
无论是最早的Hinton蒸馏（soft targets），还是最新的On-Policy Distillation（OPD，让学生自己采样答案，然后对比教师的概率分布），本质都一样：比较学生和教师在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在隐藏状态空间
LOPRD=Ex,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]LOPRD=Ex,y^[∣Llayer∣1l∑∑mt1t∑mtd1hθ,t(l)−sg(hT,t(l))22]
关键设计（公式6）：
| 组件 | 含义 | 典型设置 |
|---|---|---|
| L_layer | 蒸馏层集合 | 全部28层 |
| P(ŷ) | 监督位置 | 最后k=2000个token（答案收敛段） |
| m_t | 位置掩码 | 1[t ∈ P(ŷ)] |
| sg(·) | stop-gradient | 教师冻结 |
| d | 隐藏维度 | 1536 |
为什么监督最后2000个token？论文通过余弦相似度分析发现：学生与教师的表示分歧集中在响应尾部（chain-of-thought收敛到最终答案处），首段几乎始终接近教师（97%+相似度）。这是数据驱动的位置选择，而非人工设计。
2.2 与OPD的组合：不是替代，是互补
L=LOPD+μ⋅LOPRD\mathcal{L} = \mathcal{L}_{\text{OPD}} + \mu \cdot \mathcal{L}_{\text{OPRD}}L=LOPD+μ⋅LOPRD
两者共享同一on-policy rollout和单次教师前向传播——基础设施成本几乎为零。μ=0是纯OPD，μ=1是平衡组合，μ=10是OPRD主导。
| μ | AIME24 | 相对μ=0提升 |
|---|---|---|
| 0 | 42.3 | - |
| 1 | 47.7 | +5.4 |
| 10 | 50.2 | +7.9 |
单调提升验证了隐藏状态信号与输出空间信号的可加性——它们捕获的是不同的、互补的信息。
3. 理论双杀：零方差 + 信息瓶颈突破
3.1 Theorem 1：零方差梯度
OPRD的梯度是确定性的——给定一个rollout，损失是固定的MSE，没有采样随机性。
OPD的梯度是高方差的——即使给定rollout，logpt(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 2：LM-head信息瓶颈的量化
设W_head的奇异值为σ1≥...≥σd>0σ_1 ≥ ... ≥ σ_d > 0σ1≥...≥σd>0。
核心结论：
输出空间不可检测的隐藏状态差异：如果hθ−hT∈NWh_θ - h_T ∈ N_Whθ−hT∈NW（LM head的零空间），则输出损失ℓout=0ℓ_out = 0ℓout=0，即隐藏状态差异再大，输出分布也完全一样。
低奇异值方向的放大效应：沿最小奇异值方向vdv_dvd，隐藏状态范数与输出损失之比下界为条件数平方(σ1/σd)2(σ_1/σ_d)^2(σ1/σd)2。生产LLM中这个比值通常极大，意味着隐藏状态可以偏差数个数量级而不影响输出损失。
结论：输出空间OPD对中间层隐藏状态没有任何约束能力。OPRD恰好惩罚这些不可检测的方向，并监督任意子集的中间层。
4. 实验：又快、又省、又准
4.1 模型与数据
| 项目 | 教师 | 学生 |
|---|---|---|
| 模型 | JustRL-Deepseek-1.5B | DeepSeek-R1-Distill-Qwen-1.5B |
| 骨干 | Qwen2.5-1.5B | Qwen2.5-1.5B |
| 层数 | 28 | 28 |
| 隐藏维度 | 1536 | 1536 |
| 词汇表 | ≈151K | ≈151K |
- 训练数据：DAPO-Math-17K（数学推理prompts）
- 每prompt采样2个responses，温度1.0，最大长度16,384 tokens
- 训练：8×A100 (80G)，FSDP，500优化器步
- 评估：AIME 2024/2025、AIMO（AMC 2022+2023），Avg@16
4.2 准确率：逼近教师
| 方法 | AIME24 | AIME25 | AIMO |
|---|---|---|---|
| Teacher | 50.8 | 35.6 | 79.5 |
| Student (未修改) | 32.9 | 21.9 | 62.2 |
| OPD top-1 | 42.3 | 33.5 | 77.0 |
| OPD top-16 | 47.1 | 34.0 | 76.5 |
| OPRD (ours) | 49.8 | 34.6 | 79.1 |
关键发现：
- OPRD与教师差距：1.0 / 1.0 / 0.4点（AIMO在评估噪声范围内，视为effectively tied）
- OPD top-16（严格信息超集于top-1）也无法避免plateau，证实Theorem 1——额外输出层信息无法抵消采样噪声
- 训练动态：OPD在中期达到plateau，OPRD单调提升至接近教师
4.3 效率：Pareto三杀
| 指标 | OPD top-1 | OPD top-16 | OPRD |
|---|---|---|---|
| 峰值GPU内存(GB) | 30.2 | 45.0 | 20.5 |
| 500步训练时间(min) | 813 | 812 | 563 |
| AIME24准确率 | 42.3 | 47.1 | 49.8 |
- 1.44×训练加速（因为绕过LM head，无需materialize B×T×|V| logits张量）
- 32-54%峰值内存削减（OPD top-16需要构造top-k logits矩阵，内存开销巨大）
- 同时达到更高准确率（严格Pareto dominant）
4.4 响应更简洁
| 方法 | 收敛平均长度 |
|---|---|
| OPRD | ~5,700 tokens |
| OPD | ~7,000 tokens |
OPRD在更高准确率同时产生更简洁的推理链，进一步降低推理成本。
5. Mechanistic分析：Phase Transition假说
5.1 Loss Spike现象
所有OPD+OPRD组合运行均出现loss spike（Figure 8），推测为策略重组的phase transition。关键观察：
- 添加OPRD使spike提前到来（μ=1和μ=10早于μ=0）
- spike后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后surge：
- dip：与PG-loss spike时间重合，表示学生策略正在重组
- surge：重组后超越纯OPD基线
这支持"phase transition"假说——学生不是渐进式接近教师，而是经历一次"内部重组"后跃迁到更高质量策略。
6. 与相关工作的对比：不是BERT蒸馏的翻版
6.1 与特征蒸馏（FitNets、TinyBERT、MiniLM）的区别
| 维度 | FitNets/TinyBERT/MiniLM | OPRD |
|---|---|---|
| 监督数据 | 固定预训练/下游语料 | 学生生成的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-k=2000是基于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