LastOPD: Mitigating Collapse in Latent On-Policy Distillation

Meta description: LastOPD fixes latent supervision collapse in Qwen3 on-policy distillation, improving MATH-500 scores and training efficiency in the paper's experiments.

Tags: Qwen3, On-policy distillation, Latent supervision, LastOPD, MATH-500

Featured image: Abstract technology image: a neural pathway stabilizing from chaotic layers into a single aligned output interface, cool blue and amber tones, no text or logos.

Background: On-Policy Distillation Needs More Than Token Answers

LastOPD arrives from a narrow problem in model distillation: a small student model can imitate a larger teacher's word choices, but not necessarily the internal computation that produced them. In on-policy distillation (OPD), the student generates its own responses, and the teacher corrects those responses token by token. That correction uses the teacher's next-token distribution. It gives a useful signal, but it mainly tells the student what the teacher says, not how the teacher thinks.

Latent supervision is the attempt to fill that gap. Instead of matching only the final next-token probabilities, it aligns the student's latent states with the teacher's latent states. Recent methods such as OPRD bring this signal into on-policy distillation, aiming to give students a richer training signal while still keeping the on-policy correction loop.

Development: A Quick Gain Precedes a Sharp Collapse

The reported experiments focus on distilling Qwen3-4B and Qwen3-8B into Qwen3-1.7B-Base. When latent supervision is used alone, the early effect is dramatic: MATH-500 accuracy rises from 25 to 46 in 10 steps. That is a large jump for a small base model, and it appears to validate the promise that aligning hidden states can transfer more than token-level choices.

The later trend reverses that reading. Continued training pushes performance down to 11, with no recovery reported. The failure is not simply that the student forgets the teacher. The alignment metric keeps improving while the model's behavior worsens. In other words, the most aligned checkpoint is also the worst performing checkpoint. That separation matters because it means standard alignment scores can become misleading during training.

Further analysis points to a structural mismatch. When latent states are paired by layer depth, the same layer number does not necessarily mean the same computational role. A layer that performs one kind of aggregation in a 1.7B model may sit in a different functional position inside a 4B or 8B teacher. Continued depth-by-depth alignment may therefore pull the student toward teacher states it cannot understand. The training signal becomes more precise in alignment terms while becoming less useful in behavior terms.

Latest: LastOPD Limits the Signal to the Last Layer

LastOPD responds by narrowing where the latent signal is applied. Rather than aligning many internal layers, it uses the last-layer state, described as the common interface both language model heads read. The latent supervision is active only during a 10-step crossfade into token-level OPD. After that transition, the student returns to the more stable token-level correction path.

This design keeps the early advantage of latent supervision without letting it dominate the whole training run. The 10-step window matches the point where the reported MATH-500 gain appeared, then hands the student back to token-level OPD before the collapse pattern can take hold.

Current Stage: A More Stable Distillation Recipe

The experiments report that LastOPD improves MATH-500 over token-only OPD by 5.55 points with the Qwen3-4B teacher and by 4.02 points with the Qwen3-8B teacher. It also leads 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.

The result suggests a practical boundary for latent supervision in small-model distillation: hidden-state alignment is most useful when it is brief, targeted, and followed by token-level correction. For Qwen3 student-teacher setups, the immediate lesson is that deeper alignment is not automatically better. The next question is whether the last-layer-only signal remains stable when the teacher, student, task distribution, or training budget changes outside the reported MATH-500 setting.