Source-linked AI summary
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
TL;DR
On-policy distillation scores student-generated tokens with the teacher’s next-token distribution, leaving the latent computation behind each token unsupervised. LastOPD aligns the student’s and teacher’s last-layer states before their LM heads and crossfades that latent term into reverse top-k token-level OPD over 10 steps. Across two cross-size Qwen3 pairs, LastOPD improves MATH-500 over token-only OPD by 5.55 and 4.02 points, respectively.
Problem
On-policy distillation scores student-generated tokens with the teacher’s next-token distribution, leaving the latent computation behind each token unsupervised.
Source: §1 · p. 1Method
Results
Across two cross-size Qwen3 pairs, LastOPD improves MATH-500 over token-only OPD by 5.55 and 4.02 points, respectively.
Source: §1 · p. 3Takeaways & Limitations
The reported results support using a short-window latent signal at the shared final-layer interface rather than continued intermediate-layer alignment.
Source: §1 · p. 3, §7 · p. 10
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The introduction identifies a mismatch between token-level feedback and latent reasoning supervision, then shows that latent-only alignment can help early but fail later. LastOPD confines latent supervision to the final-layer interface and crossfades into token-level OPD, with reported gains over token-only OPD.
1 INTRODUCTION
- Latent-only OPRD raises validation MATH-500 from 25.4 to 46.2 within 10 steps, but continued optimization lowers it to 11.3 by training’s end.The 8B-teacher experiment shows the same early-gain-then-collapse pattern.Source: §1 · p. 2
- Projected cosine alignment keeps rising as accuracy falls, while teacher and student answer readouts emerge at different depths.The teacher’s answer appears only in its last two layers, whereas the student’s emerges gradually; mid-layer CKA is already 0.99 before training and is carried by a few massive-activation dimensions.Source: §1 · p. 2, Figure 1 · p. 2
- The authors argue that depth-paired layers may do different work, so continued alignment may pull the student toward teacher states it cannot understand.Their diagnostics report that masking massive dimensions or remapping layers by CKA does not prevent collapse.Source: §1 · pp. 2–3
- The central question is how to retain latent supervision’s early gain without its later collapse.The motivation is that token-level OPD transfers the teacher’s next-token distribution but does not score the latent computation upstream of the LM head.Source: §1 · pp. 1, 3
- LastOPD aligns last-layer states before the LM heads and crossfades the latent term into reverse top-k token-level OPD over 10 steps.The latent term then turns off; this design requires no intermediate-layer pairing or frozen projector.Source: §1 · p. 3
- Across two cross-size Qwen3 pairs, LastOPD improves MATH-500 over token-only OPD by 5.55 and 4.02 points, and the eight-held-out-dataset mean by 3.93 and 2.10 points, respectively.Source: §1 · p. 3
2 RELATED WORK
Prior OPD methods transmit the teacher’s next-token distribution, while latent-distillation methods align selected internal states. The paper positions LastOPD as a response to the open question of whether such targets are reachable across architectures and for how long they help.
2 RELATED WORK
- OPD trains on student-generated samples scored by the teacher, but its signal remains the teacher’s next-token distribution after the LM head.Variants change which tokens are scored or their weighting, not the basic signal.Source: §2 · p. 3
- Latent distillation has aligned student and teacher intermediate states or attention statistics since FitNets [Romero et al., 2015].LOPD, PHF, and OPRD bring latent information into on-policy distillation.Source: §2 · p. 3
- The paper identifies an open question in prior latent methods: whether cross-architecture targets are reachable and how long alignment helps, motivating LastOPD’s final-layer signal and short application window.LastOPD combines this latent signal with token-level OPD.Source: §2 · p. 3
3 PRELIMINARIES
OPD trains on student-generated responses, so the teacher’s corrections apply to contexts the student actually produces. Latent distillation adds supervision from internal states, and OPRD applies it across depth-paired layers.
3 PRELIMINARIES
- OPD has the student sample its response and the teacher score tokens on the same prefixes, placing corrections on student-generated contexts.The token-level signal is the teacher’s next-token distribution read after its LM head.Source: §3 · pp. 3–4
- The reverse top-k objective restricts next-token distributions to the student’s highest-ranked tokens, with k = 16 throughout.The restricted distributions are renormalized.Source: §3 · pp. 3–4
- Latent distillation adds a pre-LM-head signal by aligning student and teacher states; OPRD applies it on student rollouts at depth-paired layers using a frozen low-rank projector.OPRD names its cross-architecture form OPRD-Bridge and its same-architecture form OPRD-Vanilla (Yang et al., 2026c).Source: §3 · p. 4
4 WHY LATENT SUPERVISION COLLAPSES
Extended latent-only training remains collapsed, while analyses find that layer roles differ and measured alignment can be inflated by massive dimensions. Masking or remapping does not prevent collapse, whereas adding token-level supervision substantially improves the result.
What the collapse looks like
- Validation MATH-500 stays between 9.4 and 10.8 from step 70 through the 150-step OPRD-Bridge run, without recovery.Roughly two thirds of collapsed answers are short, boxed, and wrong, while fewer than one in ten loop.Source: §4 · p. 5
When each model knows and tells
- Teacher readout agreement stays near zero until 63% of depth, while the student’s rises gradually; probes find information earlier in both models.The teacher’s probe reaches 94% accuracy by layer 9, while the student’s passes 80% by layer 7 and peaks at layer 18.Source: §4 · p. 5
Can similarity guide layer pairing?
- Massive dimensions carry 43% of projector energy; masking them lowers step-one projected cosine from 0.88 to 0.29, while the last-layer CKA match remains 0.38.After masking, teacher layers 16–26 are never selected along the CKA ridge.Source: §4 · p. 5, Figure 3 · p. 5
Do masking and remapping prevent the collapse?
- Masking dimensions lowers final validation MATH-500 from 11.33 to 5.60, and CKA-ridge remapping still peaks near 50 at step 10 before collapsing by step 20.Both variants continue to supervise latent states without direct next-token supervision.Source: §4 · pp. 5–6, Figure 3 · p. 5
- Adding token-level OPD to OPRD-Bridge raises official MATH-500 from 12.12 to 53.70, close to 53.40 for token-only OPD.The authors hypothesize that continued matching can pull the student toward teacher-state components it cannot understand.Source: §4 · p. 6
5 METHODOLOGY
LastOPD applies latent supervision at the last-layer state and transitions early to token-level OPD. Its comparison covers three teacher–student pairs, with mean scores over eight datasets.
5 METHODOLOGY
- LastOPD applies latent matching only to post-normalization last-layer states, a common next-token prediction interface across models of different depths and widths.A trainable MLP bridges width differences, and gradients update the student and adapter while the teacher remains frozen.Source: §5 · p. 6
- The latent weight falls linearly to zero as the token-level OPD weight rises to one during a 10-step crossfade.At step 0 the objective is latent-only; from step 10 onward it is token-only OPD.Source: §5 · p. 6
- The projector is discarded at inference, so the adapter is not retained for inference.The passage notes that the base latent-loss weight and remaining overhead are described in Appendix A.Source: §5 · p. 6
- Table 1 compares three teacher–student pairs and reports mean columns averaging eight datasets, including GSM8K acc@1.The caption says bold and underline mark the best and second-best trained methods per column and block, respectively.Source: Table 1 · p. 7
6 EXPERIMENTS
Across three teacher–student pairs, the experiments compare LastOPD with token-only and latent-supervision baselines, then test its design choices and training dynamics. LastOPD improves held-out performance and reaches token-only OPD’s final score in fewer steps.
6.1 EXPERIMENTAL SETUP
- The comparison trains three teacher–student pairs for 62 steps on DAPO-Math-17k and evaluates on eight held-out math datasets against token-only OPD and latent-supervision baselines.Every method uses the same prompts, batches, and optimizer; scores are taken at the final checkpoint.Source: §6.1 · p. 7
6.2 OVERALL COMPARISON
- LastOPD reaches 58.95 and 53.45 on MATH-500 with Qwen3-4B and Qwen3-8B, beating token-only OPD by 5.55 and 4.02 points; it leads on five and six other datasets.Its eight-dataset mean improves over token-only OPD by 3.93 and 2.10 points, respectively.Source: §6.2 · pp. 7–8
- On the same-lineage pair, latent-only OPRD-Vanilla scores 87.22 on MATH-500, above the student’s 83.67 and token-only OPD’s 86.02; continued latent supervision performs better there.LastOPD remains within 0.58 points of the best method on the eight-dataset mean.Source: §6.2 · p. 8
6.3 ABLATION AND SENSITIVITY STUDY
- A hard switch at step 10 scores 50.45 on MATH-500 versus 58.95 for the crossfade; neither half alone recovers the full gain.The ablation compares schedule variants in Figure 5a; the reported results support combining early latent supervision with a gradual reduction in its weight.Source: §6.3 · p. 8, Figure 5 · p. 8
- Five tested variants underperform the full recipe: 5- and 15-step windows score 54.62 and 54.47, while a reversed projector scores 50.00 and a Gram-matrix loss 46.23.Adding a reference KL term lowers the score by 0.4 points.Source: §6.3 · pp. 8–9
- Masking massive dimensions lowers all three always-on recipes but changes LastOPD by only 0.4 points, from 59.72 to 59.35 on validation MATH-500.Figure 5d compares the always-on recipes and LastOPD with and without the mask.Source: Figure 5 · p. 8, §6.3 · p. 9
6.4 GENERAL ANALYSIS
- LastOPD reaches token-only OPD’s final MATH-500 score by step 30 with both teachers: 55.4 versus 51.9 for 4B, and 51.7 versus 50.4 for 8B.With the 4B teacher, a 150-step run stays near 57 from step 90, while OPRD-Bridge remains between 9 and 11.Source: §6.4 · p. 9
- At similar last-layer readout agreement, OPRD-Bridge scores 11.33 and LastOPD 59.72; partial sums become readable at median layer 14 for LastOPD, versus 18.5 untrained and 17 for token-only OPD.The student’s final answer appears at the same relative depth as before, while intermediate results surface earlier under LastOPD.Source: §6.4 · pp. 9–10, Figure 6 · p. 9
7 CONCLUSION
The paper identifies collapse and misaligned optimization as failures of latent supervision in OPD, and proposes LastOPD to address them. Across the reported experiments, LastOPD improves results on most datasets and reaches token-only OPD’s final score in about half the steps.
7 CONCLUSION
- Latent supervision can help early but harm later, while its alignment metric improves even as the student collapses.The paper studies these failures across model sizes.Source: §7 · p. 10
- LastOPD applies latent supervision only to the last-layer state and only during a 10-step crossfade into token-level OPD.The student still tells as it goes and reaches intermediate results earlier, rather than adopting the teacher’s pattern of knowing early and telling late.Source: §7 · p. 10
- Across three teacher–student pairs, LastOPD outperforms existing methods on most datasets and reaches token-only OPD’s final score in about half the steps.The reported experiments also find that intermediate results surface earlier under LastOPD.Source: §7 · p. 10
A EXPERIMENTAL DETAILS
The appendix specifies the model pairs, training and evaluation protocols, baselines, and extended diagnostic runs. It also describes a 150-step continuation used to plot training-time validation scores.
A EXPERIMENTAL DETAILS
- The cross-size pairs are Qwen3-4B or Qwen3-8B teachers with a Qwen3-1.7B-Base student; a same-lineage JustRL-1.5B/R1-Distill-1.5B pair serves as a control.All pairs share a tokenizer, and the teacher remains frozen.Source: A EXPERIMENTAL DETAILS · p. 13
- Each step uses 32 prompts with four responses per prompt, temperature 1.0, an 8192-token limit, and a learning rate of 10^-5; methods share prompts, batches, and optimizer.The 62 steps yield about 7.9k rollouts; LastOPD needs teacher latent states only during its first Tw steps.Source: A EXPERIMENTAL DETAILS · p. 13
- Final-checkpoint evaluation uses temperature 0.7 and top-p 0.95, with dataset-specific sample counts and token limits; no best intermediate checkpoint is selected.MATH-500 is used for diagnostics and choosing Tw, while the other seven datasets provide evidence independent of that choice.Source: A EXPERIMENTAL DETAILS · p. 13
- Baselines include token-only OPD, latent-only OPRD-Bridge, OPRD-Bridge plus token loss, and LastOPD-always, which keeps both loss terms at fixed weights.Same-lineage variants use OPRD-Vanilla names and do not require a projector.Source: A EXPERIMENTAL DETAILS · p. 13
B.1 LONG RUNS (PANEL A)
- The long-run diagnostic extends the training recipe to 150 optimizer steps and plots training-time MATH-500 validation every 10 steps using eight samples per problem.All other settings remain unchanged.Source: §B.1 · p. 13
B.2 LAYERWISE READOUTS (PANEL B)
The appendix details probes of layerwise predictions, intermediate values, cross-model representations, and large activations. These diagnostics specify how the reported layer-readout and CKA measurements are constructed.
B.2 LAYERWISE READOUTS (PANEL B)
- J-Lens decodes each layer’s latent state into next-token logits; the layerwise readout example uses one-digit arithmetic expressions and marks final answers, intermediate products, and other outputs.The expressions use a two-shot prefix, and the readout is the top-1 token at each layer’s final position.Source: §B.2 · pp. 13–14
B.3 READOUT AGREEMENT (PANEL C, SOLID LINES)
- Readout agreement measures the fraction of positions whose layer prediction has settled on the token the model eventually emits, across 192 positions.The sample comprises 64 MATH-500 problems and three mid-generation positions per problem.Source: §B.3 · p. 14
B.4 VALUE PROBE (PANEL C, DASHED LINES)
- The value probe tests whether a layer’s latent state predicts the unspoken partial sum a + b across 729 arithmetic expressions, using held-out test accuracy.The target has 17 possible values, and chance accuracy is 6%.Source: §B.4 · p. 14
B.5 LAYER-TO-LAYER CKA (PANELS D AND E)
- Layer-to-layer linear CKA is computed on 200 on-policy texts of up to 512 tokens, with the same sampled token positions across all layers and models.A second computation removes each model’s top-1% massive channels.Source: §B.5 · p. 14
B.6 MASSIVE ACTIVATIONS
- Massive dimensions are each model’s top-1% hidden channels by mean absolute activation, measured across layers and tokens on 64 on-policy texts.Each model’s largest channel is 32 to 83 times its median channel; the two same-lineage models share 93% of their top channels.Source: §B.6 · p. 14, Table 2 · p. 14
C ADDITIONAL RESULTS
The appendix tests ablation settings and compares training and output behaviors. Higher latent weighting and a CKA-based layer map do not prevent collapse, while LastOPD’s final alignment and outputs differ from collapsed runs.
C.1 FULL ABLATION TABLE
- Table 3 reports Qwen3-4B ablations at the final checkpoint on official MATH-500, with a gap column comparing variants against full LastOPD (58.95).The reversed-projector, Gram-matrix, and repetition-penalty rows use the always-on reference (53.77); the masked block uses validation scores.Source: §C.1 · p. 14, Table 3 · p. 15
C.2 RAISING THE LATENT WEIGHT ON THE 4B PAIR
- At weight 1000, LastOPD ends at 57.57 versus 59.72 at weight 1, while LastOPD-always ends at 33.90 versus 54.75 after falling to 13.33 by step 20.The always-on variant scores 25.30 at step 10; the authors conclude that only the crossfade keeps the gain.Source: §C.2 · p. 15, Figure 7 · p. 15
C.3 REINFORCEMENT LEARNING WITHOUT A TEACHER
- Teacher-free GRPO peaks at 41.1 then ends at 20.8 on DAPO prompts, and peaks at 45.9 then ends at 36.2 on DeepMath; token-term distillation runs end above their starting scores.The comparison tests whether collapse occurs under on-policy training without a teacher.Source: §C.3 · p. 15, Figure 8 · p. 16
C.4 ALIGNMENT AT THE FINAL CHECKPOINTS
- LastOPD scores 5.55 points higher than token-only OPD despite lower final last-layer CKA (0.546 vs 0.570).Collapsed OPRD-Bridge and LastOPD have similar readout agreement (0.62 vs 0.63), but score 12.12 and 58.95, respectively.Source: §C.4 · p. 16
C.5 TWO COLLAPSED ANSWERS SIDE BY SIDE
- The collapsed student preserves proof formatting but makes arithmetic errors; LastOPD answers both cited problems correctly and gets single-digit sums and products right in nine of ten probes.The collapsed student miscomputes a quotient and reports an incorrect polar-coordinate result, while LastOPD states the algorithm before applying it.Source: §C.5 · p. 16
C.6 WHAT THE COLLAPSED STUDENT SAYS
At step 62, collapsed rollouts are mostly short boxed wrong answers, while LastOPD and token-only OPD retain larger correct shares. On single-digit MATH-500 problems, LastOPD improves over the untrained student, and few baseline failures are rescued by longer decoding.
C.6 WHAT THE COLLAPSED STUDENT SAYS
- At step 62, 66.8% of collapsed rollouts are boxed and wrong, versus 15.8% boxed and correct; LastOPD and token-only OPD have correct shares of 59.6% and 56.0%.Wrong answers average 765 tokens; the wrong share rises from 36.5% at step 5 as the correct share falls from 57.0%.Source: §C.6 · pp. 16–17, Figure 10 · p. 17
C.7 WHERE THE GAIN SHOWS UP IN THE OUTPUTS
- On 80 single-digit-answer MATH-500 problems, LastOPD answers 45 correctly versus 36 for the untrained student, winning 14 and losing 5; only 3 of 44 baseline failures are rescued by a 1,600-token retry.All 45 LastOPD correct answers stop within 600 tokens.Source: §C.7 · p. 17
C.8 TRAINING DOES NOT CHANGE WHEN THE STUDENT TELLS
Across student recipes, the depth at which final answers become readable remains similar, while partial sums surface earlier in students than in the teacher. A CKA-based layer map follows the collapse pattern, and fading the latent term across the whole run gives no gain over token-only OPD.
C.8 TRAINING DOES NOT CHANGE WHEN THE STUDENT TELLS
- The teacher’s final answer becomes readable at median relative depth 0.97, versus 0.87–0.89 for student recipes; training leaves this student answer depth unchanged.For the intermediate product, the teacher median is 0.97 and student medians range from 0.74 to 0.81; across 192 positions, student recipes share a similar readout curve.Source: §C.8 · pp. 17–18, Figure 12 · p. 18, Figure 15 · p. 20
C.9 WHERE THE RUNNING VALUE SURFACES ALONG ONE CHAIN
- Both the untrained student and LastOPD surface the running value from layers 14–16, while the teacher shows little before layer 26 and reads it at rank 1 from layer 28.LastOPD brings the value to rank 1 at more positions and from slightly shallower layers than the untrained student.Source: Figure 13 · p. 18, §C.9 · p. 19
C.10 PAIRING LAYERS BY THE RAW CKA RIDGE
- Pairing student layers by the raw CKA map still peaks near 50 at step 10 and collapses by step 20, following the same lifecycle as depth pairing.The map sends student layers 1–15 to teacher layer 6 and the last four student layers to teacher layers 34 and 35.Source: §C.10 · p. 19, Figure 14 · p. 19
C.11 FADING THE LATENT TERM OVER THE WHOLE RUN
- Fading the latent weight to zero over all 62 steps ends at 51.65 on validation MATH-500, at the level of token-only OPD.The authors report no gain from keeping the latent term active for most of the run.Source: §C.11 · p. 19