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

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

##### Distillation objectives.

We compare three on-policy distillation variants, all sharing the same rollouts y ^ \hat{y} and the same single teacher forward pass per rollout: • OPD top-1 (sampled-token reverse KL): the per-position estimator ℓ t = log ⁡ p t ​ ( y ^ t ) − log ⁡ q t ​ ( y ^ t ) \ell_{t}=\log p_{t}(\hat{y}_{t})-\log q_{t}(\hat{y}_{t}) evaluated only at the sampled token y ^ t \hat{y}_{t} .

• OPD top-16 : the per-position estimator ∑ v ∈ 𝒱 16 t p t ​ ( v ) ​ [ log ⁡ p t ​ ( v ) − log ⁡ q t ​ ( v ) ] \sum_{v\in\mathcal{V}_{16}^{t}}\,p_{t}(v)\,[\log p_{t}(v)-\log q_{t}(v)] over the top- 16 16 tokens of p t p_{t} , a strictly informative-superset of OPD top-1 .

• OPRD-Vanilla (ours): the hidden-state objective ( 6 ) with ℒ layer = { 1 , … , L } \mathcal{L}_{\mathrm{layer}}\!=\!\{1,\ldots,L\} (all 28 28 layers) and 𝒫 ⁡ ( y ^ ) \mathcal{P}(\hat{y}) set to the last k = 2000 k\!=\!2000 response tokens (i.e. the suffix in which the chain-of-thought converges to a final answer); reported in the OPRD-only setting ( μ = 0 \mu\!=\!0 in ( 7 )).

##### Optimization.

All three methods are trained for 500 500 optimizer steps with AdamW (peak learning rate × 10 − 5 1\!\times\!10^{-5} , linear warm-up over 3 % 3\% of total steps, cosine decay), bf16 mixed precision, and FSDP over × 8\!\times\! A100 (80G) GPUs at a micro-batch of B = 8 B\!=\!8 and a maximum response length of T = 16,384 T\!=\!16{,}384 .

##### Evaluation.

We report Avg@16 (average accuracy across 16 16 independently sampled responses per prompt) at decoding temperature 0.7 0.7 on three competition-level mathematical reasoning benchmarks: AIME 2024 ( AIME24 , 30 30 problems), AIME 2025 ( AIME25 , 30 30 problems), and AIMO (AI-MO/aimo-validation-amc, comprising AMC 2022 and AMC 2023, 83 83 problems). Final answers are extracted with the standard boxed parser and graded by exact-match against the official solution.

#### 4.1.2 Main Results

Table 2 reports Avg@16 for the teacher, the unmodified student, and the three on-policy distillation methods; Figure 4 shows the corresponding training dynamics (discussed in detail in § 4.1.3 ). Three observations follow.

(1) Both OPD variants improve over the student but plateau noticeably below the teacher. The student starts from a 17.9 17.9 -/ 13.7 13.7 -/ 17.3 17.3 -point gap to the teacher on AIME24/AIME25/AIMO. OPD top-1 closes most of this gap on AIME25 (to within 2.1 2.1 points) but leaves 8.5 8.5 / 2.5 2.5 points on AIME24/AIMO; enriching the supervision to OPD top-16 helps substantially on AIME24 ( + 4.8 +4.8 ) and marginally on AIME25 ( + 0.5 +0.5 ) yet loses ground on AIMO ( − 0.5 -0.5 ). The absence of a clean ordering between top-1 and top-16 (more tokens in the loss is supposed to be strictly more informative) suggests that the output-space paradigm itself is the bottleneck: both variants are constrained by the LM-head information bottleneck ( Theorem 2 ), and top- k k ’s truncation bias means that enriching the support does not guarantee monotonic improvement.

