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.

翻译：在策略蒸馏（OPD）仅通过在输出空间中匹配下一词元概率来监督学生模型。这种仅输出范式存在两个局限：(1) 蒙特卡洛KL估计在大词表（如Qwen约15万词元）上的采样方差在整个训练过程中持续存在；(2) 它将教师模型视为黑箱，丢弃了语言模型头部之后的全部中间隐藏状态。我们提出在策略表示蒸馏（OPRD），通过在同一轨迹上对齐学生与教师选定层的表示，将蒸馏提升至隐藏状态空间，完全绕过语言模型头部。理论上，OPRD消除了采样方差，并提供了更丰富的逐层结构信息。实验上，OPRD在AIME 2024/2025和AIMO上缩小了学生与教师的差距，而输出空间OPD基线在低于教师水平处停滞。OPRD的训练速度比Top-k OPD快1.44倍，内存使用减少54%。代码：https://github.com/ShenzhiYang2000/OPRD。