OPRD Makes Distillation Cheaper by Supervising the Reasoning State Before the Logits

OPRD is a distillation paper with a very simple engineering smell: the teacher already computed a rich internal representation, and the training loop throws it away so the student can imitate a probability distribution over a 150,000-token vocabulary. That is not elegance. That is paying for the whole review and reading only the final approval emoji.

The paper, “OPRD: On-Policy Representation Distillation”, proposes On-Policy Representation Distillation as an alternative to output-space on-policy distillation. Instead of supervising the student only by matching next-token probabilities, OPRD aligns hidden representations between teacher and student across selected layers on the same student rollouts. The claim is not just that this is more informative. The claim is that it is cheaper, less noisy, and better aligned with the part of reasoning models that actually matters: the intermediate state before the final token choice.

For anyone operating reasoning models, coding agents, or model-routing systems, this is not an academic footnote. Distillation is becoming infrastructure. Teams want stronger teachers to compress into cheaper students, domain-tuned students to run locally, and specialized executors to handle high-volume agent work without paying frontier-model prices on every turn. If that loop is slow, memory-hungry, and noisy, the whole “use the right model for the right step” story gets expensive fast.

The output layer is a lossy API

Standard on-policy distillation supervises a student on trajectories sampled from the student itself. That matters because the student learns on its own distribution, not only on static teacher examples. But most OPD variants still operate in output space: match the teacher’s next-token distribution by sampling tokens, using the full vocabulary, or taking top-k logits.

Each option has a tax. Sampled-token OPD has Monte Carlo variance, and that variance does not magically disappear when the vocabulary is huge. The OPRD paper calls out Qwen-style vocabularies around 150K tokens, which makes sparse sampling a noisy estimate of the teacher’s preference landscape. Full-vocabulary or top-k variants reduce some approximation problems but create memory pressure, especially when traces are long. Reasoning-model distillation is where all the bad dimensions line up: long response lengths, large vocabularies, repeated rollouts, multiple evaluations, and expensive teachers.

OPRD moves the target. Instead of matching after the LM head, it aligns normalized hidden states across selected layers and response positions. The student and teacher process the same rollout, and the loss lives in representation space. In the paper’s framing, that bypasses the large-vocabulary output bottleneck and gives the student richer per-layer structural information than a next-token probability vector can provide.

There is a catch: the method currently assumes the teacher and student share an architecture template so that hidden states are meaningfully alignable. That is a big constraint. But inside same-family or self-distillation settings, it is exactly where many open-model teams already operate.

The numbers are operator-relevant, not just leaderboard polish

The experimental setup uses an R1-distill-1.5B student and a JustRL-1.5B teacher, trained for 500 optimizer steps on 8× A100 80GB with FSDP, batch size 8, and maximum response length 16,384. Evaluation covers AIME 2024, AIME 2025, and AIMO, reported as Avg@16 at temperature 0.7.

The teacher scores 50.8 on AIME 2024, 35.6 on AIME 2025, and 79.5 on AIMO. The student starts at 32.9, 21.9, and 62.2. OPRD reaches 49.8, 34.6, and 79.1 — within 1.0, 1.0, and 0.4 points of the teacher. It also beats the strongest output-space baseline in the table by 2.7, 0.6, and 2.1 points. OPD top-1 reaches 42.3, 33.5, and 77.0; OPD top-16 reaches 47.1, 34.0, and 76.5.

Those are respectable benchmark gains. The cost numbers are the reason this belongs in an engineering newsletter.

Actor-update transient memory is 20.5GB per GPU for OPRD, compared with 30.2GB for OPD top-1 and 45.0GB for OPD top-16. That is a 32% reduction versus top-1 and a 54% reduction versus top-16. Wall-clock time over 500 steps is 563 minutes for OPRD versus 813 and 812 minutes for the OPD baselines, roughly 1.44× faster. The paper also reports a behavioral side effect: OPRD converges to about 5,700 tokens per response, while OPD variants plateau around 7,000, suggesting the representation-supervised student may produce shorter reasoning traces at higher accuracy.

That last point is worth watching. Reasoning models do not only cost more because they are bigger. They cost more because they talk to themselves for a long time. If a distillation method improves accuracy while shortening traces, it compounds: cheaper training, cheaper inference, lower latency, less context pollution for downstream agent steps.

Why this matters for coding agents

Coding agents are mostly sold as UX improvements, but under the hood they are budget-routing systems. A real agent may use a planner, a code reader, a patch writer, a test runner, a reviewer, and a retry loop. If every step calls a premium reasoning model, the unit economics get ugly. The obvious fix is to distill specialized behaviors into cheaper models and route only the hard cases upward.

OPRD is interesting because it improves the economics of creating those cheaper models. Output-space distillation asks the student to imitate the teacher’s visible token choices. Representation distillation asks it to imitate some of the internal computation that made those choices possible. For reasoning-heavy tasks, that distinction should matter. The final token is often a poor summary of the search state, uncertainty, failed branches, and intermediate abstractions the teacher carried forward.

The metaphor is code review again: copying the merged diff teaches you what changed; reading the review discussion teaches you why. OPRD is closer to supervising the why, as long as the model family exposes comparable internal state.

Practically, teams should not treat OPRD as a universal replacement for distillation. If your teacher is a closed model, you cannot access hidden states. If your teacher and student architectures differ, alignment is not straightforward. If your production workload is not math-style reasoning, you need to reproduce the result on tasks that resemble your failures: repo edits, tool calls, structured data repair, multi-file refactors, or long-horizon agent plans.

But the design direction is strong. Open-weight vendors and internal platform teams should care about same-family distillation loops. If you run a Qwen-family, Llama-family, DeepSeek-family, or internal model stack, hidden-state supervision may be a better target than increasingly clever approximations of a giant logit tensor. It also suggests vendors should expose better training-time signals for enterprise adaptation rather than pretending the API response is the only useful artifact a model produces.

There is also a governance angle. Cheaper distillation makes continuous adaptation more realistic, but continuous adaptation creates lifecycle problems: which teacher produced this student, on what data, with what benchmark gates, and when should it be retired? If distillation becomes CI for models, it needs CI discipline — reproducible runs, eval thresholds, rollback paths, provenance, and cost regression tracking.

The clean take: OPRD is not “distillation solved.” It is a reminder that the logits are a narrow interface for a very expensive computation. If reasoning traces are the thing you are trying to transfer, supervising only the next token is leaving signal on the floor. For open-model stacks trying to make reasoning agents affordable, that floor is starting to look like the budget.

Sources: arXiv, OPRD code repository, DeepSeek-R1 distillation context, Qwen model-family context