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

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

Figure 9 tracks actor/pg_loss along training for the OPD top- 1 1 + + OPRD composite runs with μ ∈ { 0 , 1 , 10 } \mu\!\in\!\{0,1,10\} . Two observations stand out. First , all three runs exhibit a pronounced loss spike during training, likely reflecting a phase transition in the student’s policy as it reorganises to absorb the teacher’s behaviour (the precise mechanism is under active investigation). Crucially, adding OPRD causes this spike to arrive earlier : the μ = 1 \mu\!=\!1 and μ = 10 \mu\!=\!10 spikes precede the μ = 0 \mu\!=\!0 spike, indicating that hidden-state supervision accelerates the distillation dynamics. Second , after the spike all three curves converge to approximately zero PG loss in late training, yet the accuracy gap persists ( + 5.4 +5.4 and + 7.9 +7.9 pt over μ = 0 \mu\!=\!0 on AIME24). This directly corroborates Theorem 2 : once the policy gradient vanishes ( p t ≈ q t p_{t}\approx q_{t} ), the output-space OPD signal can no longer drive further improvement because the remaining student–teacher gap lives in the null space 𝒩 W \mathcal{N}_{W} of the LM head; only OPRD’s representation-level signal, which bypasses this bottleneck, continues to make progress.

##### (b) Top- 16 16 overlap: hidden-state alignment propagates to next-token agreement.

Following Li et al. [25] , who show that higher student–teacher top- k k overlap is a reliable predictor of distillation quality, we log val-topk/overlap_ratio , defined as | top- ​ 16 ​ ( π θ ) ∩ top- ​ 16 ​ ( π T ) | / 16 |\text{top-}16(\pi_{\theta})\cap\text{top-}16(\pi_{T})|/16 (higher is better; 1.0 1.0 means the student’s top- 16 16 set is identical to the teacher’s). Figure 10 compares OPD top- 16 16 alone against OPD top- 16 16 + ⋅ ℒ OPRD +1\!\cdot\!\mathcal{L}_{\mathrm{OPRD}} .

The OPD-only run increases the overlap nearly monotonically throughout training, but its rate of improvement visibly slows after mid-training. The OPD + + OPRD run behaves differently: it initially rises alongside OPD-only, then undergoes a sudden dip in overlap (temporally coinciding with the PG-loss spike of Figure 9 , consistent with the hypothesised phase transition), after which it rebounds rapidly and surpasses the OPD-only curve by a clear margin. The dip-then-surge pattern mirrors the PG-loss spike discussed above and is consistent with OPRD driving the student through a transient reorganisation that ultimately lands it in a higher-overlap regime than OPD alone can reach. The representation-level and output-level supervisions are therefore not redundant: hidden-state alignment translates back into measurable improvement on exactly the metric OPD top- 16 16 was designed to optimise.

##### (c) Predictive entropy.

We log actor/entropy and teacher/entropy (per-token Shannon entropy of π θ \pi_{\theta} and π T \pi_{T} on the same rollout positions) along training for the same OPD top- 1 1 + + OPRD composite runs ( μ ∈ { 0 , 1 , 10 } \mu\!\in\!\{0,1,10\} ). The teacher’s entropy curve serves as a reference: since π T \pi_{T} is frozen, any drift is induced purely by the changing rollout distribution. Figure 11 shows the three runs side by side. All three runs eventually bring the student’s entropy into close agreement with the teacher’s by the end of training (note that the teacher’s entropy also drifts upward as the rollout distribution evolves, so alignment means tracking the teacher, not returning to a fixed level). However, they differ markedly in their early dynamics: each run exhibits an entropy-increase phase in which the student–teacher gap widens before narrowing. Adding OPRD causes this entropy-increase phase to begin earlier , temporally coinciding with the PG-loss spike of Figure 9 : the μ = 10 \mu\!=\!10 onset precedes the μ = 1 \mu\!=\!1 onset, which in turn precedes the μ = 0 \mu\!=\!0 onset. This is consistent with the picture that OPRD accelerates the student’s internal reorganisation (the same phase transition visible in the PG-loss and overlap diagnostics), after which the student’s entropy converges to the teacher’s more quickly.

### 4.2 OPRD-Bridge: Cross-Architecture and Cross-Tokenizer Distillation

