OPRD: On-Policy Representation Distillation

📄 arXiv: 2606.06021v1 📥 PDF

作者: Shenzhi Yang, Guangcheng Zhu, Bowen Song, Haobo Wang, Mingxuan Xia, Xing Zheng, Yingfan Ma, Zhongqi Chen, Weiqiang Wang, Gang Chen

分类: cs.LG, cs.AI

发布日期: 2026-06-04

🔗 代码/项目: GITHUB

💡 一句话要点

提出OPRD以解决现有蒸馏方法的局限性

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 蒸馏训练 隐藏状态对齐 自然语言处理 模型压缩 深度学习

📋 核心要点

  1. 现有的基于输出的蒸馏方法存在采样方差和忽略隐藏状态的问题，限制了模型性能的提升。
  2. 本文提出OPRD，通过在隐藏状态空间中对齐学生和教师的表示，绕过语言模型头，提升了蒸馏效果。
  3. 实验结果显示，OPRD在多个基准上超越了传统方法，训练速度提高1.44倍，内存使用减少54%。

📝 摘要（中文）

在现有的基于输出的蒸馏方法中，学生模型仅通过匹配下一个标记的概率进行监督，这种方法存在两个主要限制：一是由于大词汇表（如Qwen的约15万个标记）导致的采样方差在训练过程中持续存在，二是将教师模型视为黑箱，忽略了其隐藏状态的信息。为此，本文提出了基于隐藏状态的蒸馏方法OPRD，通过在同一回合中对学生和教师的表示进行对齐，绕过了语言模型头，从而消除了采样方差并提供了更丰富的结构信息。实验证明，OPRD在AIME 2024/2025和AIMO上缩小了学生与教师之间的差距，同时训练速度提高了1.44倍，内存使用减少了54%。

🔬 方法详解

问题定义：现有的基于输出的蒸馏方法（OPD）仅通过匹配下一个标记的概率进行监督，导致采样方差和对教师模型隐藏状态信息的忽视，这限制了模型的学习效果和性能提升。

核心思路：OPRD通过在隐藏状态空间中对齐学生和教师的表示，避免了对语言模型头的依赖，从而消除了采样方差，并提供了更丰富的结构信息，增强了模型的学习能力。

技术框架：OPRD的整体架构包括教师模型和学生模型的并行训练，选择特定层进行表示对齐，利用相同的回合数据进行监督。主要模块包括隐藏状态对齐模块和损失计算模块。

关键创新：OPRD的核心创新在于将蒸馏过程提升到隐藏状态空间，通过对齐不同层的表示来获取更丰富的信息，这与传统的仅依赖输出的蒸馏方法本质上不同。

关键设计：在设计上，OPRD采用了特定层的选择策略，损失函数设计为对齐不同层的表示，同时优化了模型的训练过程以提高效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，OPRD在AIME 2024/2025和AIMO基准上显著缩小了学生与教师模型之间的性能差距，同时训练速度提高了1.44倍，内存使用减少了54%，展现出优越的效率和效果。

🎯 应用场景

OPRD的研究成果在自然语言处理、对话系统和文本生成等领域具有广泛的应用潜力。通过提升模型的蒸馏效果，能够在资源受限的环境中实现更高效的模型部署，进而推动智能助手和自动化内容生成等技术的发展。

📄 摘要（原文）

On-policy distillation (OPD) supervises the student only in output space by matching next-token probabilities. This output-only paradigm has two limits: (1) sampling variance from Monte Carlo KL estimates over large vocabularies (e.g., Qwen's ~150k tokens) persists throughout training, and (2) it treats the teacher as a black-box, discarding all intermediate hidden states after the LM head. We propose On-Policy Representation Distillation (OPRD), which lifts distillation into hidden-state space by aligning student and teacher representations across selected layers on the same rollouts, bypassing the LM head entirely. Theoretically, OPRD eliminates sampling variance and provides richer per-layer structural information. Empirically, OPRD closes the student-teacher gap on AIME 2024/2025 and AIMO, while output-space OPD baselines plateau below the teacher. OPRD also trains 1.44x faster and uses 54% less memory than top-k OPD. Code: https://github.com/ShenzhiYang2000/OPRD.