OPMD: On-Policy Mechanism Distillation Beyond Output Supervision

AI summary

Overview

Research area: Machine learning — knowledge distillation for large language models, specifically on-policy distillation and internal-representation (latent-state) supervision for reasoning tasks.

Technical level: Advanced. The paper assumes familiarity with autoregressive language models, KL-divergence-based distillation, LoRA fine-tuning, hidden-state/representation matching, and activation-patching analysis.

Scope (one sentence): The paper proposes OPMD, a two-stage distillation method that builds a "student-compatible teacher" via learned injection projectors and then trains the student to internalize the teacher-induced latent states on its own generated prefixes, evaluated on mathematical reasoning, multi-teacher distillation, and cross-family SFT.

What This Paper Is About

On-policy distillation (OPD) fixes the mismatch between training and deployment by supervising a student model on prefixes it generated itself, but all of that supervision lives in the output space: it transfers what the teacher predicts, not the intermediate computation that produces those predictions. The authors show that naively matching intermediate hidden states is not a fix, because teacher and student representations can differ geometrically, and even a well-aligned hidden state may be unusable by the student's remaining layers. The goal is to transfer teacher-derived intermediate computation in a form the student's own downstream layers can actually use, and then to internalize that computation into the student's weights so nothing extra is needed at inference.

Key Contributions

1. A student-compatible teacher construction. A learned bridge of layer-wise projectors injects frozen teacher hidden states as residual updates into a frozen student backbone, and the projectors are trained purely by their downstream functional effect (reverse KL to the teacher, or an SFT objective when tokenizers differ) rather than by hidden-state reconstruction error.

2. On-policy mechanism internalization. A native student is trained to reproduce the student-compatible teacher's latent trajectories on prefixes the student itself samples, using a stop-gradient MSE objective at selected layers; the teacher and projectors are discarded at inference.

3. Broad empirical validation. The framework is tested across three teacher–student settings on AMC, AIME 2024, and AIME 2025, extended to two domain teachers (math and code) and to cross-family supervised fine-tuning (Qwen3-4B to Llama3.1-8B-Instruct) without logit-level KL supervision.

4. Layer-level behavioral analysis. Weight-patching experiments and a two-problem solution-route case study probe whether different layer groups internalize different aspects of the reasoning mechanism, distinguishing selecting a solution strategy from executing it successfully.

Main Findings

• Latent states carry actionable task information. Cross-model activation patching on AMC 2023 showed that replacing a base model's hidden state at a selected layer with the corresponding state from its RL-optimized counterpart substantially improved downstream accuracy, leaving all subsequent computation unchanged.

• Geometric alignment is not enough. When patching heterogeneous teacher–student pairs, teacher states mapped by an MSE-trained projector could severely disrupt student generation, whereas states from a task-trained projector yielded substantial gains across patching depths.

• Highest average accuracy across all three settings. Against SFT, OPD, and OPRD-Bridge at 512 student updates, OPMD achieved the best average accuracy in all three teacher–student configurations. AMC gains over OPD were 3.76, 7.04, and 2.66 percentage points respectively.

• Gains on the hardest benchmarks. In the RL-Math to Qwen3-1.7B-Base setting, OPMD also improved AIME 2024 by 4.37 points and AIME 2025 by 3.95 points over OPD (reaching 9.79% and 7.08%).

• A stronger teacher does not automatically give a stronger student. Switching the teacher from Qwen3-4B to RL-Math lowered OPD's AMC accuracy from 35.31% to 33.12%, but raised OPMD's from 39.07% to 40.16%, indicating OPMD exploits a specialized teacher more effectively.

• 512 updates do not exhaust the benefit. Extending to 1,024 updates, OPMD reached its highest observed accuracy at 1,024 on AMC and AIME 2025; AIME 2024 instead peaked at 768 updates, so the benefit of longer training is task-dependent.

• Better accuracy at lower token cost. OPD's responses grew substantially longer during training with only modest AMC improvement. OPMD produced higher accuracy with shorter responses at every evaluated update count, and even at 256 updates exceeded OPD's accuracy at 512 and 768 updates while generating fewer tokens. Both used the same 7,168-token evaluation limit.

