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, Wei Lin, Guojun Yin
cs.AI
2026-09-29
SAKI routes reverse-KL vs teacher Top-1 using maximal-coupling accept/correct events; 1.7B Mean@8/Pass@8 rises from TRB's 27.9/44.6 to 29.0/47.5 on seven math benchmarks.
On-policy distillation (OPD) trains a student on prefixes it actually generates, which cuts the train-test state mismatch of teacher-forced or frozen trajectories. When the student is much weaker than the teacher, early errors compound. Later prefixes can sit in regions the teacher would almost never visit, so the teacher's conditional signal is well-defined but a poor match for the behavior that makes the teacher strong.
Trust-Region Behavior Blending (TRB) already builds an intermediate rollout distribution q inside a student-centered KL ball, so visited states move toward the teacher without jumping onto teacher trajectories. TRB still applies the same reverse-KL (RKL) objective at every token. RKL weights tokens by student mass, so teacher-preferred tokens with little student support get weak gradients. SAKI asks a complementary question: positions where the guided policy can keep a student proposal may not need the same update as positions where it has to override that proposal.
SAKI (Supervision Allocation with KL-constrained Interpolation) ties trajectory construction and token-level losses to one maximal coupling.
A correction is the event that the cheapest coupling realizing q cannot keep this proposal. The router is the sampler, not an extra conflict heuristic. Under the same routing rule, teacher-mode beats sampling the target from T by +0.62 Mean@8 and +1.56 Pass@8 on the 1.7B student.
Exact q needs the teacher distribution at every step. An engine-resident speculative verifier drafts K=8 tokens, commits through the first rejection, samples the exact residual there, and drops the rest of the block. That matches sequential exact-q sampling. Under a matched 64×8 workload it runs at 3,276 tok/s versus 776 for the external loop (4.22×), about 42% of student-only generation at 7,760 tok/s.
The teacher is Qwen3-4B-Base trained with GRPO on DAPO-Math-17K. Students are Qwen3-0.6B-Base and Qwen3-1.7B-Base. All methods share 200 steps, 64 prompts × 8 responses, and student-only evaluation. Benchmarks: MATH-500, HMMT-Feb26, AIME 2026/2025, AMC 2023, Minerva Math, OlympiadBench.
| Method | 1.7B Mean@8 | 1.7B Pass@8 | 0.6B Mean@8 | 0.6B Pass@8 |
| Teacher 4B | 41.3 | 57.2 | 41.3 | 57.2 |
| Student base | 3.0 | 16.4 | 1.7 | 10.3 |
| OPD | 27.3 | 45.0 | 17.1 | 31.6 |
| SKD | 24.5 | 41.2 | 12.4 | 29.6 |
| TRB | 27.9 | 44.6 | 17.2 | 33.6 |
| SAKI | 29.0 | 47.5 | 18.4 | 35.6 |
Versus TRB, the 1.7B student gains +1.1 Mean@8 and +2.9 Pass@8; the 0.6B student gains +1.2 and +2.0. Mean@8 beats TRB on 13 of 14 student-benchmark pairs; the miss is 1.7B AIME25 (8.3 vs 8.8). Versus SKD the gaps are +4.5/+6.3 and +6.0/+6.0. ExOPD sits with OPD at 27.1/44.4 on 1.7B and 17.7/33.2 on 0.6B.
Placement controls freeze the number of teacher-mode updates per exact-q trajectory. Random-TM scores 28.30/45.50. TV-Weighted-TM, which samples positions by local TV(p,q), scores 28.24/46.77. SAKI scores 29.00/47.50. On a frozen probe of 2,048 prefixes, step 50 already shows +4.63 pp on teacher Top-1 probability (C1) and +4.54 pp on Top-16 mass (C16) over TRB; after corrections stop at step 51, step 200 still holds +4.26 and +2.94. The C16 edge over random placement grows from +0.25 pp in the lowest initial-conflict quartile to +1.32 in the highest.
This is an objective-side patch on TRB, not a new mixing formula. Teams already running teacher-guided OPD can keep the same exact-q engine and change the loss only at correction indices. A reference implementation is public. The lift is incremental: +1.1 Mean@8 on 1.7B against a strong TRB baseline, with a larger +2.9 on Pass@8. The placement controls say those points come from where teacher-mode lands, not from using it more often.
Engine-resident speculation moves exact-q from 776 to 3,276 tok/s. In a representative training step, q generation still takes 189.5 of 232.6 seconds; student weight sync is about 1.1%. Online teacher verification does not go away. The method fits math-reasoning distillation where that rollout cost is already acceptable and weak-student prefixes are the pain point.
The paper states that the optimized backend cannot remove the intrinsic cost of online teacher verification. Throughput stays below student-only generation and still depends on response length, correction rate, and committed tokens per wave. Teacher and student are Qwen3 Base-to-Base because public checkpoints disagree on termination tokens; instruct models or cross-family vocabularies need extra alignment. All tasks are math. Training is 200 steps with ε nonzero only for the first 50. There are no numbers for non-math domains, longer runs, or larger teachers.
Macro averages win; per-benchmark results do not. On 0.6B Minerva, Pass@8 is 32.0, below TRB's 34.9 and ExOPD's 36.0. At correction positions the training target is the teacher argmax while the token written into the trajectory is a residual sample, so later prefixes and the local gradient can point at different tokens. The paper calls this a split between exploration steering and parameter updates, but never ablates whether the supervision target must equal the trajectory token. Every exact-q variant reuses TRB's ε schedule, which isolates routing from behavior tuning and also means SAKI never got its own trust-region sweep.