SAKI: Maximal-Coupling-Routed Teacher Supervision for On-Policy Distillation
Miteto Wei, Xiaohan Wang, Zehao Chen, Jiajun Chai, Sichao Liu, Li Wang, Haoyuan Xu, Zhaoyu Hu and 2 others
Meituan · KTH Royal Institute of Technology
Published September 29, 2026 · on the feed September 30, 2026 · ▲ 91 · arXiv 2609.36601
The 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.
The Brief
2:13A short podcast summary of the paper. Written and voiced by AI.
Transcript
SAKI. Maximal-Coupling-Routed Teacher Supervision for On-Policy Distillation. By Miteto Wei and nine others, from Meituan and the KTH Royal Institute of Technology. Published on the twenty-ninth of September. Its code repository is public and real: eighteen source files.
On-policy distillation trains a student on trajectories it generated itself, which reduces the mismatch between the states seen in training and those seen at test time. The difficulty is a weak student. It may visit prefixes that are misaligned with the teacher, and there the teacher's supervision is less representative.
SAKI, short for Supervision Allocation with KL-constrained Interpolation, combines a teacher-guided rollout constrained by KL divergence with maximal coupling. Coupling produces events at each token: the student's token is accepted, or it is corrected. SAKI reuses those realized events to route supervision, token by token.
Accepted positions keep sampled-token reverse-KL supervision. Correction positions instead receive direct supervision on the teacher's highest-probability token. Under maximal coupling, the probability of a correction is exactly the total variation distance between the two distributions. One trust-region radius therefore does two jobs: it controls how far the rollout deviates, and it upper-bounds how often intervention and specialized supervision occur. The authors also build a speculative verifier that runs inside the engine. It preserves the exact trajectory distribution and coupling semantics, and improves rollout throughput on a matched workload by four point two two times.
Across seven mathematical reasoning benchmarks, SAKI improves on the matched teacher-guided baseline in mean at eight and pass at eight, for students of one point seven billion and of zero point six billion parameters. Placement controls and a fixed-prefix analysis further support correction-triggered routing as a supervision signal that adapts to conflict.
Rather than handing every position identical supervision, SAKI lets the rollout itself mark where student and teacher disagree, and spends the stronger signal there. Because a single radius governs both deviation and correction, how far the student wanders and how often it is pulled back are set by the same dial.
Paper Quest · 7 chapters · about 40 min
Play this paper
Learn it by doing: make the calls the authors made, guess the result before you see it, and defend it in a viva graded on your own words. Chapter 1 is right here.
Chapter 1 🧩 · The problem
A weak student, a strong teacher, and a student that keeps wandering off. That is the wall.
You are teaching a small student model to reason like a much stronger teacher. You let the student write its own answers and ask the teacher what it would say at each step, so training matches what the student sees at test time.
The road ahead
- 1 · The problem ~6 min Letting a weak student learn on its own answers keeps training honest, but the teacher's advice on the student's strange drafts may teach less.
- 2 · What came before ~3 min
- 3 · The key idea ~9 min Lean the rollout toward the teacher within a KL cap, and let each token's accept-or-correct outcome choose how that token is trained.
- 4 · Choose your own research ~3 min
- 5 · Build the method ~8 min Build a KL-capped teacher-guided blend, sample it by maximal coupling, then pick each token's loss from whether it was accepted or corrected.
- 6 · The experiment ~6 min You’ll guess: How much did the engine-resident speculative verifier raise rollout throughput on a matched workload?
- 7 · The viva ~5 min Three questions from the examiners, answered in your own words. Win it for the paper’s trophy.
Free, with Google or an email address. Your stars and progress are kept.