OPRD: On-Policy Representation Distillation

OPRD: Lifting On-Policy Distillation from Output Space to Hidden-State Space

Current methods for teaching large language models (LLMs) often rely on on-policy distillation. In this setup, a smaller "student" model learns by mimicking a larger "teacher" model. The student generates its own responses and checks them against the teacher's logic. However, this process is currently limited. It only evaluates the student based on the final words it chooses. A new paper proposes moving beyond this superficial feedback. We can look at the "internal thoughts"—the intermediate hidden states (mathematical vectors representing concepts)—of the teacher model.

Can we teach a model by its thoughts rather than its words?

The core challenge in post-training LLMs is transferring reasoning capabilities. We want to move these skills from massive models into smaller, deployable ones. On-policy distillation (OPD) is a vital tool for this. It allows the student to learn from its own mistakes in real-time. The student samples its own tokens. It then calculates the divergence (a measure of how much two probability distributions differ) between its choices and the teacher's. This provides a dense, token-level signal.

Is the final word choice a complete proxy for reasoning? The researchers investigate if the current paradigm is inherently flawed. They focus on the "output space" (the final probability distribution over the vocabulary). They ask if this approach creates a performance ceiling. This ceiling stems from statistical noise and information loss. Shifting supervision into the "hidden-state space" might break this ceiling.

The cracks in the output-only paradigm

Most current methods match the teacher's next-token probabilities. There are three primary ways to do this. Researchers use sampled-token, full-vocabulary, or top-$k$ (focusing on the most likely tokens) methods. While these vary in complexity, they share a critical blind spot.

First, the most common method—sampled-token distillation—relies on a Monte Carlo estimate (approximating a value via random sampling). The vocabulary of models like the Qwen series is massive. It contains roughly 150,000 tokens. Estimating the correct direction for improvement using only one sample introduces significant variance (randomness in the signal). As the student resembles the teacher, this noise drowns out the learning signal. This causes performance to plateau or oscillate .

Second, the authors identify an informational bottleneck. Every LLM ends with a language-model (LM) head. This is a final linear layer that projects high-dimensional internal representations into the vocabulary space. The researchers argue this projection is "ill-conditioned." It compresses the teacher's rich, $d$-dimensional internal states into a much lower-rank output. Consequently, the student is graded only on the sliver of knowledge that survives this projection. It receives no signal regarding the structural reasoning stored in the teacher's intermediate layers .

Moving the goalposts to the hidden states

To resolve these tensions, the researchers propose On-Policy Representation Distillation (OPRD). OPRD does not ask the student to match final word probabilities. Instead, it asks the student to match the teacher's internal hidden states. The investigators implement this by selecting specific layers and response positions. They target the "tail" of a response where reasoning usually culminates. They minimize the mean-squared error (the average squared difference) between the student's and teacher's representations.

The experimental design was rigorous. The authors used a student model (R1-distill-1.5B) and a teacher model (JustRL-1.5B). Both share the same architecture. They tested OPRD against established output-space baselines on three math benchmarks: AIME 2024, AIME 2025, and AIMO. This setup isolated the effect of the supervision target. It kept the computational budget and generated text identical across all methods.

Closing the gap through internal alignment

The findings suggest that shifting to representation-level supervision is transformative. The authors report that OPRD effectively closes the gap between the student and the teacher. On the AIMO benchmark, OPRD achieved 79.1% accuracy. This essentially tied with the teacher. In contrast, all output-space OPD baselines fell significantly short [Table 2].

The training dynamics revealed why this happens. Standard OPD variants hit a performance wall. OPRD demonstrated monotonic improvement. It continued to climb even in late-stage training . This supports the authors' theoretical claim. OPRD provides a deterministic, zero-variance gradient. Standard OPD suffers from a signal-to-noise ratio collapse as the student nears the teacher .

The method also proved remarkably efficient. The OPRD loss is calculated before the expensive LM-head projection. This avoids materializing massive vocabulary-sized tensors. The researchers found that OPRD trains 1.44$\times$ faster. It also reduces peak transient memory for actor updates by up to 54% compared to top-$k$ OPD [Table 3]. This memory saving is hardware-relevant. It could allow for larger batches or longer contexts on the same GPU. Furthermore, OPRD models produced shorter reasoning chains. This suggests internal alignment encourages efficient logic rather than mere imitation .

A new axis of supervision

This work proves that internal representations are an under-utilized resource. We should treat the teacher as a structured source of layered computation. We can guide students to inhabit the same conceptual geometry as the teacher.

The efficiency gains suggest a path toward scaling distillation. OPRD could help merge multiple specialized teachers into one student. It does this without the memory explosion seen in full-vocabulary matching. This makes it an attractive option for production environments.

One significant hurdle remains: the "same-architecture" requirement. The authors note that hidden states from different model sizes are nearly orthogonal (meaning they share almost no common direction). Currently, OPRD requires the student and teacher to share the same depth and width. Future research could investigate learnable projection heads. These might allow OPRD to bridge the gap between different model sizes.

Figures from the paper