aiqbench analysis
OPRD addresses critical bottlenecks in on-policy distillation by shifting supervision from output probabilities to hidden-state alignment, effectively eliminating sampling variance in large-vocabulary models. This approach not only closes the performance gap on complex benchmarks like AIME but also delivers significant efficiency gains, training 1.44x faster with 54% less memory usage. By bypassing the LM head to utilize intermediate layer representations, the method offers a theoretically sound and practically superior alternative to existing output-space paradigms.
OPRD: On-Policy Representation Distillation
arXiv:2606.06021v3 Announce Type: replace Abstract: 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.
Abstract: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: this https URL.
Machine Learning (cs.LG); Artificial Intelligence (cs.AI)
arXiv:2606.06021 [cs.LG]
(or arXiv:2606.06021v3 [cs.LG] for this version)
Full text under the publisher's open licence · 147 words
Related coverage
OPRD: On-Policy Representation Distillation
arXiv · Jun 9, 2026
OPRD: On-Policy Representation Distillation
arXiv · Jun 6, 2026
Eliciting Weak-to-Strong Generalization with On-Policy Reverse Distillation
arXiv · Sep 9, 2026
OPMD: On-Policy Mechanism Distillation Beyond Output Supervision
arXiv · Oct 15, 2026