• Multi-teacher distillation transfers. With RL-Math and RL-Code teachers into Qwen3-1.7B-Base, OPMD improved over M-OPD by 4.31 points on AMC, 0.61 on HumanEval+, and 6.61 on MBPP+ — the shared student improved on both domains, and the uneven code gains show the benefit is task-dependent rather than uniform.

• Mechanism supervision works without logit-level KL. In cross-family SFT (Qwen3-4B to Llama3.1-8B-Instruct), SFT+MD beat plain SFT on AMC (25.31% vs. 22.50%) and AIME 2025 (1.81% vs. 1.46%), but decreased on AIME 2024 (3.88% vs. 4.33%).

• Partial layers recover much of the gain. Patching blocks 1–7 from the step-1024 OPMD student into raw Qwen3-1.7B-Base raised accuracy from 14.84% to 31.88%; patching blocks 8–14 reached 34.06%; the full OPMD student reached 40.68%. Both patches also shortened responses (1,785 and 1,936 tokens vs. 3,147 for the raw student and 2,462 for the full student), separating brevity from accuracy.

• Strategy preference and execution dissociate. On two selected AMC problems (q17, q22), the block 8–14 patch moved solution routes closer to those of the full OPMD student but retained a large accuracy gap — for q17, 62.50% vs. 81.25%; for q22, 31.25% vs. 50.00%. The authors read this as different layer groups retaining different aspects of the reasoning mechanism.

• Favorable cost profile. On two NVIDIA H100 GPUs, OPMD reached 38.96% AMC accuracy at 768 student updates versus 36.50% for OPD at 1,024 updates, with Stage I costs included. A shorter Stage I also suffices: with RL-Math, 32 Stage I updates gave 40.00% AMC accuracy versus 40.16% after 96 updates.

Methodology in Plain English

The method separates two problems that prior representation-matching work conflates: making the teacher's information usable, and making it the student's own.

Stage I — build a teacher the student can actually use. Take a strong teacher and a student, and freeze both. Insert a small learned projector at each of several chosen teacher-layer/student-layer pairs. The projector maps the teacher's hidden vector into the student's hidden dimension and adds it as a residual into the student's forward pass at that layer. Because the student's backbone is frozen, the only way the projectors can help is by expressing teacher information in a form the student's existing downstream layers can act on. Crucially, the projectors are trained by what the injected network outputs: reverse KL against the teacher's next-token distribution over the full vocabulary when the two models share a tokenizer, or a supervised next-token objective on teacher-generated responses tokenized by the student when they do not. No hidden-state reconstruction loss is used. The result is a "student-compatible teacher" — the frozen teacher plus projectors plus frozen student.

Stage II — make the student do it alone. The student-compatible teacher still depends on the teacher and projectors, so a second stage trains the native student (LoRA adapters, rank and alpha 64, zero dropout, AdamW at learning rate 1e-6) to match the student-compatible teacher's intermediate states. The student samples its own responses; both the native student and the frozen student-compatible teacher are evaluated on those same student-generated textual prefixes at each response-token position without shifting, and the student minimizes a normalized squared-distance loss to the (stop-gradient) teacher-induced states at selected layers. Because the targets were produced inside the frozen student backbone, they already live in the student's representation space and need no cross-model mapping. In the on-policy experiments this stage uses only the latent-state objective, without a downstream task loss.

At inference, the teacher and all projectors are removed; only the native student remains.

The evaluation uses DAPO-Math-17K for training and AMC, AIME 2024, and AIME 2025 for testing, reporting avg@16 (mean final-answer accuracy over 16 sampled responses per problem; checkpoint analyses used 40 AMC problems and 30 problems per AIME benchmark at temperature 0.7 and top-p 0.95). Baselines are output-level OPD, SFT on teacher responses, and OPRD-Bridge. A separate SFT variant, SFT+MD, combines the latent-state MSE (coefficient 0.5) with response negative log-likelihood (coefficient 1) on teacher-response replay data, with no KL on logits.

Why This Matters