(2) OPRD effectively closes the student–teacher gap. OPRD reaches 49.8 49.8 on AIME24, 34.6 34.6 on AIME25, and 79.1 79.1 on AIMO, leaving only 1.0 1.0 / 1.0 1.0 / 0.4 0.4 points to the teacher, all within the variance of 16 16 -sample Avg@16 evaluation, so the AIMO result is effectively a tie with the teacher (underlined in Table 2 ). Relative to the better OPD baseline on each benchmark, OPRD gains + 2.7 +2.7 / + 0.6 +0.6 / + 2.1 +2.1 points; relative to the unmodified student it gains + 16.9 +16.9 / + 12.7 +12.7 / + 16.9 +16.9 points. The advantage is most striking on AIMO, where OPRD recovers essentially all of the 17.3 17.3 -point student–teacher gap that no output-space variant fully bridges. OPRD’s gradient is conditionally deterministic ( Theorem 1 ), avoiding the late-stage variance collapse that limits OPD; it also exposes per-layer structural information that the LM-head projection compresses away ( Theorem 2 ), supervising directions in 𝒩 W \mathcal{N}_{W} that any output-space objective treats as invisible.

#### 4.1.3 Training Dynamics

The end-of-training numbers in Table 2 are only one slice of the story; we now examine how each method gets there. Three complementary views (per-step accuracy, response-length behaviour, and OPRD’s own internal alignment metric) together paint a consistent picture of OPD stalling in the late-training regime predicted by Theorem 1 , while OPRD continues to make progress.

##### Accuracy curves: OPRD climbs monotonically, OPD plateaus.

Figure 4 compares OPRD step-by-step against OPD top- 1 1 (top row) and OPD top- 16 16 (bottom row) on all three benchmarks; raw curves are drawn at α = 0.5 \alpha\!=\!0.5 and the solid curve with markers is the 5 5 -step centred rolling mean. The two methods in each panel share the same initialisation and quickly enter qualitatively different regimes: both OPD variants lift accuracy in the first few dozen steps but then plateau or oscillate without further improvement , whereas OPRD continues to climb essentially monotonically until it reaches the teacher level. Enriching the OPD supervision from top- 1 1 to top- 16 16 narrows the asymptotic gap to OPRD on AIME24 but does not change the qualitative shape: OPD top- 16 16 also plateaus, and on AIMO it does so ∼ 2.6 \sim\!2.6 points below OPRD despite passing strictly more output-distribution information into the loss. This is the SNR-collapse prediction of Theorem 1 in pictures: as p t → q t p_{t}\to q_{t} , the OPD gradient’s signal-to-noise ratio collapses and additional output-layer information cannot rescue the per-token sampling noise; only OPRD’s deterministic, hidden-state-level signal continues to make progress.

##### Behavioural view: OPRD produces shorter, more efficient reasoning.

The accuracy curves answer whether a method keeps improving; a complementary question is how the policy changes. Figure 5 reports the mean rollout length response_length/mean for the same three runs. OPRD converges to a mean response length of ∼ 5,700 {\sim}5{,}700 tokens, substantially shorter than the ∼ 7,000 {\sim}7{,}000 tokens produced by both OPD variants. Combined with OPRD’s higher accuracy ( Table 2 ), this indicates that hidden-state supervision guides the student toward more concise reasoning chains: the student learns to reach the correct answer with fewer tokens rather than relying on longer, less directed exploration. This also translates to a practical inference-time efficiency gain, since shorter responses require proportionally less compute at deployment.

##### Internal view: OPRD’s own loss is being optimised end-to-end.

A final, internal diagnostic is whether the representation-level loss OPRD is supposed to minimise actually decreases along training. Figure 6 plots rep/cosine_similarity , the cosine similarity between π θ \pi_{\theta} ’s and π T \pi_{T} ’s hidden states averaged across all transformer layers and OPRD-supervised positions, for the OPRD-only run from Table 2 . The curve rises sharply in the first few dozen steps and then drifts upward steadily for the rest of training. Two consequences follow: (i) the OPRD objective is well-conditioned for end-to-end optimisation at this scale: the gradient produced by ( 6 ) is consistent enough to monotonically pull the supervised hidden states towards the teacher’s; (ii) the downstream gains of Table 2 are matched by a corresponding internal trend: OPRD is improving on exactly the quantity its loss is defined on, confirming that the improvement, not a coincidental rollout-distribution shift, drives the gains.

