OPRD: On-Policy Representation Distillation Moves Beyond Output-Space Knowledge Distillation

Forum topic · 小凯 · 2026-06-07

Summary

Researchers from Zhejiang University and Ant Group propose OPRD (On-Policy Representation Distillation), a method that supervises a student model in the hidden-state space of a teacher rather than only on next-token output distributions. The authors identify two fundamental flaws of output-space distillation such as On-Policy Distillation (OPD): single-sample Monte Carlo estimates of KL divergence suffer variance collapse as the student approaches the teacher, and the LM head projection has an effective null space, so hidden states can differ drastically while producing identical outputs. OPRD adds a deterministic MSE loss that aligns student and teacher hidden states on the student's own sampled rollouts, focusing on the last ~2000 tokens where representational divergence concentrates. On AIME 2024/2025 and AIMO math reasoning benchmarks using a 1.5B Qwen-based student, OPRD reaches 49.8 on AIME24 versus a teacher score of 50.8, achieving a 1.44x training speedup and 32-54% peak GPU memory reduction compared to OPD, while producing shorter reasoning chains. Combined OPD+OPRD training shows additive gains and a phase-transition-like loss spike, supporting the claim that representation-level supervision unlocks improvements invisible in output space.

OPRD: On-Policy Representation Distillation Moves Beyond Output-Space Knowledge Distillation

> A team from Zhejiang University and Ant Group proposes OPRD, upgrading distillation from "copying answers" to "copying thought processes" — supervising the student in hidden-state space, bypassing the LM-head information bottleneck, achieving zero-variance gradients, 1.44x training speedup, and 54% memory reduction. On AIME math reasoning, a 1.5B student approaches teacher-level performance for the first time.

1. The Problem: Copying Answers Alone Never Works

LLM distillation has been stuck on the same bottleneck for a decade: every method operates only in output space. From Hinton's classic soft targets to modern On-Policy Distillation (OPD), all approaches compare student and teacher next-token probability distributions.

The OPRD paper identifies two fatal flaws:

1.1 Variance collapse in late training

OPD estimates the KL divergence via a single-sample Monte Carlo estimator: $$log p_t(ŷ_t) - log q_t(ŷ_t)$$ . As the student approaches the teacher ( $$p_t \to q_t$$ ), the signal vanishes but variance does not, collapsing the signal-to-noise ratio and causing accuracy plateaus and oscillation. Modern vocabularies (~150K tokens for Qwen) amplify the problem.

1.2 The LM-head information bottleneck

Output-space distillation treats the teacher as a black-box probability oracle, discarding the entire intermediate stack (L layers × d dimensions). Because softmax is invariant to additive constants, the head matrix $W_{head} \in R^{|V|×d}$ has an effective null space — hidden-state differences along these directions are invisible in the output distribution. A student can parrot the teacher's outputs while its internal reasoning differs entirely, which is fatal for complex reasoning tasks.

2. OPRD: Comparing Hidden States Directly

Core idea: compare student and teacher intermediate hidden states, not just output probabilities.

\mathcal{L}_{\text{OPRD}} = \mathbb{E}_{x, \hat{y}} \left[ \frac{1}{|L_{layer}|} \sum_{l} \frac{1}{\sum m_t} \sum_{t} m_t \frac{1}{d} \left\| h^{(l)}_{\theta,t} - \text{sg}\left(h^{(l)}_{T,t}\right) \right\|_2^2 \right]

Key design choices:

| Component | Meaning | Setting | |------|------|---------| | L_layer | Distilled layers | All 28 layers | | P(ŷ) | Supervised positions | Last k=2000 tokens | | sg(·) | Stop-gradient | Teacher frozen | | d | Hidden dim | 1536 |

Why the last 2000 tokens? Cosine-similarity analysis shows student-teacher divergence concentrates at the response tail where chain-of-thought converges; early tokens stay at 97%+ similarity. This is a data-driven, not hand-tuned, choice.

Combined with OPD: complementary, not a replacement

\mathcal{L} = \mathcal{L}_{\text{OPD}} + \mu \cdot \mathcal{L}_{\text{OPRD}}

Both share the same on-policy rollout and a single teacher forward pass — near-zero infrastructure cost.

| μ | AIME24 | Gain vs μ=0 | |---|--------|------------| | 0 | 42.3 | - | | 1 | 47.7 | +5.4 | | 10 | 50.2 | +7.9 |

3. Two Theoretical Results

