SAKI: Maximal-Coupling-Routed Teacher Supervision for On-Policy Distillation

On-policy distillation (OPD) reduces train–test state mismatch by training a student on its own generated trajectories. However, a weak student may initially visit poor, teacher-misaligned prefixes, forcing the teacher to provide supervision on states that it would rarely generate under its own policy. Teacher-guided rollout policies improve the visited state distribution, but typically retain the same per-prefix reverse-KL objective. We introduce SAKI (Supervision Allocation with KL-constrained Interpolation), which routes token-level supervision using realized accept/correction events from maximal coupling. We construct a geometrically interpolated behavior policy inside a student-centered KL trust region and realize it through maximal coupling with the student. The coupling preserves a student proposal whenever possible and exposes a correction event precisely when realizing the guided policy requires an intervention. Accepted positions retain reverse-KL supervision, whereas correction positions switch to direct supervision on the teacher's highest-probability token. Because the correction probability is exactly $\text{TV}(p_t,q_t)$, the same trust-region radius constrains rollout deviation and upper-bounds intervention and specialized-supervision frequency. An engine-resident speculative verifier preserves the exact-$q$ trajectory distribution and coupling semantics while improving matched-workload rollout throughput by $4.22\times$. Across seven mathematical reasoning benchmarks, our method improves the matched teacher-guided baseline in Mean@8 and Pass@8 for both 1.7B and 0.6B students. Placement controls show that correction-triggered routing outperforms both equal-budget random and TV-weighted placement. Fixed-prefix analysis further shows persistent teacher alignment, with larger gains over random placement at higher initial student–teacher disagreement.