OPRD: On-Policy Representation Distillation

Shenzhi Yang, Guangcheng Zhu, Bowen Song, Haobo Wang, Mingxuan Xia +5 more
6/4/2026
cs.LGcs.AI

No confident code implementation found for this paper.

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.

View on arXivView PDF

Code Implementations

No confident code match yet

We couldn't find an author-owned or strongly-evidenced community implementation for this paper. Any repos shown below are weak matches — verify before relying on them.

No code implementations found yet.

Know of an implementation? Let us know in the comments below!

Similar Papers

When Do We Need On-Policy Distillation? Distilling on Offline Student Rollouts Is Often Better

91% match2026

Less is More: Early Stopping Rollout for On-Policy Distillation

90% match2026

SAGE-OPD: Selective Agent-Guided Intervention for Multi-Turn On-Policy Distillation

90% match2026

LastOPD: Taming Collapse in Latent On-Policy Distillation

90% match2026

Does On-Policy Distillation Need to Complete Every Rollout?

89% match2026

DTDD: Divergence-Triggered Dynamic Distillation for Reliable On-Policy Supervision

89% match2026

Cite this paper

@article{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    = {http://arxiv.org/abs/2606.06021v1}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2606.06021.svg)](https://codeofpaper.com/paper/2606.06021)

Discussion