OPRD: On-Policy Representation Distillation — part 6/9

Cleaned-text code points [78225, 93947). Bibliography omitted here; source.md includes it.

(1) Representation-space alignment bypasses the vocabulary barrier. Standard OPD, which relies on the shared vocabulary as its alignment channel, collapses when the two tokenizers are incompatible. OPRD-Bridge instead aligns through the learned projector pair in representation space, lifting the student from near-zero performance ( 2.8 % 2.8\% Avg@16) to meaningful accuracy ( 5.0 % 5.0\% ).

(2) Best@16 reveals substantial capability transfer. While Avg@16 gains are modest, Best@16 tells a more encouraging story: the student achieves 13.3 % 13.3\% on both AIME24 and AIME25 (up from 6.7 % 6.7\% and 3.3 % 3.3\% ), and 37.3 % 37.3\% on AIMO (up from 3.5 % 3.5\% ), representing a 4.7 × 4.7\times improvement in average Best@16 ( 21.3 % 21.3\% vs 4.5 % 4.5\% ). This indicates that the bridge successfully transfers the teacher’s reasoning capability into the student’s backbone, even across a tokenizer boundary. The gap between Best@16 and Avg@16 is consistent with the LM-head calibration hypothesis from § 4.2.2 : the backbone has acquired teacher-like representations, but the LM head (which was pre-trained with a different tokenizer’s embedding table) has not been calibrated to translate these representations into consistent token-level outputs.

This experiment validates the core theoretical promise of OPRD-Bridge: by shifting alignment from the output space to the representation space, distillation becomes decoupled from the vocabulary, enabling knowledge transfer across incompatible tokenizers.

## 5 Discussion

##### From representational alignment to cross-architecture distillation.

The unifying insight behind OPRD is that the quality of hidden-state distillation is governed by representational alignment between teacher and student (§ 3.3 ). When alignment is high (same architecture, shared initialisation), the full ℝ d \mathbb{R}^{d} channel is already an effective all-pass filter and OPRD-Vanilla suffices. When alignment is low (different depth, width, or training history), OPRD-Bridge constructs a low-rank, low-pass channel that transmits the shared structure and suppresses architecture-specific noise. Both regimes are instances of the same principle; the bridge rank r r simply controls the bandwidth of the supervision channel.

##### High-value application 1: multi-model RL merging.

In the same-architecture setting (e.g., multi-model RL merging where all models share a backbone), OPRD-Vanilla addresses a pressing practical pain point. In large-scale RL pipelines that merge multiple reward models or policy checkpoints, full-vocabulary OPD is the natural distillation objective but incurs prohibitive memory cost: materialising the [ B , T , | 𝒱 | ] [B,T,|\mathcal{V}|] logit tensor for | 𝒱 | |\mathcal{V}| demands extremely high transient GPU memory, often requiring extensive infrastructure modifications [ 5 ] . The common workaround, top- k k OPD, reduces memory but introduces a truncation bias (tail tokens are ignored) and remains subject to the LM-head information bottleneck analyzed in Theorem 2 . OPRD offers a third path: it simultaneously mitigates the variance problem (by providing a deterministic, hidden-state-level gradient) and dramatically reduces memory and wall-clock cost (by never materialising the vocabulary-sized tensor), making it an attractive drop-in component for multi-model RL consolidation.

##### High-value application 2: on-policy self-distillation (OPSD).

OPRD is a natural fit for on-policy self-distillation , where the teacher is constructed from the student itself by injecting privileged information (e.g., ground-truth solutions, step-level verification signals) into the prompt. Because the teacher and student share exactly the same weights, the same-architecture requirement is satisfied by construction, and the hidden-state alignment signal is maximally informative. In this setting OPRD can replace the reverse-KL computation in the output space with a cheaper and lower-variance representation-level objective, while retaining the full benefit of privileged-information guidance.

##### Toward alignment-aware pre-training.

Our framework reveals a clear division of labour: OPRD-Vanilla exploits existing representational alignment, while OPRD-Bridge constructs it post hoc via a low-rank projection. A natural question is whether the alignment itself could be established earlier, during pre-training, so that the bridge becomes unnecessary. Concretely, if a large model (e.g., Qwen3-32B) and a small model (e.g., Qwen3-1.7B) were pre-trained with an explicit alignment objective, for instance through shared initialisation of common layers, periodic representation-matching regularisation, or co-distillation, then the resulting pair would already inhabit a well-aligned representation space despite differing in depth and width. Post-training distillation could then proceed via OPRD-Vanilla (with a simple dimension adapter when d S ≠ d T d_{S}\neq d_{T} ), bypassing the bridge entirely: the full ℝ d S \mathbb{R}^{d_{S}} channel would serve as a high-quality all-pass filter, combining the fidelity of full-vocabulary distillation with the efficiency of representation-level supervision. In the language of our filter analogy, alignment-aware pre-training would shift the “cutoff frequency” upward, turning directions that are currently noise (and must be filtered out by the bridge) into signal. This perspective reframes the bridge not as a permanent architectural component but as a compensator for misalignment that better pre-training could eliminate.