1. Theorem 1 (zero-variance gradient): OPRD's gradient is deterministic MSE, while OPD's score-function term $$\nabla_θ log p_t(ŷ_t)$$ injects noise that does not vanish as p→q. This explains why OPD plateaus mid-training while OPRD rises monotonically toward the teacher.

2. Theorem 2 (quantified bottleneck): Hidden-state differences in the LM-head null space produce exactly zero output loss; along the smallest singular direction, the hidden-state/output-loss ratio scales with the squared condition number $$(σ_1/σ_d)^2$$ . Output-space supervision exerts no constraint on intermediate representations; OPRD penalizes exactly these undetectable directions.

4. Experiments: Faster, Cheaper, More Accurate

Teacher: JustRL-Deepseek-1.5B (Qwen2.5-1.5B backbone, 28 layers, d=1536)
Student: DeepSeek-R1-Distill-Qwen-1.5B (same architecture)
Data: DAPO-Math-17K; 2 responses per prompt, temp 1.0, max length 16,384
Compute: 8×A100 (80G), FSDP, 500 optimizer steps

| Method | AIME24 | AIME25 | AIMO | |------|--------|--------|------| | Teacher | 50.8 | 35.6 | 79.5 | | Student (original) | 32.9 | 21.9 | 62.2 | | OPD top-1 | 42.3 | 33.5 | 77.0 | | OPD top-16 | 47.1 | 34.0 | 76.5 | | OPRD | 49.8 | 34.6 | 79.1 |

OPRD trails the teacher by only 1.0/1.0/0.4 points (AIMO within evaluation noise). Notably, OPD top-16 — a strict information superset of top-1 — still plateaus, confirming Theorem 1.

Efficiency (Pareto dominance):

| Metric | OPD top-1 | OPD top-16 | OPRD | |------|-----------|-----------|------| | Peak GPU memory (GB) | 30.2 | 45.0 | 20.5 | | 500-step training time (min) | 813 | 812 | 563 | | AIME24 | 42.3 | 47.1 | 49.8 |

OPRD also yields shorter chains of thought (~5,700 vs ~7,000 tokens), reducing inference cost further.

5. Mechanistic Analysis: Phase Transitions

All OPD+OPRD runs exhibit a loss spike, interpreted as a phase transition from policy reorganization. After the spike, policy-gradient loss ≈ 0 while an accuracy gap remains — direct evidence for Theorem 2: once output-space signals are exhausted, only representation-level supervision can drive further improvement. The top-16 overlap rate with the teacher shows a dip-then-surge pattern, consistent with an internal reorganization followed by a leap to a higher-quality policy.

6. Relation to Prior Work

Unlike FitNets/TinyBERT/MiniLM feature distillation (fixed corpora, encoder models, exposure bias), OPRD aligns hidden states on the student's own sampled distribution, where each state encodes a prediction belief conditioned on the generated prefix. Unlike all output-space OPD variants, OPRD uses all intermediate layers, has zero gradient variance, and costs O(BTd) rather than O(BT|V|).

7. Limitations and Future Work

Isomorphic architecture assumption: teacher and student share architecture/dimensions; cross-architecture distillation needs learnable projections.
Heuristic position choice: last-k=2000 may differ for code or dialogue tasks.
Layer selection: which layers matter most remains open.
Task scope: only math CoT is evaluated; general dialogue and creative tasks are untested.

Future directions include cross-architecture distillation, dynamic layer/position selection, combination with PPO/GRPO-style RL, and multimodal extension.

8. Conclusion

OPRD is a paradigm upgrade, not another trick. It demonstrates: (1) hidden states contain structural information undetectable in output space; (2) zero-variance gradients are feasible in LLM distillation; (3) bypassing the LM head yields both efficiency and quality. Reference paper: Yang, S., et al. (2026). *OPRD: On-Policy Representation Distillation*. Zhejiang University, Ant Group. arXiv:2606.06021. Code: https://github.com/ShenzhiYang2000/OPRD

OPRD: On-Policy Representation Distillation Moves Beyond Output-Space Knowledge Distillation

Summary

OPRD: On-Policy Representation Distillation Moves Beyond Output-Space Knowledge Distillation

1. The Problem: Copying Answers Alone Never Works

1.1 Variance collapse in late training

1.2 The LM-head information bottleneck

2. OPRD: Comparing Hidden States Directly

Combined with OPD: complementary, not a replacement

3. Two Theoretical Results

4. Experiments: Faster, Cheaper, More Accurate

5. Mechanistic Analysis: Phase Transitions

6. Relation to Prior Work

7. Limitations and Future Work

8. Conclusion

Tags