On-policy distillation (OPD) corrects a student on the responses it writes, but its signal is the teacher's next-token distribution: it tells the student what the teacher says but misses how it thinks. Latent supervision promises the missing part by aligning the student's latent states to the teacher's. Recent methods such as OPRD bring this signal into on-policy distillation. However, we observe two failures of this recipe when distilling Qwen3-4B and Qwen3-8B into Qwen3-1.7B-Base. Early gain, late collapse: latent supervision alone lifts MATH-500 accuracy from 25 to 46 in 10 steps, but subsequent training degrades performance down to 11 with no recovery. Better alignment, worse behavior: although the alignment metric steadily improves throughout this collapse, the most aligned model turns out to be the worst performing. Further analysis suggests a mismatch in how the latent signal is applied: layers paired by depth play different roles in the two models, so continued alignment may pull the student toward teacher states it cannot understand. To address this, we propose LastOPD, which applies the latent signal only at the last-layer state, the common interface both LM heads read, and only during a 10-step crossfade into token-level OPD. This keeps the useful part of the latent signal and hands the student to token-level supervision before the collapse sets in. Extensive experiments show that LastOPD improves MATH-500 over token-only OPD by 5.55 and 4.02 points with the 4B and 8B teachers, leads on most held-out datasets, and reaches the final score of token-only OPD in about half the steps. Code is available at https://github.com/Muyiiiii/LastOPD.
Figures & tables
Figure 1: Two failure phenomena of latent supervision. (a) Early gain, late collapse : latent-only distillation (OPRD) beats vanilla OPD within 10 steps and then collapses sharply (top), with representational changes concentrated in the student’s later layers (bottom). (b) Better alignment, worse behavior : during training the optimized projected cosine keeps rising as accuracy falls (top), and before any training the CKA between mid layers is already 0.99, carried by a few massive-activation dimensions (bottom). (c) Per-layer J-lens readout: the student’s answer surfaces gradually, whereas the teacher’s appears only in its last two layers.
Figure 2: What the collapse looks like, and where the models truly correspond. (a) Validation MATH-500 for OPRD-Bridge and LastOPD runs extended to 150 steps. (b) J-Lens top-1 readouts on “ 3∗2+2= ”. The uninformative 0–40% interval is compressed. (c) Per-layer readout agreement (solid lines) and running-value probe accuracy (dashed lines). (d) Raw CKA over all layer pairs. (e) The same map with massive dimensions removed. Details in Appendix B .
Figure 3: Inflated alignment, and the difference that matters. (a) Projected cosine per layer pair, with and without the massive dimensions. (b) Validation MATH-500 at the final step against last-layer readout agreement. Hollow: masked retrain. Squares: same-lineage pair. (c) Same-lineage CKA with the ridge (red). (d) Self-CKA against initialization.
Figure 4: From output-only and layerwise latent OPD to LastOPD. (a) Vanilla OPD matches next-token distributions and leaves the teacher’s latent states unused. (b) Layerwise latent OPD aligns depth-paired layers for the whole run and provides no direct token-level supervision. (c) LastOPD aligns only the last-layer state and crossfades to token-level OPD within the first 10 steps.
Method
MATH-500
AIME24
AIME25
AMC23
Minerva
OlympiadBench
AIMO
GSM8K
Mean
avg@8
best@8
avg@16
best@16
avg@16
best@16
avg@8
best@8
avg@4
best@4
avg@4
best@4
avg@16
best@16
acc@1
avg
best
Qwen3-4B → Qwen3-1.7B-Base
Teacher (Qwen3-4B)
82.85
94.2
22.08
56.7
23.54
46.7
69.06
90.0
31.34
38.2
47.48
60.3
60.47
84.3
91.51
53.54
70.2
Student (Qwen3-1.7B-Base)
24.45
72.2
2.71
10.0
0.62
10.0
18.44
50.0
6.80
18.4
9.33
16.6
9.19
44.6
38.51
13.76
32.5
Token-only OPD
53.40
81.2
9.17
23.3
5.21
20.0
27.50
60.0
13.14
24.6
20.78
36.9
26.58
63.9
67.78
27.95
47.2
OPRD-Bridge (latent-only)
12.12
33.8
0.21
3.3
0.00
0.0
3.12
20.0
7.35
16.9
2.63
7.4
3.69
28.9
8.72
4.73
14.9
Table 1: Main comparison across three teacher–student pairs. Qwen3-4B and Qwen3-8B distill into Qwen3-1.7B-Base, and JustRL-1.5B into R1-Distill-1.5B, whose latent-only recipe is OPRD-Vanilla. Both Mean columns average all eight datasets, with the acc@1 of GSM8K entering each. Bold marks the best trained method per column and block and underline the second.
Figure 5: Ablation and sensitivity study (Qwen3-4B teacher, official MATH-500 except (d)). (a) Schedule variants, with token-only OPD dashed. (b) Alternative forms of the latent loss. (c) Crossfade window and the reference KL term. (d) The always-on recipes and LastOPD with and without the massive-activation mask, on validation MATH-500 at the final step.
Figure 6: General analysis. (a, b) Validation MATH-500 every 10 steps with the Qwen3-4B and the Qwen3-8B teacher, sharing one legend. (c) Last-layer readout agreement with the teacher against validation MATH-500. (d) First layer at which a partial sum enters the top-5 J-Lens readout, over 247 positions from 60 arithmetic chains. Positions never read out are dropped per model.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Model
hidden d
top-1% channels
median channel
largest channel
largest/median
Qwen3-1.7B-Base
2048
20
159.6
7567.0
47.4
Qwen3-4B
2560
25
36.4
2079.3
57.2
Qwen3-8B
4096
40
53.6
4448.4
83.0
JustRL-1.5B
1536
15
48.0
1532.2
31.9
R1-Distill-1.5B
1536
15
45.6
1818.3
39.9
Appendix
Table 2: Massive activations in the five models used in this paper: mean absolute activation per hidden channel over 64 on-policy texts of 512 tokens, all layers and tokens. The top-1% channels are the massive dimensions removed in panel e of Figure 2 .
Configuration
Latent term
Token term
Schedule
MATH-500
Gap
Schedule: when the latent term is on
LastOPD (full)
last layer
reverse top-16
latent 1→0 , token 0→1 over 10 steps
58.95
—
hard switch
last layer
reverse top-16
latent only for 10 steps, then token only
50.45
− 8.50
fade only
last layer
reverse top-16
latent 1→0 , token constant
50.90
− 8.05
raise token only
—
reverse top-16
no latent term, token 0→1
54.15
− 4.80
raise token, hold latent 10 steps
last layer
reverse top-16
latent constant for 10 steps then off, token 0→1
55.73
− 3.22
Appendix
Table 3: Ablations with the Qwen3-4B teacher, official MATH-500 at the final checkpoint. Each block changes one thing about LastOPD. The last column is the gap to the full method (58.95). The reversed-projector, Gram-matrix, and repetition-penalty rows also change the schedule, keeping both terms on throughout, so their like-for-like reference is the always-on row (53.77). The repetition-penalty row comes from the earlier regime (penalty 1.05). In the masked block both columns are validation MATH-500 and the last column is masked minus unmasked.
Figure 7: Latent weight 1 against 1000 on the Qwen3-4B pair. Validation MATH-500 at step 62 for LastOPD and LastOPD-always. Solid: weight 1. Hatched: weight 1000.
Figure 8: Reinforcement learning without a teacher. Training-time validation MATH-500 every 10 steps for GRPO on the DAPO prompts and on DeepMath, against token-only OPD and LastOPD with the Qwen3-4B teacher.
Figure 9: Case-study statistics. (a) Greedy accuracy of the untrained student and LastOPD on the 80 MATH-500 problems with a single-digit answer, under the same chat prompt with a boxed instruction. (b) How the untrained student’s 44 failures break down after a 1600-token retry.
Figure 10: Four kinds of MATH-500 rollout. (a) Share of rollouts of each kind at step 62 for a replicate OPRD-Bridge run, the same run with only its LM head retrained on teacher text, token-only OPD, and LastOPD. Problems whose answer is one or two characters are left out of the unboxed kinds. (b) The same shares along the OPRD-Bridge run.
Figure 11: When the answer becomes readable, on ten one-digit prompts. Each dot is one prompt and the bar is the median. (a) Depth at which the final answer first becomes the top-1 J-Lens readout. (b) Depth at which the intermediate product first enters the top-5 readout. Hollow dots: never.
Figure 12: Readout agreement with the final output by relative depth , for the untrained student, the three trained students, and the teacher.
Figure 13: Reading the running partial sum along one chain. For 1+2+3−2+3−4+4−5= the top row of each panel lists the token read in and, below it, the partial sum at that token. Each cell is the rank of that partial sum in the layer’s J-Lens readout, blank when it falls outside the top 99. The untrained student and LastOPD read the running value out from the middle layers on, while the teacher shows little before its last ten layers.
Figure 14: Pairing layers by the raw CKA ridge. (a) The layer map used by the run, against pairing by depth. (b) Its validation trajectory against the token-only OPD and LastOPD final scores.
Figure 15: Layerwise J-Lens readout on all ten one-digit prompts. Each panel repeats Figure 2 b for one prompt and five models: the untrained student, LastOPD, token-only OPD, OPRD-Bridge, and the teacher. Depth is normalized per model, the 0–40% interval is compressed to the first layer, and green marks the final answer, amber the intermediate product, and grey anything else.
On-policy distillation (OPD) supervises the student exclusively in the output space by matching next-token distributions. This paradigm suffers from two limitations: (i) a high-variance gradient estimator whose signal-to-noise ratio collapses as the student approaches the teacher, and (ii) an LM-head information bottleneck that discards the teacher's intermediate hidden states. We propose On-Policy Representation Distillation (OPRD), the first method to lift on-policy distillation into the hidden-state space. OPRD aligns student and teacher representations across selected layers on the same on-policy rollouts, providing dense, deterministic, per-layer supervision while bypassing the LM head entirely. Theoretically, OPRD provides a deterministic per-sample gradient, removing the token-level estimation variance that plagues OPD, and exposes structural information that any output-space objective necessarily discards. Empirically, OPRD closes the student-teacher gap on competition mathematics benchmarks (AIME 2024, AIME 2025, and AIMO), where every output-space baseline plateaus below the teacher, while training 1.44x faster and using up to 54% less memory. We further extend OPRD to the cross-architecture setting via OPRD-Bridge. By exploiting the observation that heterogeneous models share a low-rank representational structure, we construct a frozen projector pair that aligns representations across arbitrary depth and width mismatches, shifting the alignment from the output space (which depends on a shared vocabulary) to the representation space. We validate OPRD-Bridge on both cross-architecture (Qwen3-4B -> Qwen3-1.7B-Base) and cross-tokenizer (Phi-4-mini-reasoning -> Qwen3-1.7B-Base) settings, demonstrating successful knowledge transfer even when the vocabulary-based alignment channel is unavailable. Code: https://github.com/ShenzhiYang2000/OPRD.
On-policy distillation (OPD) trains a student on its own generated responses using dense, token-level supervision from a stronger teacher. Vanilla OPD treats all teacher signals equally, assuming that the teacher's supervision is equally important for every token. However, teacher signals at different tokens may have very different effects on the student's performance: some correct important reasoning errors, while others have little effect on the final answer. Motivated by this observation, we introduce Dr. OPD (OPD Done Right), which defines the optimal weighted OPD to maximize the student's performance. We formulate Dr. OPD as a bilevel optimization problem in which the student learns from weighted teacher supervision, while the weights are selected to maximize the expected reward of the resulting student. To solve Dr. OPD, we develop an efficient iterative solver that updates the token weights and student policy alternatively. At each round, it updates weights in closed form and then takes one gradient step on the resulting weighted OPD objective. Under regularity conditions, we show that this weighted update achieves a higher expected reward than a vanilla OPD update. Empirically, across strong-to-weak and same-size distillation on math and code, Dr. OPD consistently outperforms all evaluated baselines. In particular, in the strong-to-weak distillation setting, Dr. OPD improves average math performance by 9.7 points over vanilla OPD, and enables the smaller student to surpass its larger teacher.
On-policy distillation transfers reasoning capabilities by training a student model on its own generated trajectories using token-level feedback from a teacher. However, we identify a critical bottleneck, \textbf{Supervision Fidelity Decay (SFD)}: as student-generated prefixes lengthen, the teacher's next-token distribution becomes less confident and less discriminative. Consequently, the teacher-dependent corrective signal in reverse-KL distillation weakens, causing student drift to compound across long reasoning chains. To mitigate SFD, we introduce \textbf{Lookahead Group Reward (\ours{})}. Building on the insight that next-step teacher confidence reflects the discriminative strength of future reverse-KL supervision, \ours{} evaluates the student's top-K candidate tokens by the teacher confidence they induce at the subsequent step and assigns a group-normalized reward. To maintain computational efficiency, we further design an entropy-triggered tree-attention mechanism. Across six math and code benchmarks, \ours{} improves mean@8 by \textbf{2.57} points over OPD for a 7B student, with gains increasing in longer-generation and reaching +\textbf{4.92} points on AIME-26 at 39k tokens.
Yanjiang Liu, Jie Lou, Xinyan Guan +7
University of Chinese Academy of Sciences · 2Chinese Information Processing Laboratory Institute of Software, Chinese Academy of Sciences 2University of Chinese Academy of Sciences, Beijing, China · 3Xiaohongshu