#### 4.1.4 Efficiency

The OPRD loss path is computed entirely before the LM head: it never materializes the [ B , T , | 𝒱 | ] [B,T,|\mathcal{V}|] logits tensor on the student side, never invokes the | 𝒱 | |\mathcal{V}| -way log_softmax , and never backpropagates through W head ∈ ℝ | 𝒱 | × d W_{\mathrm{head}}\in\mathbb{R}^{|\mathcal{V}|\times d} for the distillation term. As a direct consequence, OPRD-only training is strictly cheaper than any output-space OPD variant at the same rollout/teacher budget. Table 3 quantifies this on the same training configuration as our main results.

Memory. The actor-update transient footprint ( Δ \Delta peak, the most direct proxy for the loss path’s own cost since always-resident state is subtracted out) is 30.2 30.2 GB for OPD top-1 and 45.0 45.0 GB for OPD top-16, vs. only 20.5 20.5 GB for OPRD, a 32 32 % and 54 54 % reduction, or equivalently a 1.47 × 1.47\times and 2.20 × 2.20\times ratio. The gap is dominated by the [ B , T , | 𝒱 | ] [B,T,|\mathcal{V}|] logits tensor (and its gradient buffer for top- k k ), which scales with | 𝒱 | ≈ 151 |\mathcal{V}|\!\approx\!151 K but is entirely absent in the OPRD-only loss path. The roughly 10 10 – 25 25 GB of saved transient memory is hardware-relevant: on 80 80 GB-class accelerators it is enough to either enlarge the micro-batch or extend the context at the same hardware budget.

Wall-clock. At identical schedules ( 500 500 steps, same rollout, same teacher forward pass), OPRD finishes in 563 563 minutes vs. 813 813 / 812 812 minutes for OPD top-1 / top-16, a 31 31 % wall-clock reduction, equivalent to a 1.44 × 1.44\times speed-up. We attribute this to the fact that the two OPD variants take essentially the same time, consistent with the observation that the cost is dominated by the [ B , T , | 𝒱 | ] [B,T,|\mathcal{V}|] matrix multiplication and log_softmax rather than by the top- k k slicing itself.

Putting it together. Combining Table 3 with the accuracy results in Table 2 , OPRD strictly Pareto-dominates both OPD baselines on this benchmark suite: at ∼ 69 % \sim\!69\% of the wall-clock and 46 46 – 68 % 68\% of the actor-update transient memory ( Δ \Delta peak), it reaches accuracies that are + 0.6 +0.6 to + 2.7 +2.7 points above the better OPD baseline and effectively close the gap to the teacher. These efficiency gains are a secondary consequence of OPRD’s design (the primary motivation, as developed in § 3.2 , is informational; see Theorem 2 ), but they make OPRD a more economical training objective in practice as well. We note that our current implementation reuses the existing OPD training framework without OPRD-specific infrastructure optimisation (e.g., the teacher still computes and discards the full logits tensor even though OPRD does not consume it). With a dedicated implementation that eliminates these redundant computations, we expect both peak memory and wall-clock to decrease further.

#### 4.1.5 Mechanistic Analysis

The experiments above show that OPRD outperforms OPD; we now ask why . We first study the effect of composing OPRD with OPD via the mixing weight μ \mu , and empirically motivate the choice of supervised positions. We then track three complementary diagnostics along training for the composite runs ℒ OPD + μ ⋅ ℒ OPRD \mathcal{L}_{\mathrm{OPD}}+\mu\cdot\mathcal{L}_{\mathrm{OPRD}} with μ ∈ { 0 , 1 , 10 } \mu\!\in\!\{0,1,10\} to reveal a consistent mechanistic picture: OPRD pre-aligns the student’s hidden states to the teacher’s, which propagates back to (a) a smaller residual policy-gradient signal, (b) higher next-token top- k k agreement, and (c) a student exploration distribution whose shape matches the teacher’s.

