让大模型蒸馏更高效:用隐藏层对齐替代输出层匹配
OPRD: On-Policy Representation Distillation

- 在相同采样路径上对齐师生模型的隐藏层表示,实现密集确定性监督
- 在数学竞赛题上超越所有输出空间基线,训练快1.44倍,内存少54%
- 支持跨架构和跨分词器的知识迁移,突破传统词汇依赖限制
在策略蒸馏(OPD)中,学生模型仅通过匹配下一个词分布进行监督,存在两个缺陷:(i) 当学生接近教师时梯度信号噪声比急剧下降;(ii) 语言模型头部造成中间隐藏状态信息丢失。本文提出在策略表示蒸馏(OPRD),首次将蒸馏从输出空间扩展到隐藏状态空间。OPRD在同一批采样轨迹上对选定层的学生与教师表示进行对齐,提供密集、确定性的逐层监督,完全绕过语言模型头部。理论上,OPRD给出确定性的样本级梯度,消除原有输出空间蒸馏中的令牌级估计方差,并暴露输出目标必然丢弃的结构信息。实验上,OPRD在竞赛数学基准(AIME 2024、AIME 2025、AIMO)上显著缩小师生差距,所有输出空间基线均未达到教师性能,同时训练速度提升1.44倍,内存消耗最多降低54%。进一步提出OPRD-Bridge,利用异构模型共享低秩表示结构的特性,构建冻结投影对,实现任意深度与宽度差异下的表示对齐,将对齐从依赖共同词汇表的输出空间转移到表示空间。在跨架构(Qwen3-4B → Qwen3-1.7B-Base)和跨分词器(Phi-4-mini-reasoning → Qwen3-1.7B-Base)场景下验证成功知识迁移,即使无词汇对齐通道亦可有效传递知识。
原文摘要 · Abstract (English)
On-policy distillation (OPD) supervises the student exclusively in the output space by matching next-token distributions. This paradigm suffers from two limitations: (i) a high-variance gradient estimator whose signal-to-noise ratio collapses as the student approaches the teacher, and (ii) an LM-head information bottleneck that discards the teacher's intermediate hidden states. We propose On-Policy Representation Distillation (OPRD), the first method to lift on-policy distillation into the hidden-state space. OPRD aligns student and teacher representations across selected layers on the same on-policy rollouts, providing dense, deterministic, per-layer supervision while bypassing the LM head entirely. Theoretically, OPRD provides a deterministic per-sample gradient, removing the token-level estimation variance that plagues OPD, and exposes structural information that any output-space objective necessarily discards. Empirically, OPRD closes the student-teacher gap on competition mathematics benchmarks (AIME 2024, AIME 2025, and AIMO), where every output-space baseline plateaus below the teacher, while training 1.44x faster and using up to 54% less memory. We further extend OPRD to the cross-architecture setting via OPRD-Bridge. By exploiting the observation that heterogeneous models share a low-rank representational structure, we construct a frozen projector pair that aligns representations across arbitrary depth and width mismatches, shifting the alignment from the output space (which depends on a shared vocabulary) to the representation space. We validate OPRD-Bridge on both cross-architecture (Qwen3-4B -> Qwen3-1.7B-Base) and cross-tokenizer (Phi-4-mini-reasoning -> Qwen3-1.7B-Base) settings, demonstrating successful knowledge transfer even when the vocabulary-based alignment channel is unavailable. Code: https://github.com/ShenzhiYang2000/OPRD.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。