Zhejiang University's OPRD Trains Smarter AI 1.44x Faster Using Hidden States
OPRD moves LLM distillation from output probabilities into hidden states, closing the student-teacher gap while training 1.44x faster and using 54% less GPU memory.
- OPRD aligns student and teacher hidden states instead of output token probabilities, bypassing the LM head entirely.
- Closes the student-teacher gap on AIME 2024/2025 and AIMO benchmarks, where all output-space OPD baselines plateau several points below the teacher.
- Trains 1.44x faster and uses 54% less GPU memory than top-16 OPD by never materializing the [batch, seq, vocab] logit tensor.
- Theoretically proven to have zero gradient variance given a rollout, while standard OPD's signal-to-noise ratio collapses to zero as training progresses.
- Requires student and teacher to share the same architecture; cross-architecture support is left as future work.
- Code is available at github.com/ShenzhiYang2000/OPRD; paper at arxiv.org/abs/2606.06021.
On-policy distillation (OPD) has become a standard ingredient in post-training pipelines for large language models. The idea is elegant: let the student generate its own responses, then score each token against the teacher's probability distribution, creating a dense, token-level training signal. Production systems from Qwen3 to DeepSeek-V4 all rely on some form of it. But a new paper from Zhejiang University and Ant Group argues that every existing OPD variant is fundamentally bottlenecked by the same design choice: they all supervise the student exclusively through the model's output layer, throwing away everything the teacher computed internally.
The paper introduces On-Policy Representation Distillation (OPRD), the first method to move on-policy distillation from the output space into the hidden-state space. The result is a training objective that is simultaneously more accurate, faster, and cheaper than the best existing baselines.
Why the output layer is a bottleneck
To understand the problem, it helps to know how standard OPD works. At each training step, the student samples a full response, then both the student and the frozen teacher are run forward on that same sequence. At each token position, you compare the student's next-token probability distribution to the teacher's, and penalize the divergence. The most common variant, called sampled-token OPD, only looks at the single token the student actually sampled, making it a single-sample Monte Carlo estimate of the KL divergence.
The paper identifies two structural problems with this paradigm:
- Variance collapse. The sampled-token estimator is unbiased but noisy. Early in training, when the student and teacher disagree a lot, the signal dominates the noise and the student improves quickly. But as the student converges toward the teacher, the signal shrinks while the variance stays constant. The signal-to-noise ratio collapses, and the gradient becomes dominated by noise. The student stalls well below the teacher.
- The LM head is an information bottleneck.
This story is for Pro members
You've reached the end of the free preview. Upgrade to AlphaSignal Pro to read the full article - and everything else behind the paywall.
Comments
