Researchers have proposed SAKI (Supervision Allocation with KL-constrained Interpolation), a new approach to on-policy distillation (OPD) that aims to fix a key weakness: weak students often generate trajectories that drift into prefixes where teacher supervision is less representative.
SAKI combines a KL-constrained teacher-guided rollout with maximal coupling, then reuses the realized accept/correction events to route token-level supervision. Accepted positions keep the sampled-token reverse-KL supervision, while correction positions receive direct supervision on the teacher's highest-probability token.
Under maximal coupling, the paper shows that the correction probability is exactly the total variation distance between the student and teacher distributions at each step. This means a single trust-region radius controls both rollout deviation and the upper bound on intervention and specialized-supervision frequency.
The authors also implement an engine-resident speculative verifier that preserves the exact trajectory distribution and coupling semantics. It improves matched-workload rollout throughput by 4.22x.
Across seven mathematical reasoning benchmarks, SAKI improves over the matched teacher-guided baseline in Mean@8 and Pass@8 for both 1.7B and 0.6B students. Placement controls and fixed-prefix analysis further support correction-triggered routing as a conflict-adaptive supervision signal.