OPRD: On-Policy Representation Distillation Lets Students Peek at the Teacher's Hidden States

Forum topic · 小凯 · 2026-06-07

Summary

Researchers from Zhejiang University and Ant Group propose OPRD (On-Policy Representation Distillation), a new LLM distillation paradigm that supervises the student in the teacher's hidden-state space instead of only matching output token probabilities. The paper identifies two flaws of output-space on-policy distillation (OPD): high-gradient-variance single-sample KL estimates whose signal-to-noise ratio collapses as student and teacher converge, and an information bottleneck caused by the LM head's effective null space, which makes many hidden-state differences invisible in the output distribution. OPRD adds a deterministic zero-variance MSE loss on intermediate hidden states over the last ~2000 tokens of student-generated rollouts, sharing the same rollout and one teacher forward pass. On AIME math reasoning with a 1.5B Qwen2.5-based student, OPRD reaches 49.8 on AIME24 versus the teacher's 50.8, with 1.44x training speedup and 32-54% peak memory reduction compared to OPD. Code and the paper are publicly available.

OPRD: On-Policy Representation Distillation

> A Zhejiang University + Ant Group team 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 Distillation Dilemma: Copying Answers Is Never Enough

LLM distillation typically compares the student and teacher only in output space — next-token probability distributions. Whether classic Hinton-style soft targets or recent On-Policy Distillation (OPD, where the student samples responses and matches the teacher's probabilities), the core operation is the same.

The OPRD paper identifies two critical flaws:

1.1 Variance Catastrophe

OPD estimates the KL divergence with a single-sample Monte Carlo estimate: log p_t(ŷ_t) − log q_t(ŷ_t). As the student approaches the teacher (p_t → q_t), the signal goes to zero but variance does not. The signal-to-noise ratio collapses in later training, causing accuracy plateaus and oscillation — a "variance wall" that more training cannot break. Modern vocabularies (~150K tokens for Qwen) amplify the problem.

1.2 Information Bottleneck

Output-space distillation treats the teacher as a black-box probability oracle, discarding the entire intermediate computation stack (L layers × d-dim hidden states). Because softmax is invariant to additive constants, the LM head projection W_head ∈ R^{|V|×d} has an effective null space: hidden-state differences along null-space directions are completely undetectable in the output distribution, yet may represent very different internal reasoning. The student can learn to parrot the output distribution while thinking entirely differently internally — fatal for complex math and code reasoning.

2. OPRD: From "Copying Answers" to "Copying Reasoning"

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

The loss is a masked MSE over hidden states:

\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]

| Component | Meaning | Typical setting | |---|---|---| | L_layer | distilled layers | all 28 layers | | P(ŷ) | supervised positions | last k=2000 tokens | | m_t | position mask | 1[t ∈ P(ŷ)] | | sg(·) | stop-gradient | teacher frozen | | d | hidden dim | 1536 |

Why the last 2000 tokens? Cosine-similarity analysis shows student–teacher representation divergence concentrates in the tail of the response (where chain-of-thought converges to the final answer); early segments stay above 97% similarity. The choice is data-driven, not hand-tuned.

2.2 Complementary to OPD

\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 accuracy rises monotonically with μ:

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

This confirms hidden-state and output-space signals are additive and complementary.

3. Theory: Zero Variance + Bottleneck Breakthrough

Theorem 1 (zero-variance gradients). OPRD's gradient is deterministic — a fixed MSE for a given rollout. OPD's gradient is high-variance: the score-function term ∇_θ log p_t(ŷ_t) keeps injecting noise even as p → q. This explains why OPD plateaus mid-training while OPRD climbs monotonically toward the teacher.

Theorem 2 (LM-head bottleneck). With singular values σ_1 ≥ ... ≥ σ_d of W_head: 1. If h_θ − h_T lies in the null space N_W, the output loss is zero regardless of the magnitude of the hidden-state difference. 2. Along the smallest-singular-value direction, the ratio of hidden-state norm to output loss is lower-bounded by the squared condition number (σ_1/σ_d)² — typically enormous in production LLMs.

