LastOPD: Taming Collapse in Latent On-Policy Distillation
Jie Yang, Zhengyu Fang, Zelin Xu, Jiarui Sun, Xiran Fan, Junpeng Wang, Liang Wang, Qinghua Liu, Yiwei Cai, Yan Zheng
cs.LG, cs.CL
2026-09-24
LastOPD aligns only last-layer states and fades into token-level OPD in 10 steps. Qwen3-4B/8B into 1.7B: +5.55/+4.02 MATH-500 vs token-only OPD, matched in about half the steps.
On-policy distillation (OPD) lets the student write its own answers and has the teacher score every token on those prefixes. The signal is the next-token distribution after the teacher's LM head: what to say next, not how it was computed.
Latent distillation tries to supply the missing computation by aligning hidden states. OPRD puts that idea inside OPD, pairing layers by relative depth and bridging width with a frozen low-rank projector. Across size, those paired layers may not do the same work.
Distilling Qwen3-4B and Qwen3-8B (36 layers, width 2560 and 4096) into Qwen3-1.7B-Base (28 layers, width 2048) produces two failures. Latent-only OPRD-Bridge lifts validation MATH-500 from 25.4 to 46.2 in 10 steps, about 28 points ahead of token-only OPD at the same step, then falls to 11.3, below the 24.45 start. Extending the run to 150 steps leaves the score between 9.4 and 10.8. About two thirds of the collapsed answers are short, boxed, and wrong. Asked for the gcd of 128 and 144, the student still writes a Euclidean-algorithm outline with headings and a box, then divides 128 by 16 and reports quotient 1 remainder 0. Form survives; carries and remainders do not.
In the same run the projected cosine that the loss actually optimizes rises from 0.97 to 0.98 (4B) and from 0.95 to 0.98 (8B), while 4B MATH-500 drops about 35 points. Mid-layer CKA is already 0.99 before training, with scores 24.5 versus 82.9. The alignment metric goes up as the student gets worse.
J-Lens decodes each layer into next-token logits. The teacher knows early and tells late: a linear probe reaches 94% on partial sums by layer 9, yet readout agreement stays near zero until 63% of depth. The untrained student starts saying the answer from layer 12. Pairing by depth then asks the student to match states it cannot read. Masking massive activations (a few huge-magnitude channels) drops the collapsed endpoint from 11.33 to 5.60. Remapping layers along the raw CKA ridge still peaks near 50 at step 10 and collapses by step 20. Adding the token term to OPRD-Bridge recovers 53.70, next to token-only OPD at 53.40. Left on for the whole run, the latent term adds nothing on top of tokens.
LastOPD changes two knobs: where the latent signal sits, and how long it stays on.
Training still uses the student's own samples: DAPO-Math-17k, 62 optimizer steps, about 7.9k rollouts, Qwen3 with thinking off. The projector is dropped at inference. After step 10 the teacher only returns next-token probabilities on the student's top-k candidates.
The useful slice of the latent signal arrives early. Keeping the match on pulls the student toward teacher states it cannot use. LastOPD spends that slice at the shared interface, then switches to token-level supervision before collapse sets in.
Three teacher-student pairs, one run per cell, final checkpoint only.
| Method | MATH-500 | 8-dataset mean |
| 4B→1.7B LastOPD | 58.95 | 31.88 |
| 4B token-only OPD | 53.40 | 27.95 |
| 4B OPRD-Bridge | 12.12 | 4.73 |
| 4B LastOPD-always | 53.77 | 28.88 |
| 8B→1.7B LastOPD | 53.45 | 28.38 |
| 8B token-only OPD | 49.43 | 26.28 |
| Same-lineage 1.5B OPRD-Vanilla | 87.22 | 63.49 |
| Same-lineage 1.5B LastOPD | 83.42 | 62.91 |
With 4B and 8B teachers, LastOPD beats token-only OPD by 5.55 and 4.02 MATH-500 points, and by 3.93 and 2.10 on the eight-dataset mean. It leads on five and six of the other seven held-out sets, by up to 8.0 and 6.5 points. Under the 4B teacher, AIME24 and AIME25 do not go up (7.50 vs 9.17, 3.96 vs 5.21). LastOPD-always, which keeps the last-layer term on, scores 53.77 and 48.95; the fade adds another 5.18 and 4.50. Layerwise latent plus token finishes at 53.70 and 47.05, essentially token-only.
Ablations use the 4B teacher and official MATH-500. A hard switch at step 10 scores 50.45, worse than never using latents. Fade-only 50.90, raise-token-only 54.15, windows of 5 and 15 steps 54.62 and 54.47. A reversed projector scores 50.00; a Gram-matrix loss scores 46.23. Masking massive dimensions barely moves LastOPD (validation 59.72 to 59.35) and knocks 5 to 13 points off every always-on recipe.
After the fade, validation MATH-500 jumps from 23.1 to 46.6 in ten steps with the 4B teacher, then 55.4 by step 30, already past token-only's 51.9 at step 62. With the 8B teacher, step 30 is 51.7 against a token-only finish of 50.4. About half the steps. A 150-step LastOPD run sits near 57 from step 90 on.
The same-lineage pair reverses the ranking. JustRL-1.5B into R1-Distill-1.5B has diagonal CKA 0.983, against 0.684 for the 4B cross-size pair. Latent-only OPRD-Vanilla hits 87.22, above token-only at 86.02. LastOPD finishes at 83.42, slightly below the untrained student's 83.67. When layers already correspond, a fixed 10-step fade costs accuracy.
For people distilling across size: do not pair every layer by depth, and do not leave the latent term on for the whole run. Ten steps of last-layer alignment, then token-level KL, adds 4 to 5.5 MATH-500 points over token-only OPD on Qwen3 4B/8B into 1.7B, in about half the steps, and needs teacher hidden states only for those first ten steps.
CKA at 0.99 and a rising projected cosine both fail to track whether the student can still do the math. On same-lineage, same-depth pairs, leaving latent supervision on still wins. LastOPD is for the cross-size case.
There is no dedicated limitations section. Training and evaluation are almost entirely math, and Qwen3 runs with thinking off. Code, general instruction following, and thinking mode are untested for the same early-gain late-collapse pattern.
Cross-size evidence is one Qwen3 pair plus a 1.5B same-lineage control. Tw=10 is chosen from the MATH-500 peak of latent-only training, and MATH-500 is also the diagnostic set; the other seven datasets are the independent evidence. Every cell is a single run.
On the same-lineage 1.5B pair, LastOPD hurts MATH-500. Treating a 10-step fade as a universal recipe regresses when layers already match.
The 'knows early, tells late' picture comes from a linear probe on arithmetic and ten one-digit expressions. That is thin evidence for general reasoning traces. A hard switch is worse than skipping the latent term, so the handoff is more brittle than the choice to use latents at all. Teacher-free GRPO on the same prompts also drifts, from 41.1 down to 20.8; every distillation run that keeps a token term still finishes above the starting point.