| title | OPRD: On-Policy Representation Distillation |
|---|---|
| authors | Shenzhi Yang, Guangcheng Zhu, Bowen Song, Haobo Wang, Mingxuan Xia, Xing Zheng, Yingfan Ma, Zhongqi Chen, Weiqiang Wang, Gang Chen |
| date | 2026-06-04 |
| venue | arXiv (preprint, not peer reviewed) |
| id | arXiv:2606.06021 |
| upvotes | 8 (Hugging Face Daily Papers, 2026-11-01) |
| links | abs pdf hf code |
| 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. |
cite
@misc{yang2026oprd,
title={OPRD: On-Policy Representation Distillation},
author={Shenzhi Yang and Guangcheng Zhu and Bowen Song and Haobo Wang and Mingxuan Xia and Xing Zheng and Yingfan Ma and Zhongqi Chen and Weiqiang Wang and Gang Chen},
year={2026},
eprint={2606.06021},
archivePrefix={arXiv},
url={https://arxiv.org/abs/2606.06021}
}