## 6 Related Work

OPRD draws on three main lines of research: classical knowledge distillation, on-policy distillation, and feature-level / intermediate-representation distillation. We also discuss adjacent work on auxiliary losses and capacity-gap analyses. Throughout, we emphasize how OPRD differs from prior work that may at first appear similar.

##### Output-Space Knowledge Distillation.

The idea of compressing a large model into a smaller one by matching their output distributions dates back to Hinton et al. [12] . In the sequence-modelling setting, Kim and Rush [20] showed that training a student on teacher-generated translations is an effective form of sequence-level knowledge transfer; subsequent work applied the same principle to pre-trained language models [ 30 , 17 , 35 ] and to instruction-following LLMs via supervised fine-tuning on teacher rollouts [ 4 , 31 , 36 ] . A common thread across all these methods is that supervision is provided (i) off-policy , on data the student did not generate, and (ii) exclusively in the output space , at the LM-head logits or the softmax distribution derived from them. The first property introduces exposure bias [ 2 ] ; the second confines the learning signal to the ill-conditioned image of W head W_{\mathrm{head}} , leaving hidden-state deviations along its effective null space entirely unpenalised ( Theorem 2 ). OPRD departs from both properties simultaneously.

##### On-Policy Distillation.

The exposure-bias problem motivated a shift toward on-policy training. MiniLLM [ 8 ] optimised a reverse-KL objective on student-sampled responses via policy gradient, observing that the mode-seeking property of reverse KL discourages the student from placing mass where the teacher assigns low probability. GKD [ 1 ] generalised this to a family of divergences that interpolate between on- and off-policy data. More recently, Yang et al. [41] reinterpreted OPD through the lens of KL-constrained RL, revealing that the teacher’s per-token log-probability ratio serves as an implicit dense reward. Building on these foundations, OPD has been adopted in several production post-training pipelines [ 5 , 39 , 47 , 37 , 21 , 18 , 16 , 7 , 13 ] and extended to self-distillation settings where the teacher is derived from the student itself via privileged information [ 14 , 48 , 11 , 32 , 43 , 29 , 19 , 42 , 40 , 24 , 6 ] . A key observation, however, is that the entire design space explored so far (sampled-token, top- k k , and full-vocabulary variants) concerns only how many output tokens to supervise per position; the supervision itself never leaves the output space. OPRD is, to our knowledge, the first on-policy method whose learning signal originates strictly before the LM head , operating on the student’s own trajectories.

##### Feature / Intermediate-Representation Distillation.

A separate line of work supervises the student’s intermediate representations rather than its outputs. Early instances include FitNets [ 28 ] , which match a single “hint” layer of the student to the teacher; attention-transfer [ 46 ] , which matches per-pixel attention maps in CNNs; and FSP-matrix distillation [ 44 ] , which matches Gram matrices between layers. For BERT-style language models, TinyBERT [ 17 ] and MobileBERT [ 33 ] extend this idea by jointly matching hidden states and attention maps across all layers, and MiniLM/MiniLMv2 [ 35 , 34 ] match self-attention relation matrices. At first glance, OPRD may look like a straightforward port of these ideas to autoregressive LLMs, but two structural differences set it apart:

• On-policy vs. off-policy supervision. FitNets, TinyBERT, MiniLM, and their successors compute the feature-matching loss on fixed inputs from a pre-training or downstream corpus, i.e. inputs the student does not generate. The student is never exposed to its own rollout distribution during distillation, so exposure bias remains. OPRD, by contrast, computes the hidden-state loss on student-generated sequences y ^ ∼ π θ ( ⋅ ∣ x ) \hat{y}\sim\pi_{\theta}(\cdot\mid x) that evolve as training progresses. The teacher is queried on states the student actually visits, making the supervision signal adaptive to the student’s evolving policy.

• Encoder representations vs. autoregressive prefix representations. Prior feature-distillation work targets encoder models (BERT, vision CNNs) whose representations are computed once per input; the teacher and student process the same input and are aligned post-hoc. In the autoregressive LLM setting, each hidden state h ⋅ , t ( l ) h_{\cdot,t}^{(l)} (for either model) encodes the model’s belief just before predicting token y ^ t \hat{y}_{t} , conditional on the entire sampled prefix y ^ < t \hat{y}_{<t} . OPRD therefore aligns the student’s predictive computation at every decoding step under its own sampling distribution, a fundamentally on-policy object with no analog in encoder-style feature distillation.

