cs.LGSep 23, 2026
SaveLastOPD: Taming Collapse in Latent On-Policy Distillation
Organizations: University of Illinois at Chicago · Visa Research · Case Western Reserve University · University of Florida · The Ohio State University
Abstract
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 “ ”. 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 | 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 |
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 , token 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 , token constant | 50.90 | 8.05 |
| raise token only | — | reverse top-16 | no latent term, token | 54.15 | 4.80 |
| raise token, hold latent 10 steps | last layer | reverse top-16 | latent constant for 10 steps then off, token | 55.73 | 3.22 |
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 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.