AIAny
Icon for item

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

Selectively admits dense token-level teacher supervision only after a prompt-level verifier audit, routing prompts that fail the audit to verifier-grounded trajectory supervision instead — reducing harmful updates from confidently wrong teachers and improving teacher GPU utilization.

Introduction

Why this matters: on-policy distillation (OPD) speeds post-training by using a frozen teacher to provide dense token-level supervision on the student's own rollouts, but a confidently wrong teacher can drive harmful, mode-seeking reverse-KL updates. The core insight is simple and operational: verify teacher reliability per prompt before accepting dense supervision; when the teacher fails the check, fall back to a verifier-grounded trajectory objective instead.

Key Findings
  • Prompt-level gating improves downstream performance: across 4B and 35B student models in math, code, and instruction-following, the gated method (TGOPD) outperforms vanilla OPD in all six single-domain settings and raises multi-domain benchmark averages.
    • Largest gains on code (≈+3.0 at 4B, +2.9 at 35B), then instruction-following (+1.6 / +1.9) and math (+1.5 / +1.2), showing the approach helps most where teacher confidence is least informative.
  • Waste-reducing verification: using otherwise-idle teacher capacity to run small verifier-scored probe rollouts raises teacher-node GPU utilization from 9.8% to 78.9% and cluster utilization from 51.5% to 69.5% in measured runs, with modest end-to-end overhead.
  • Safer supervision routing: instead of blending signals, each prompt is routed exclusively to dense OPD when the audit passes or to verifier-grounded GRPO when it fails, avoiding strong updates from confidently incorrect teacher outputs.
Who It Helps and Trade-offs

Great fit if you: need to distill a large language or sequence model from a frozen teacher but worry about teacher errors corrupting student updates (especially in coding domains or other areas where a teacher can be confidently wrong). It is also suitable when you can afford the small extra runtime to run verifier-scored probes on otherwise-idle teacher capacity.

Look elsewhere if you: cannot run any additional teacher work (no idle teacher capacity), have no reliable verifier for your task, or your training regime absolutely forbids branching supervision rules per prompt — TGOPD relies on a verifier and on using teacher probes to estimate per-prompt reliability.

How It Works

For each prompt, the teacher produces a small set of probe rollouts that a verifier scores; the prompt-level pass rate (fraction of probes judged correct) becomes an online reliability estimate. If the pass rate meets a threshold τ, the system applies dense token-level OPD supervision for that prompt. If it falls below τ, the teacher signal is withheld and the update uses verifier-grounded GRPO (coarser trajectory-level supervision). The probes are scheduled to run primarily when teacher nodes would otherwise be idle, converting wasted capacity into reliability checks and improving overall cluster efficiency.

Information

  • Websitearxiv.org
  • AuthorsZhiwei Zhang, Zechen Sun, Fei Zhao, Kang Peng, Bin Liang, Huayu Deng, Yao Hu, Kam-Fai Wong, Mu Chuan
  • Published date2026/09/02

More Items

Combines sparse verifier outcomes with dense privileged‑hindsight token scoring to learn an outcome‑calibrated, normalized distribution over complete responses for on‑policy self‑improvement. Key features: sign‑gated guidance (retain/reverse/disable per verifier advantage), profiled trajectory balance with one log‑partition per rollout group, and explicit correction against false‑positive self‑guidance.

Analyzes on-policy self-distillation for language-model reasoning, diagnosing “collapse” and framing it as controlled by three levers: where token-level signals apply, what privileged information the teacher sees, and how teacher dynamics evolve.

Uses lightweight discrete-diffusion adapters to draft blocks of tokens in parallel while keeping the original autoregressive model as the final arbiter, enabling lossless speedups (up to 3×) and compatibility with existing AR weights via a Diffusion Distillation phase.