OPMD: On-Policy Mechanism Distillation Beyond Output Supervision

OPMD. On-Policy Mechanism Distillation Beyond Output Supervision. By Yingdong Shi and five others, from ShanghaiTech University. Published on the twenty-ninth of September. No code repository is linked to it.

On-policy distillation, or OPD, improves a language model by supervising a student on prefixes sampled from the student's own policy. Its supervision, though, stays in the output space. What gets transferred is what the teacher predicts. The intermediate computation that produces those predictions is never supervised. An obvious repair would be to distill intermediate representations directly, for example by representation matching. The authors argue that this falls short for two reasons. Teacher and student representations may differ geometrically. And even hidden states that are well aligned need not be usable by the computation that follows them inside the student.

On-Policy Mechanism Distillation, or OPMD, transfers intermediate computation derived from the teacher while explicitly accounting for whether the student can use it. Compatibility with the student is treated as part of the target, not as an afterthought. It answers both objections to plain representation matching at once. Geometry that differs between the two models is no longer something to force into agreement, and a hidden state only helps if the student can carry it forward.

OPMD works in two stages. First, it constructs a student-compatible teacher: something that transforms the teacher's representations into intermediate states the student's downstream computation can use. Then it has the student internalize those states, learning them from on-policy trajectories, in keeping with the setting it extends. The order matters. The translation into a usable form happens before anything is learned, so what the student absorbs from its own sampled prefixes is already something its later layers can work with, rather than a copy of the teacher's internals.

Across three teacher and student settings, OPMD achieves the highest average accuracy among the distillation baselines evaluated, in every one of them. It also generalizes to distillation from multiple teachers, and to supervised fine-tuning across model families. Further analyses report that mechanism supervision yields more effective reasoning, and meaningfully changes both the student's intermediate computation and its solution strategies. The authors take that as support for the claim that it transfers information output supervision alone does not.

Where OPD teaches a student which answer a teacher would give, OPMD reaches for how the teacher got there, translated first into a form the student can run. The authors present mechanism-level supervision as the way to move on-policy distillation past outputs, letting students learn inner mechanisms they can actually use and internalize.