##### Composing OPD with OPRD ( μ \mu sweep).

Equation 7 suggests that OPRD can also be added on top of existing OPD objective, rather than used as a standalone replacement. We test this for the simplest output-space baseline, sampled-token OPD (i.e. OPD top- 1 1 ), by training the composite loss ℒ OPD + μ ⋅ ℒ OPRD \mathcal{L}_{\mathrm{OPD}}+\mu\cdot\mathcal{L}_{\mathrm{OPRD}} for μ ∈ { 0 , 1 , 10 } \mu\in\{0,1,10\} , keeping all other knobs identical to § 4.1.1 . Figure 7 shows that AIME24 avg@16 rises monotonically with μ \mu : from the vanilla OPD top- 1 1 baseline at 42.3 42.3 ( μ = 0 \mu\!=\!0 ), to 47.7 47.7 with a light OPRD contribution ( μ = 1 \mu\!=\!1 , + 5.4 +5.4 pt, already exceeding the OPD top- 16 16 baseline of 47.1 47.1 from Table 2 ), to 50.2 50.2 with a stronger contribution ( μ = 10 \mu\!=\!10 , + 2.5 +2.5 pt further, essentially matching the teacher’s 50.8 50.8 ). The trend confirms two things: (i) the hidden-state signal that OPRD exposes is additive to the output-space signal that OPD already uses, consistent with the information-bottleneck view of Theorem 2 ; and (ii) the improvement is monotonic in μ \mu within the swept range, so the composition is robust to the mixing weight and does not require careful tuning. We therefore view ℒ OPD + μ ⋅ ℒ OPRD \mathcal{L}_{\mathrm{OPD}}+\mu\cdot\mathcal{L}_{\mathrm{OPRD}} as a drop-in upgrade for existing OPD pipelines.

##### Where does the student diverge from the teacher? (motivation for last- k k supervision).

A natural design question for OPRD is which response positions the projector 𝒫 ⁡ ( y ^ ) \mathcal{P}(\hat{y}) in ( 6 ) should select. We answer this empirically by directly measuring where along the response the student and teacher representations still disagree. For the student initialisation π θ ( 0 ) = \pi_{\theta}^{(0)}\!=\! R1-distill-1.5B and the teacher π T = \pi_{T}\!=\! JustRL-1.5B , we sample on-policy rollouts from the student, run both models forward on each rollout, and compute the cosine similarity between their last-layer hidden states, restricted to either the first k k or the last k k tokens of every response. Figure 8 reports this similarity as a function of k k .

Two patterns emerge. (i) The early response is already teacher-aligned. The first- k k curve stays above 97 % 97\% for every k ≤ 1600 k\!\leq\!1600 and peaks at 97.69 % 97.69\% at k = 1600 k\!=\!1600 , meaning the prompt-following preamble and the opening of the chain-of-thought are essentially already matched by the student; there is little headroom for representation-level supervision to act on. (ii) The late response is where the gap lives. The last- k k curve starts at only 91.65 % 91.65\% for k = 50 k\!=\!50 and remains ≥ 4 \geq\!4 points below the first- k k curve until k k approaches the full response length, at which point both curves converge to the whole-sequence similarity of 95.42 % 95.42\% by construction. Almost all of the student–teacher representational disagreement is concentrated in the tail of the response, precisely where the chain-of-thought commits to a final answer.

This directly motivates our default choice 𝒫 ⁡ ( y ^ ) = last- ​ k \mathcal{P}(\hat{y})\!=\!\textsc{last-}k in § 4.1.1 : supervising the last k k tokens targets exactly the positions in which the student still deviates from the teacher, while sparing compute on the early positions where the signal has already been absorbed. It also explains why a small budget ( k = 2000 ≪ | y ^ | k\!=\!2000\ll|\hat{y}| on average) suffices to recover the gains reported in Table 2 , since OPRD’s representation loss is not diluted across positions that carry no residual signal.

##### (a) Policy-gradient loss: OPRD accelerates distillation and validates the bottleneck theory.