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在KL信任域插值的教师引导rollout上，用最大耦合的接受与纠错事件分配token监督，纠错处改为教师Top-1。七项数学基准上，1.7B相对TRB的Mean@8/Pass@8提升1.1/2.9。
On-policy distillation（OPD）让学生在自己滚出来的轨迹上问教师，用来减轻训练前缀和推理前缀不一致。学生很弱时，早期走偏会把后续状态推到教师自己几乎不会去的前缀上。教师在那儿的条件分布仍然有定义，但不一定代表它真正强的推理习惯。
Trust-Region Behavior Blending（TRB）已经在学生中心的KL球里造中间行为分布q，让rollout朝教师靠、又离学生不太远。它改的是「在哪儿问」，每个位置仍用同一套reverse-KL（RKL）：token按学生概率加权，学生几乎采不到的教师偏好词梯度很弱。美团这篇补目标侧：接受学生提案的位置，和必须改提案的位置，不必用同一种监督。
SAKI（Supervision Allocation with KL-constrained Interpolation）把轨迹怎么走和损失怎么分，绑在同一次最大耦合上。
校正事件的含义是「最小干预实现q时，这次提案留不住」。路由信号来自采样过程本身，不用另设冲突阈值。附录里同一套路由下，教师mode比从T随机抽样做目标，在1.7B上再高0.62 Mean@8、1.56 Pass@8。
在线构造q要每步看教师全词表分布。他们把K=8的投机块验证做进推理引擎：提案只提交到第一次拒绝，拒绝处抽精确残差，后续投机状态全部丢掉，轨迹分布和逐token精确q采样一致。匹配负载下引擎内3276 tok/s，相对外部循环776 tok/s为4.22倍，约为纯学生7760 tok/s的42%。
教师是Qwen3-4B-Base在DAPO-Math-17K上用GRPO训出的；学生是Qwen3-0.6B/1.7B-Base。200步、每步64条prompt×8条回复，评测一律学生单独生成。基准覆盖MATH-500、HMMT-Feb26、AIME 2026/2025、AMC 2023、Minerva Math、OlympiadBench。
| 方法 | 1.7B Mean@8 | 1.7B Pass@8 | 0.6B Mean@8 | 0.6B Pass@8 |
| 教师 4B | 41.3 | 57.2 | 41.3 | 57.2 |
| 学生底座 | 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 |
相对TRB，1.7B为+1.1/+2.9，0.6B为+1.2/+2.0。Mean@8在14个学生×基准组合里13个超过TRB；例外是1.7B的AIME25（8.3对TRB的8.8）。相对SKD，1.7B为+4.5/+6.3，0.6B为+6.0/+6.0。ExOPD在1.7B上为27.1/44.4、0.6B上为17.7/33.2，与OPD同一档。
位置对照把每条exact-q轨迹上的教师mode次数钉死：Random-TM得28.30/45.50，按局部TV加权的TV-Weighted-TM得28.24/46.77，SAKI是29.00/47.50。固定2048个前缀的探针上，第50步相对TRB，学生在教师Top-1上的概率C1高4.63个百分点，Top-16质量C16高4.54；关掉校正监督之后到第200步仍高4.26和2.94。相对随机放置，C16优势从最低冲突四分位的+0.25升到最高四分位的+1.32。
这是TRB之上的目标侧补丁，行为混合公式没改。已经在跑教师引导OPD的人，可以复用同一套exact-q rollout，只在校正位置换损失。参考实现已公开。幅度是渐进的：1.7B相对TRB只多1.1 Mean@8，Pass@8多2.9；位置对照说明多出来的点主要来自放在哪，不是多放了几次。
引擎内投机把exact-q从776 tok/s拉到3276。代表训练步里q生成仍占189.5/232.6秒，学生权重同步只占约1.1%。在线教师验证这笔账还在。更适合已经接受这笔rollout成本、又嫌弱学生早期前缀太脏的数学推理蒸馏。
论文写明：系统优化消不掉在线教师验证的固有成本，吞吐仍低于纯学生生成，并随回复长度、校正率和每波提交token数变化。教师学生都锁Qwen3 Base对Base，因为公开checkpoint的终止token约定不同；换instruct或跨系列词表要对齐额外工作。任务全是数学，200步、ε只在前50步非零，没有非数学、更长训程或更大教师的数字。
宏平均赢了，单基准没有全面压过。0.6B的Minerva Pass@8是32.0，低于TRB的34.9和ExOPD的36.0。校正位置的监督目标是教师argmax，写入轨迹的是残差样本，后续前缀和梯度目标可以不是同一个词。论文把它解释成探索与更新解耦，但没有单独消融「监督目标是否必须等于轨迹token」。所有exact-q变体共用TRB那套ε日程，有利于隔离路由效应，也意味着没测过为SAKI单独调信任域会不会更好。