Impact on research. The paper reframes distillation from "match the teacher's outputs" or "match the teacher's vectors" to "produce intermediate states the student can computationally use." Its preliminary activation-patching evidence — that a task-trained projector succeeds where an MSE-trained projector disrupts generation — is a concrete argument that representational closeness is the wrong objective, which challenges a broad family of feature-distillation methods that optimize hidden-state distance (Equation 2 in the paper). The multi-teacher and cross-family SFT results further suggest mechanism supervision is a general add-on rather than a single-benchmark trick.

Real-world applications:

• Small on-device reasoning models. Distilling a large teacher's reasoning mechanism into 1.7B–3B students with shorter outputs is directly relevant to latency- and cost-sensitive deployment, where the paper reports better accuracy at fewer generated tokens.
• Specialized domain teachers. Routing prompts to a math teacher or a code teacher with matching projector banks lets one student absorb multiple specialists, relevant to enterprise assistants that must cover several domains in one model.
• Cross-vendor model consolidation. The Qwen3-4B to Llama3.1-8B-Instruct experiment shows mechanism supervision working across model families and tokenizers, useful when a company wants to compress a licensed teacher into its own architecture.
• Cheaper distillation pipelines. Since Stage I can be shortened (32 vs. 96 updates gave 40.00% vs. 40.16% AMC) and Stage II matches OPD's update throughput with a vLLM implementation, the method is plausible for routine production training rather than one-off research runs.

Industry relevance. Post-training budgets are dominated by rollout and teacher-query cost. A method that extracts more from the same teacher, shortens student outputs, and keeps inference free of any auxiliary modules addresses both the training bill and the serving bill. The finding that a stronger teacher does not automatically yield a stronger student is also practically important: teacher selection cannot be done by benchmark score alone.

Future Directions

• Scale and breadth. The limitations section states plainly that the evaluation covers a limited set of teacher–student pairs and tasks; whether gains hold for larger models, more architectures, and tasks beyond mathematics and code generation is untested.
• Understanding layer specialization. The two case studies (q17, q22) were selected because the OPMD student beat Patch 1–7 and do not establish general layer specialization or causal mechanisms. A systematic mapping of which layer groups carry strategy selection versus execution competence is open.
• Closing the strategy–execution gap. The block 8–14 patch reproduced the full model's solution-route preferences while trailing far behind in accuracy (e.g., 31.25% vs. 50.00% on q22). What additional computation is missing, and whether it can be supervised directly, is unresolved.
• Choosing transfer-layer pairs and projector objectives. The paper defines a set of teacher–student transfer-layer pairs and two projector objectives (reverse KL for shared tokenizers, SFT for differing tokenizers) but does not report a systematic study of how pair selection or objective choice affects the outcome.

Target Audience

Researchers and engineers working on LLM post-training, distillation, and reasoning-model compression who already understand on-policy distillation and hidden-state supervision and want a method that treats intermediate computation as a transferable, usability-constrained target. It is also relevant to interpretability researchers interested in activation patching and weight patching as tools for localizing reasoning behavior, and to practitioners deciding whether representation-level distillation is worth the added engineering over output-level OPD. Readers without background in distillation objectives and transformer internals will find the method section demanding.

Authors’ abstract

On-policy distillation (OPD) improves language-model by supervising a student on prefixes sampled from its own policy, but its supervision remains confined to the output space. Consequently, OPD transfers what a teacher predicts without supervising the intermediate computation that produces those predictions. Simply distilling intermediate representations like representation matching is insufficient since teacher and student representations may differ geometrically, and even well-aligned hidden states need not be usable by the student's subsequent computation. We introduce On-Policy Mechanism Distillation (OPMD), which transfers teacher-derived intermediate computation while explicitly accounting for its compatibility with the student. OPMD first constructs a student-compatible teacher that transforms teacher representations into intermediate states usable by the student's downstream computation, then it internalizes these states from on-policy trajectories. OPMD achieves the highest average accuracy among the evaluated distillation baselines in all three teacher--student settings, and generalizes to multi-teacher distillation and cross-family supervised fine-tuning. Further analyses show that mechanism supervision yields more effective reasoning and meaningfully alters the student's intermediate computation and solution strategies, supporting that it transfers information beyond output supervision alone. Together, these results establish mechanism-level supervision as an effective way to move on-policy distillation beyond output supervision, enabling students to learn inner mechanisms they can use and internalize.