Output-space OPD therefore exerts no constraint on many hidden-state directions; OPRD penalizes exactly these directions.

4. Experiments: Faster, Cheaper, More Accurate

Setup

| | Teacher | Student | |---|---|---| | Model | JustRL-Deepseek-1.5B | DeepSeek-R1-Distill-Qwen-1.5B | | Backbone | Qwen2.5-1.5B | Qwen2.5-1.5B | | Layers / hidden dim | 28 / 1536 | 28 / 1536 | | Vocab | ≈151K | ≈151K |

Data: DAPO-Math-17K; 2 responses per prompt, temperature 1.0, max length 16,384 tokens
Training: 8×A100 (80G), FSDP, 500 optimizer steps
Eval: AIME 2024/2025, AIMO (AMC 2022+2023), Avg@16

Accuracy: Approaching the Teacher

| Method | AIME24 | AIME25 | AIMO | |---|---|---|---| | Teacher | 50.8 | 35.6 | 79.5 | | Student (unmodified) | 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). OPD top-16 — a strict information superset of top-1 — still plateaus, confirming Theorem 1: extra output-layer information cannot cancel sampling noise.

Efficiency: Pareto Triple Win

| 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 |

1.44x training speedup (no B×T×|V| logits materialization)
32–54% peak memory reduction
Strictly Pareto-dominant on all three axes

More Concise Responses

OPRD converges to ~5,700-token responses vs ~7,000 for OPD — higher accuracy with shorter reasoning chains.

5. Mechanistic Analysis: Phase Transition Hypothesis

All OPD+OPRD runs show a loss spike, interpreted as a phase transition of policy reorganization. Notably, after the spike the PG loss ≈ 0, yet the accuracy gap persists (+5.4/+7.9 points) — direct evidence for Theorem 2: once policy-gradient signal vanishes, remaining improvements live in the LM head's null space, reachable only via representation-level supervision. The top-16 overlap rate between student and teacher also shows a dip-then-surge pattern, consistent with the student undergoing an internal reorganization before jumping to a higher-quality policy.

6. Relation to Prior Work

| Dimension | FitNets / TinyBERT / MiniLM | OPRD | |---|---|---| | Supervision data | fixed corpora | student-generated rollouts | | Exposure bias | present | eliminated (on-policy) | | Model type | encoders (BERT, CNN) | autoregressive LLMs | | Representations | one-shot | conditioned on sampled prefix |

| Dimension | OPD (all variants) | OPRD | |---|---|---| | Supervision space | logits | hidden states | | Gradient variance | high (REINFORCE) | zero (deterministic MSE) | | Teacher info used | final distribution only | all intermediate layers | | Memory | O(BT\|V\|) | O(BTd) |

7. Limitations and Future Work

Limitations: requires teacher–student architectural homogeneity (no projector needed); the last-k=2000 suffix heuristic may differ across tasks; layer selection is default-all; experiments are limited to math CoT reasoning.

Future directions: learnable projectors for cross-architecture distillation; dynamic layer/position selection; combination with RL methods like PPO/GRPO; extension to multimodal models.

8. Conclusion

OPRD is an upgrade to the distillation paradigm itself, demonstrating that:

1. Output space is not the only supervision channel — hidden states contain structurally undetectable-but-important information. 2. Zero-variance gradients are feasible in LLM distillation — deterministic MSE replaces high-variance REINFORCE estimates. 3. Bypassing the LM head yields dual efficiency/quality gains — 1.44x speedup, 54% memory cut, near-teacher accuracy.

> "We show that hidden-state supervision is not just a richer signal source, but a necessary condition for breaking the LM-head information bottleneck. When output-space signal is exhausted, optimization in representation space is just beginning." — OPRD authors

Reference: Yang, S., Zhu, G., Song, B., Wang, H., Xia, M., Zheng, X., Ma, Y., Chen, Z., Wang, W., & Chen, G. (2026). *OPRD: On-Policy Representation Distillation*. Zhejiang University, Ant Group. arXiv:2606.06021.

Code: https://github.com/ShenzhiYang2000/OPRD