We now evaluate OPRD-Bridge in settings where direct hidden-state comparison is impossible without the bridge mechanism introduced in § 3.3 : first in the cross-architecture setting, where the student and teacher differ in depth and width but share a vocabulary; then in the cross-tokenizer setting, where the two models additionally use completely disjoint tokenizers.

#### 4.2.1 Experimental Setup

##### Models.

We use Qwen3-4B [ 39 ] as the (frozen) teacher and Qwen3-1.7B-Base [ 39 ] as the student. The teacher has L T = 36 L_{T}\!=\!36 transformer layers with hidden dimension d T = 2560 d_{T}\!=\!2560 ; the student has L S = 28 L_{S}\!=\!28 layers with d S = 2048 d_{S}\!=\!2048 . Both share the same | 𝒱 | ≈ 151 |\mathcal{V}|\!\approx\!151 K vocabulary, but differ in depth ( 29 % 29\% more layers) and width ( 25 % 25\% wider hidden states), making naïve OPRD-Vanilla inapplicable ( d S ≠ d T d_{S}\neq d_{T} , L S ≠ L T L_{S}\neq L_{T} ). In the cross-tokenizer experiment (§ 4.2.4 ), we additionally test with Phi-4-mini-reasoning [ 38 ] ( L T = 32 L_{T}\!=\!32 , d T = 3072 d_{T}\!=\!3072 , | 𝒱 | ≈ 200 |\mathcal{V}|\!\approx\!200 K, tiktoken-based tokenizer) as teacher, paired with the same Qwen3-1.7B-Base student, where the two models have completely disjoint tokenizers.

##### Bridge construction (Stage 1).

We sample 2000 2000 prompts from DAPO-Math-17K, generate student rollouts, and collect both models’ hidden states. The teacher projector P T ( l ) ∈ ℝ r × d T P_{T}^{(l)}\in\mathbb{R}^{r\times d_{T}} is obtained via PCA (eigendecomposition of the hidden-state covariance) at each layer and permanently frozen. The student projector P S ( l ) ∈ ℝ r × d S P_{S}^{(l)}\in\mathbb{R}^{r\times d_{S}} (linear) is trained for 20 20 epochs to minimize Eq. ( 13 ) with both model backbones frozen, then also frozen. Layer correspondence follows proportional spacing (Eq. ( 11 )). Unless otherwise stated, we use rank r = 8 r\!=\!8 .

##### Distillation (Stage 2).

On-policy prompts are drawn from DAPO-Math-17K. For each prompt the student samples 2 2 responses at temperature 1.0 1.0 with max length 16,384 16{,}384 tokens; global batch size is 8 8 prompts per step. We supervise all 28 28 student layers and apply the position mask to the last 2000 2000 response tokens. The loss is normalized MSE (L2-normalize both projected vectors before MSE). By default we use OPRD-Bridge only. All methods are trained with AdamW (peak learning rate × 10 − 5 1\!\times\!10^{-5} , cosine decay), bf16 mixed precision, and FSDP over × 8\!\times\! A100 (80G) GPUs.

##### Baselines.

We compare against the unmodified Qwen3-1.7B-Base checkpoint (no distillation), OPD top-1 (sampled-token reverse KL), and OPD top-16. Note that output-space OPD baselines can be applied cross-architecture because they only require a shared vocabulary; they therefore serve as the natural comparison for OPRD-Bridge.

##### Evaluation.

We evaluate at decoding temperature 0.7 0.7 with 16 16 independently sampled responses per prompt on three competition-level mathematical reasoning benchmarks: AIME 2024 ( 30 30 problems), AIME 2025 ( 30 30 problems), and AIMO (AMC 2022/2023, 83 83 problems). We report four metrics: Avg@16 (average accuracy across the 16 16 samples), Best@16 (accuracy of the best sample, measuring the capability ceiling), Dist-4g (ratio of unique 4-grams to total 4-grams, measuring lexical diversity), and Avg Len (mean response length in tokens, measuring generation efficiency). Final answers are extracted with the standard boxed parser and graded by exact-match. We save checkpoints periodically and report results from the checkpoint that achieves the best average performance across AIME24, AIME25, and AIMO.

#### 4.2.2 Main Results

Tables 4 and 5 report the main cross-architecture distillation results. Three observations emerge from different facets of the evaluation.

