OPRD: On-Policy Representation Distillation

Hugging Face Daily Papers Papers
knowledge-distillation on-policy-learning representation-learning large-language-models reinforcement-learning reasoning

Summary

OPRD proposes a new knowledge distillation method that aligns student and teacher hidden states across layers during on-policy rollouts, eliminating sampling variance from token-space KL estimation. Empirically, OPRD outperforms output-space baselines on math reasoning benchmarks (AIME 2024/2025, AIMO) while being 1.44x faster and using 54% less memory.

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.
Original Article

Similar Articles

Spend Teacher Tokens Where They Matter: Success-Referenced On-Policy Distillation

arXiv cs.AI

The paper introduces Success-Referenced On-Policy Distillation (SR-OPD), which uses a student's own successful rollouts as references to selectively allocate teacher supervision, achieving comparable reasoning performance on math benchmarks using only 3.46–5.02% of the teacher-input tokens required by vanilla on-policy distillation.

OPMD: On-Policy Mechanism Distillation Beyond Output Supervision

Hugging Face Daily Papers

The paper introduces On-Policy Mechanism Distillation (OPMD), which transfers teacher intermediate computations compatible with the student's architecture during on-policy distillation, outperforming output-only distillation baselines across multiple teacher-student settings.

Latent-MOPD: Latent Multi-Teacher On-Policy Distillation

arXiv cs.LG

The paper introduces Latent-MOPD, the first representation-level multi-teacher on-policy distillation method for LLMs, which integrates multiple specialist teachers through both their output distributions and hidden states without extra teacher training, outperforming token-only, representation-only, and averaging baselines across nine math, code, and logic benchmarks.

Understanding Off- vs On-Policy Distillation: A Tale of Distinct Training Objectives

arXiv cs.LG

This UCLA paper theoretically analyzes on-policy distillation (OPD) versus off-policy/SFT-style distillation, showing how forward and reverse KL objectives aggregate multiple teachers into arithmetic vs geometric mixtures, establishing logarithmic regret bounds, and explaining both the benefits and fragility of OPD (e.g., sensitivity to low-probability teachers and incorrect-prefix bias).

On-Policy Delta Distillation

Hugging Face Daily Papers

The paper introduces On-Policy Delta Distillation (OPD^2), a new distillation reward called the delta signal that captures the difference between a teacher model and its base model before reasoning tuning, providing a more direct signal for transferring reasoning capabilities. Experiments across math, science, and code benchmarks show OPD^2 consistently outperforms conventional on-policy distillation.