Verify Before You Distill: Prompt-Level Teacher Gating for On-Policy Distillation

Verify Before You Distill. Prompt-Level Teacher Gating for On-Policy Distillation. By Zhiwei Zhang and eight others, from the AllSpark Team. Published on the second of September. No code repository is linked to it.

On-policy distillation, or OPD, speeds up post-training by having a frozen teacher supervise the student densely, token by token, on the student's own rollouts. Vanilla OPD applies that supervision to every prompt alike. It never asks whether the teacher is reliable on the prompt in front of it. That matters because of how the loss behaves: reverse KL is mode-seeking, so a teacher that is confidently wrong does not produce a weak update. It produces a strong one, pointed the wrong way. The obvious checks do not close the gap either. Distributional proxies such as entropy, or how closely teacher and student likelihoods agree, measure uncertainty or agreement. None of them verifies whether the teacher's answer is actually correct.

Teacher-Gated On-Policy Distillation, TGOPD, rests on one principle: verify the teacher's reliability at the level of the prompt before admitting its dense supervision at all. The teacher has to earn the right to teach each prompt.

Reliability is estimated from a small set of teacher probes, each scored by a verifier. Every prompt is then routed exclusively down one of two paths. If the reliability check passes, the prompt gets dense OPD from the teacher. If it fails, the prompt goes to verifier-grounded GRPO instead, so learning still happens, but it is anchored to checked outcomes rather than to a teacher that may be wrong.

There is a systems angle too. In asynchronous OPD, teacher capacity otherwise sits idle, and TGOPD uses that idle capacity to run its reliability estimation.

Students at two scales were tested, four billion and thirty-five billion parameters, across mathematics, code and instruction following. TGOPD beats Vanilla OPD in all six single-domain settings. Under multi-domain training it reaches higher averages over seven benchmarks at both scales. On the compute side, in the measured single-domain run at four billion parameters, GPU utilization on the teacher node rises from nine point eight percent to seventy-eight point nine percent.

The fix is a gate, placed before the loss rather than inside it. A teacher that can be wrong with confidence is dangerous under a mode-seeking objective, so each prompt is checked first and sent where the signal can be trusted. And the checking pays for itself in hardware that was sitting unused.