OPRD: On-Policy Representation Distillation

arxiv:2606.06021

Original source · Paper JSON

Post-training · 2026-06-04

Cites (10) · Cited by (0)

Complete cleaned Markdown · Source Markdown including bibliography

Read in parts

1. Part 1
2. Part 2
3. Part 3
4. Part 4
5. Part 5
6. Part 6
7. Part 7
8. Part 8
9. Part 9

Contents

• OPRD: On-Policy Representation Distillation
• Abstract
• 1 Introduction
• 2 Background and Problem Setup
• 2.1 Notation
• 2.2 The On-Policy Distillation Framework
• Setup.
• Objective.
• 2.3 Three Output-Space Variants
• (a) Sampled-token OPD ( S t = { y ^ t } S_{t}=\{\hat{y}_{t}\} ).
• (b) Full-vocabulary OPD ( S t = 𝒱 S_{t}=\mathcal{V} ).
• (c) Top- k k OPD ( S t = TopK ⁡ ( p t , k ) S_{t}=\mathrm{TopK}(p_{t},k) ).
• 2.4 A Shared Structural Limitation
• 3 On-Policy Representation Distillation
• 3.1 The OPRD Objective
• 3.2 Why OPRD Works
• Theorem 1 (Zero-variance gradient) .
• Comparison with top- k k : no truncation bias.
• Theorem 2 (Hidden-state information beyond the LM head) .
• 3.3 OPRD-Bridge: Cross-Architecture Extension
• 3.3.1 The Shared Low-Rank Structure
• 3.3.2 Stage 1: Bridge Construction
• 3.3.3 Stage 2: Bridge-Based Distillation
• 3.3.4 Design Philosophy
• Theorem 3 (Optimality and bias–variance trade-off of the low-rank bridge, informal) .
• 4 Experiments
• 4.1 OPRD-Vanilla: Same-Architecture Distillation
• 4.1.1 Experimental Setup
• Models.
• Training data.
• Distillation objectives.
• Optimization.
• Evaluation.
• 4.1.2 Main Results
• 4.1.3 Training Dynamics
• Accuracy curves: OPRD climbs monotonically, OPD plateaus.
• Behavioural view: OPRD produces shorter, more efficient reasoning.
• Internal view: OPRD’s own loss is being optimised end-to-end.
• 4.1.4 Efficiency
• 4.1.5 Mechanistic Analysis
• Composing OPD with OPRD ( μ \mu sweep).
• Where does the student diverge from the teacher? (motivation for last- k k supervision).
• (a) Policy-gradient loss: OPRD accelerates distillation and validates the bottleneck theory.
• (b) Top- 16 16 overlap: hidden-state alignment propagates to next-token agreement.
• (c) Predictive entropy.
• 4.2 OPRD-Bridge: Cross-Architecture and Cross-Tokenizer Distillation
• 4.2.1 Experimental Setup
• Models.
• Bridge construction (Stage 1).
• Distillation (Stage 2).
• Baselines.
• Evaluation.
• 4.2.2 Main Results
• 4.2.3 Ablation Studies
• Effect of bridge rank r r .
• 4.2.4 Cross-Tokenizer Distillation
• 5 Discussion
• From representational alignment to cross-architecture distillation.
• High-value application 1: multi-model RL merging.
• High-value application 2: on-policy self-distillation (OPSD).
• Toward alignment-aware pre-training.
• 6 Related Work
• Output-Space Knowledge Distillation.
• On-Policy Distillation.
• Feature / Intermediate-Representation Distillation.
• Hint Learning, Auxiliary Losses, and Distribution Matching.
• 7 Conclusion and Future Work
• Future Work.
• Appendix A Notation Summary
• Appendix B Formal Theoretical Guarantees
• B.1 Setup and Assumptions
• Definition 1 (Stochastic gradient estimators) .
• Assumption 1 (Standard regularity) .
• B.2 Variance of Sampled-Token OPD
• Lemma 1 (Score-function decomposition of OPD gradient) .
• Proof.
• Theorem 4 (OPD gradient variance lower bound) .
• Proof.
• B.3 OPRD Gradient Is Deterministic
• Theorem 5 (OPRD has zero conditional variance) .
• Proof.
• Corollary 1 (Variance gap) .
• B.4 Signal-to-Noise Ratio Collapse of Sampled-Token OPD
• Definition 2 (Signal-to-noise ratio) .
• Theorem 6 (SNR collapse for OPD, SNR stability for OPRD) .
• Proof.
• Remark 1 (Interpretation: late-stage stagnation of pure OPD) .
• Sub-summary (Perspective 1).
• B.5 Formal Results for Theorem 2 : LM-Head Information Bottleneck
• Definition 3 (Effective null space of the LM head) .
• Lemma 2 (Softmax kernel) .
• Proof.
• Theorem 7 (Null-direction identity; formal version of ( 9 )) .
• Proof.
• Lemma 3 (Local Lipschitzness of output-space losses) .
• Proof sketch.
• Theorem 8 (Spectral gap; formal version of ( 10 )) .
• Proof.
• Remark 2 (Intermediate layers) .
• Sub-summary ( Theorem 2 ).
• Overall summary of § B .
• B.6 Formal Results for Theorem 3 : Optimality of the Low-Rank Bridge
• B.6.1 Setup
• Assumption 2 (Monotone alignment decay) .
• B.6.2 Rate–Distortion Optimality
• Theorem 9 (PCA bridge is rate–distortion optimal) .
• Proof.
• B.6.3 Bias–Variance Decomposition of Distillation Error
• Definition 4 (Distillation error at rank r r ) .
• Theorem 10 (Bias–variance decomposition) .
• Proof.
• Instructions for reporting errors