LastOPD: Taming Collapse in Latent On-Policy Distillation

LastOPD. Taming Collapse in Latent On-Policy Distillation. By Jie Yang and nine others, from the University of Illinois at Chicago, Case Western Reserve University, and three other institutions. Published on the twenty-third of September. Its code repository is a placeholder: a readme, and nothing to run yet.

On-policy distillation, OPD, corrects a student model on the responses the student itself writes. Its signal is the teacher's next-token distribution. That tells the student what the teacher says, but it misses how the teacher thinks. Latent supervision promises to supply the missing part, by aligning the student's internal latent states to the teacher's, and recent methods such as OPRD bring that signal into on-policy distillation.

The authors observe two failures of that recipe when distilling Qwen3-4B and Qwen3-8B into Qwen3-1.7B-Base. The first is early gain, then late collapse. Latent supervision alone lifts MATH-500 accuracy from twenty-five to forty-six within ten steps, but further training drags it down to eleven, with no recovery. The second is better alignment with worse behavior. The alignment metric keeps improving right through the collapse, and the most aligned model turns out to be the worst performing one.

Their analysis points to a mismatch in how the latent signal is applied. Layers are paired between the two models by depth, but layers at the same depth play different roles in each. Continued alignment may therefore pull the student toward teacher states it cannot understand.

LastOPD keeps the latent signal but narrows it in two ways, one in space and one in time. In space, it applies the signal only at the last-layer state, which is the common interface that both models' language-model heads read. In time, it applies the signal only during a short crossfade into ordinary token-level OPD.

Each restriction answers one of the two observations. Restricting the signal to the last layer answers the depth-pairing mismatch, because the last-layer state is the one place where both models are, by construction, feeding the same kind of reader. Restricting it in time answers the early gain and late collapse, because the gain arrived within the first ten steps and the damage came afterwards.

The crossfade lasts ten steps. During it, the student gets the useful early push from latent supervision, and then it is handed over to token-level supervision before the collapse has a chance to set in. The design keeps the part of the latent signal that was useful, and stops applying it before the point where it had begun doing harm.

Extensive experiments compare against token-only OPD. With the four billion parameter teacher, LastOPD improves MATH-500 by five point five five points. With the eight billion parameter teacher, the improvement is four point zero two points. It leads on most held-out datasets, not only on the one used to diagnose the collapse. And it reaches the final score of token-only OPD in about half the steps, so the gain shows up in speed as well as in the final score.

The paper's two observed failures belong together. A signal can keep getting better by its own measure while the model it trains gets worse, and the authors' analysis suggests the cause is alignment asked of layers whose roles differ. LastOPD's answer is not to abandon latent supervision but to keep it where the two models genuinely share an interface, and only for as long as it helps.