Saki improves on-policy distillation for stronger student models
SAKI: Maximal-Coupling-Routed Teacher Supervision for On-Policy Distillation
Artificial Intelligence
Summary
Training AI models to imitate a teacher model can be tricky when the student model makes mistakes early on, leading it off the right path. The authors introduce SAKI, a clever way to guide the student by mixing the teacher’s suggestions and corrections based on how different the student’s behavior is from the teacher. This method uses math to measure when and how often to correct the student, making sure training stays focused and effective. They also created a fast system to run this process more efficiently and showed it improves performance on math reasoning tasks.
What this means in practice
- •For language model engineers: Train smaller language models to better mimic large teacher models by selectively correcting student predictions during training.
- •For ai system developers: Increase throughput and accuracy of model distillation pipelines in applications needing fast and reliable reasoning capabilities.
Authors
Miteto Wei, Xiaohan Wang, Zehao Chen, Jiajun Chai, Sichao Liu, Li Wang, Haoyuan Xu, Zhaoyu Hu, Wei Lin, Guojun Yin
Abstract
On-policy distillation (OPD) reduces train-test state mismatch by training a student on its own generated trajectories, but weak students may visit teacher-misaligned prefixes where supervision is less representative. We introduce SAKI (Supervision Allocation with KL-constrained Interpolation), which combines a KL-constrained teacher-guided rollout with maximal coupling and reuses realized accept/correction events to route token-level supervision. Accepted positions retain sampled-token reverse-KL supervision, while correction positions receive direct supervision on the teacher's highest-probability token. Under maximal coupling, the correction probability is exactly TV(p_t, q_t), so the same trust-region radius controls rollout deviation and upper-bounds intervention and specialized-supervision frequency. We further implement an engine-resident speculative verifier that preserves the exact-q trajectory distribution and coupling semantics while improving matched-workload rollout throughput by 4.22x. Across seven mathematical reasoning benchmarks, SAKI improves 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.