(1) OPRD-Bridge matches OPD’s capability ceiling (Best@16). While OPRD-Bridge’s Avg@16 is lower than OPD top-1 on some benchmarks, its Best@16 tells a different story: on AIME24 and AIME25, OPRD-Bridge achieves 20.0 20.0 and 13.3 13.3 , matching OPD top-1 exactly and substantially outperforming OPD top-16. Best@16 reflects the upper bound of the student’s capability after distillation (the best answer the model can produce across 16 16 attempts) and is therefore a more direct measure of how much teacher knowledge has been successfully transferred. The parity in Best@16 indicates that OPRD-Bridge raises the student’s capability ceiling to the same level as output-space distillation, despite operating through a rank- 8 8 subspace rather than the full | 𝒱 | |\mathcal{V}| -dimensional vocabulary simplex.

(2) OPRD-Bridge produces substantially more diverse reasoning (Dist-4g). OPRD-Bridge achieves Dist-4g scores of 22.6 22.6 / 17.2 17.2 / 21.3 21.3 on AIME24/AIME25/AIMO, dramatically higher than both OPD top-1 ( 6.2 6.2 / 6.4 6.4 / 15.6 15.6 ) and OPD top-16 ( 9.9 9.9 / 11.7 11.7 / 13.6 13.6 ), and approaching the teacher’s own diversity levels ( 24.0 24.0 / 25.4 25.4 / 41.6 41.6 ). This indicates that hidden-state supervision through the bridge produces reasoning chains with far less repetition than output-space distillation. The OPD baselines, by contrast, appear to induce repetitive patterns, likely because the high-variance output-space gradient causes the student to “loop” through similar token sequences rather than making steady forward progress in its reasoning.

(3) OPRD-Bridge generates much shorter, more efficient responses (Avg Len). OPRD-Bridge converges to mean response lengths of 5,909 5{,}909 / 5,234 5{,}234 / 4,855 4{,}855 tokens, which are 2 2 – 3 × 3\times shorter than OPD top-1 ( 19,120 19{,}120 / 17,931 17{,}931 / 11,008 11{,}008 ) and 1.5 1.5 – 2.4 × 2.4\times shorter than OPD top-16 ( 12,222 12{,}222 / 10,175 10{,}175 / 11,833 11{,}833 ). Remarkably, OPRD-Bridge’s response lengths are close to the teacher’s own ( 6,542 6{,}542 / 4,426 4{,}426 / 2,384 2{,}384 ), suggesting that the bridge successfully transfers the teacher’s concise reasoning style in addition to its knowledge. Combined with the Best@16 parity, this means OPRD-Bridge achieves the same capability ceiling at a fraction of the inference cost, a significant practical advantage for deployment.

Summary. The Avg@16 gap between OPRD-Bridge and OPD reflects a difference in consistency (how often the model reaches its best answer), not in capability (what the best answer is). OPRD-Bridge compensates with dramatically better diversity and efficiency: it produces teacher-like reasoning chains that are concise, non-repetitive, and reach the same performance ceiling. This profile (matching capability ceiling with superior efficiency) makes OPRD-Bridge particularly attractive for inference-constrained deployment scenarios where shorter, more reliable responses are preferred over longer, repetitive ones.

(4) Composing OPRD-Bridge with OPD further raises the capability ceiling. The Avg@16 gap of OPRD-Bridge relative to OPD has a clear mechanistic explanation: OPRD-Bridge only trains the student backbone (hidden states), while the LM head W head W_{\mathrm{head}} (the final mapping from hidden states to token distributions) receives no gradient. Even when the backbone representations are well-aligned with the teacher’s, the untrained LM head may not translate this alignment into precise per-token output distributions, reducing sampling consistency. OPD, by contrast, directly supervises the output distribution and therefore calibrates the LM head implicitly through backpropagation.