##### Hint Learning, Auxiliary Losses, and Distribution Matching.

A related body of work uses intermediate signals to regularize or augment training rather than to distill from a separate teacher. Deeply-supervised nets [ 23 ] attach auxiliary classifiers to intermediate layers of a single model; DINO [ 3 ] aligns hidden states across augmented views of the same input in self-supervised learning; representation engineering [ 49 ] steers or interprets hidden states without explicit teacher supervision. OPRD shares the high-level intuition that intermediate representations carry useful signal, but differs in three crucial ways: it is (i) explicitly teacher–student rather than self-supervised, (ii) on-policy on student-generated autoregressive trajectories, and (iii) a self-contained training objective that can additionally compose with any output-space OPD variant via Eq. ( 7 ).

## 7 Conclusion and Future Work

We presented OPRD , the first on-policy distillation method that supervises the student in the hidden-state space rather than at the LM-head output. The central thesis is that all existing OPD variants (sampled-token, top- k k , and full-vocabulary) share two practical limitations inherent to the output-space paradigm: the dominant sampled-token variant suffers from a high-variance REINFORCE-style gradient estimator whose signal-to-noise ratio collapses as the student approaches the teacher, while top- k k and full-vocabulary variants trade this variance for a truncation bias or prohibitive memory cost; and all variants are subject to an LM-head projection that acts as an information bottleneck, compressing the teacher’s full stack of intermediate hidden states through an ill-conditioned, softmax-invariant mapping. By moving supervision from the output of the LM head to its input, OPRD yields a deterministic per-sample gradient that removes the token-level estimation variance by construction, and exposes per-position, per-layer structural information that any output-space objective necessarily discards. Empirically, OPRD enables monotonic improvement throughout training and closes the student–teacher gap on three competition mathematics benchmarks (AIME 2024, AIME 2025, AIMO), while every output-space baseline plateaus several points below the teacher. On the same hardware budget, OPRD is strictly Pareto-dominant: 1.44 × 1.44\times faster wall-clock training and up to 54 % 54\% less actor-update transient memory than top- k k OPD, because its loss path never materialises the [ B , T , | 𝒱 | ] [B,T,|\mathcal{V}|] logits tensor.

We further introduced OPRD-Bridge , which extends representation distillation to the cross-architecture setting by constructing a frozen low-rank bridge between heterogeneous teacher and student representations. By exploiting the empirical finding that models of different depth and width share a low-rank representational structure, OPRD-Bridge shifts alignment from the output space to the representation space, decoupling distillation from the vocabulary. We validated this capability on both cross-architecture (Qwen3-4B → \to Qwen3-1.7B) and cross-tokenizer (Phi-4-mini-reasoning → \to Qwen3-1.7B) settings, demonstrating successful knowledge transfer even when the vocabulary-based alignment channel is unavailable.

##### Future Work.

Several directions follow naturally from our framework: • Beyond mathematical reasoning. Our experiments focus on long-CoT math benchmarks. Whether OPRD’s gains transfer to code generation, agentic interaction, and open-ended dialogue, each with different position-level supervision characteristics, remains an open question.

• Cross-modal distillation. Our cross-tokenizer experiment (§ 4.2.4 ) validates that OPRD-Bridge transfers knowledge across incompatible tokenizers by aligning in representation space. The natural next step is cross-modal distillation (e.g., from a vision-language teacher to a language-only student), where the representation-space bridge may serve as the only viable supervision channel.

• Adaptive layer and position selection. We use uniform layer weighting and a simple last- k k position heuristic. Adaptively weighting layers and positions based on where the student–teacher gap is largest or where the gradient signal is most informative could further sharpen supervision.

• On-policy representation self-distillation (OPRSD). As discussed in § 5 , OPRD is a natural fit for self-distillation with privileged information, where the same-architecture requirement is satisfied by construction. Scaling OPRSD to multi-turn and multi-task settings is a promising next step.

• Understanding the phase transition. Our mechanistic analysis (§ 4.1.5 ) reveals a PG-loss spike and associated entropy/overlap dynamics when OPRD is active. Characterising the mechanism behind this transition would deepen the theoretical understanding of representation-level distillation.

• Attention-map distillation. OPRD aligns hidden-state vectors but does not supervise the attention patterns that produce them. Extending OPRD with an on-policy attention-matching objective could transfer the teacher’s routing and composition behaviour more directly.