This suggests a natural two-stage recipe: first use OPRD-Bridge to align the backbone representations, then fine-tune with OPD to calibrate the output distribution. The last row of Tables 4 and 5 (“ ↪ \hookrightarrow OPD top-1”) reports exactly this: we take the OPRD-Bridge checkpoint and continue training with OPD top-1. The results confirm the complementarity: (i) Best@16 rises from 30.8 30.8 to 33.9 \mathbf{33.9} (averaged across benchmarks), surpassing both standalone OPD top-1 ( 30.8 30.8 ) and OPD top-16 ( 27.1 27.1 ) and establishing a new best among all methods. This validates that OPRD-Bridge pre-aligns the backbone into a better initialisation from which OPD can extract more capability. (ii) Avg@16 improves from 11.2 11.2 to 13.6 13.6 , partially closing the gap to standalone OPD top-1 ( 15.3 15.3 ), consistent with the LM-head calibration hypothesis. (iii) The generation quality metrics shift toward a middle ground: Dist-4g decreases from 20.4 20.4 to 15.0 15.0 (still above OPD’s 9.4 9.4 – 11.7 11.7 ) and Avg Len increases from 5,333 5{,}333 to 7,169 7{,}169 (still well below OPD’s 11,410 11{,}410 – 16,020 16{,}020 ), reflecting the trade-off between output-level calibration and the concise reasoning style inherited from the bridge stage.

The two-stage pipeline thus combines the strengths of both approaches: OPRD-Bridge provides a high-quality backbone initialisation with teacher-like reasoning structure, and OPD fine-tunes the output distribution to improve consistency, yielding the highest capability ceiling overall.

#### 4.2.3 Ablation Studies

We ablate the key design choices of OPRD-Bridge to understand their individual contributions.

##### Effect of bridge rank r r .

Figure 12 reports the average cosine similarity between P S ​ ( h S ) P_{S}(h_{S}) and P T ​ ( h T − μ T ) P_{T}(h_{T}-\mu_{T}) across all layer pairs, before and after Stage 1 bridge training, for ranks r ∈ { 1 , 2 , 4 , 8 , 16 , 32 , 64 , 128 , 256 , 512 , 1024 , 2048 } r\in\{1,2,4,8,16,32,64,128,256,512,1024,2048\} . Three findings emerge.

(i) The bridge training is essential. Before Stage 1, the projected representations are nearly orthogonal at all ranks (cosine similarity ≤ 7.5 % \leq 7.5\% ), confirming that the two models’ hidden states are incommensurable without the learned alignment. After training, similarity jumps to 72.3 % 72.3\% even at rank 1 1 and exceeds 90 % 90\% for r ≥ 4 r\geq 4 .

(ii) The optimal rank is low. Similarity peaks at r = 8 r\!=\!8 ( 95.0 % 95.0\% ) and then declines monotonically: 94.1 % 94.1\% at r = 32 r\!=\!32 , 87.9 % 87.9\% at r = 256 r\!=\!256 , and 77.0 % 77.0\% at full rank r = 2048 r\!=\!2048 . This validates the low-pass filter interpretation from § 3.3.4 : the top- r r principal components correspond to the “low-frequency” modes where teacher and student are highly aligned; beyond this sweet spot, additional dimensions introduce architecture-specific variation (“high-frequency noise”) that the linear projector cannot align, diluting the overall similarity.

(iii) The rank controls the signal-to-noise trade-off. Too small a rank ( r = 1 r\!=\!1 – 2 2 ) under-represents the shared structure, leaving useful information on the table. Too large a rank ( r ≥ 64 r\geq 64 ) forces the bridge to transmit misaligned directions, degrading the supervision signal. The plateau at r = 4 r\!=\!4 – 16 16 ( 93.7 93.7 – 95.0 % 95.0\% ) defines the practical operating range; we default to r = 8 r\!=\!8 throughout.

#### 4.2.4 Cross-Tokenizer Distillation

Since OPRD-Bridge performs alignment in representation space, it remains applicable even when teacher and student use entirely different tokenizers. We verify this by distilling Phi-4-mini-reasoning (Microsoft, | 𝒱 | ≈ 200 |\mathcal{V}|\!\approx\!200 K, tiktoken-based BPE; L T = 32 L_{T}\!=\!32 , d T = 3072 d_{T}\!=\!3072 ) into Qwen3-1.7B-Base ( | 𝒱 | ≈ 151 |\mathcal{V}|\!\approx\!151 K, Qwen BPE; L S = 28 L_{S}\!=\!28 , d S = 2048 d_{S}\!=\!2048 ). All training details follow § 4.2.1 exactly, except that we set bridge rank r = 4 r\!=\!4 .

Table 6 reports the results. Two observations stand out.