跳到主要内容
返回时间线
arXiv来源发表:

SAKI 用最大耦合的接受/纠正事件路由教师监督,在七个数学推理基准上把 1.7B 与 0.6B 学生的 Mean@8/Pass@8 同时推高到同规模最佳

核心概要

SAKI 在 KL 约束的教师引导 rollout 中用最大耦合实现中间行为分布,并把实际发生的接受/纠正事件当作 token 级监督路由器:接受位置保留采样 token 的反向 KL,纠正位置改为直接监督教师最高概率 token,纠正概率恰为 TV(p_t,q_t) 因而受同一信任域半径上界约束;配套的引擎内投机验证器保持精确 q 轨迹分布与耦合语义,把同等工作负载下的 rollout 吞吐提升 4.22 倍,并在七个数学推理基准上让 1.7B 与 0.6B 学生的 Mean@8 与 Pass@8 均优于匹配的教师引导基线。

AI-generated editorial illustration: SAKI: Maximal-Coupling-Routed Teacher Supervision for On-Policy Distillation

深度剖析

SAKI 把最大耦合产生的接受/纠正事件直接当作监督路由器:接受位置沿用采样 token 的反向 KL,纠正位置切换为对教师 Top-1 token 的负对数似然监督。 此前的教师引导 rollout(如 TRB)只改变监督被查询的状态分布,仍保留同一套逐前缀反向 KL 目标;SAKI 则利用实现该引导分布时天然暴露的耦合事件,把“在哪里监督”与“监督什么”解耦,无需额外的 token 选择启发式或路由阈值。 论文给出 Proposition 1 与 Proposition 2:纠正概率恰为 TV(p_t,q_t),且是这些边缘分布所有耦合中的最小干预概率;Corollary 1 进一步说明信任域半径同时约束 rollout 偏离并上界干预与专门监督频率。附录 J.2 报告退火过程中观测纠正概率低于 Pinsker 上界并在半径归零时消失。

在七个数学推理基准上,SAKI 在 1.7B 与 0.6B 两个学生规模上都取得最佳宏平均 Mean@8 与 Pass@8。 相对共享同一精确 q rollout 与退火调度的 TRB,1.7B 学生从 27.9/44.6 提升到 29.0/47.5,0.6B 学生从 17.2/33.6 提升到 18.4/35.6;相对 SKD 的增益更大。Mean@8 在 14 个学生–基准组合中有 13 个优于 TRB。 学生为 Qwen3-0.6B-Base 与 Qwen3-1.7B-Base,教师为同一 Qwen3-4B-Base-GRPO,训练数据为 DAPO-Math-17K,200 步预算,rollout 批为 64 个提示各 8 条回答,评测在 MATH-500、HMMT-Feb26、AIME 2026、AIME 2025、AMC 2023、Minerva Math、OlympiadBench 上以温度 1.0、最大长度 16,384 采样八条回答。

纠正触发的位置选择本身携带信息:在同等教师模式监督预算下,它优于随机放置与按 TV 加权的放置。 Random-TM 达到 28.30 Mean@8 / 45.50 Pass@8,TV-Weighted-TM 达到 28.24 / 46.77,SAKI 达到 29.00 / 47.50;说明标量冲突度量有用,但已实现的耦合纠正提供更强的路由信号。 两个对照在每条精确 q 轨迹上匹配 SAKI 的教师模式更新数量;TV-Weighted-TM 中纠正掩码只决定每条轨迹的监督预算而不决定位置。固定前缀探针(2,048 个位置、213 个提示)显示优势从最低冲突四分位到最高冲突四分位递增,Q4–Q1 差为 0.87 至 1.29(95% 置信区间),主要由教师支持的非 argmax token 承载。

引擎内投机块验证在保持精确 q 轨迹分布与耦合语义的同时,把同等工作负载下的 rollout 吞吐从 776 tokens/s 提升到 3,276 tokens/s。 相对外部循环实现为 4.22 倍加速,并达到仅学生生成理论上限 7,760 tokens/s 的 42% 以上;附录 K 的实现演进显示逐 token 原型仅 5 tokens/s、早期外部块路径为 165–299 tokens/s。 Proposition 3 证明:只要提案前缀与顺序最大耦合一致、只提交到首次拒绝、纠正从精确残差分布采样并丢弃后续投机状态,生成轨迹与顺序采样同分布;附录 K.5 从桥构造、耦合、训练状态与路由四个维度做正确性校验。

启示与展望

该结果面向以数学推理为验证场景的 on-policy 蒸馏:学生为 Qwen3-0.6B-Base 与 Qwen3-1.7B-Base,教师为同一 Qwen3-4B-Base-GRPO,训练数据为 DAPO-Math-17K,200 步预算,评测覆盖 MATH-500、HMMT-Feb26、AIME 2026、AIME 2025、AMC 2023、Minerva Math 与 OlympiadBench。方法本身不依赖数学内容,因此可直接用于其他需要在线教师查询的蒸馏流程;引擎内投机验证器为需要精确 q rollout 的实现提供了可复用的执行结构。主训练调度中纠正监督是暂态的:信任域半径退火到零后,方法回到纯学生 rollout 的反向 KL 训练,因此该机制被设计为训练早期的引导手段。

纠正监督在半径退火到零后停止,固定前缀探针显示对齐优势在随后 149 步纯反向 KL 训练中仍然可见,但这一持久性在更长训练与更大模型上的表现仍是开放问题。放置对照与固定前缀分析目前主要在 1.7B 学生上展开,0.6B 规模是否呈现同样的冲突自适应模式尚待确认。吞吐数字依赖响应长度、纠正率与每次验证波次提交的 token 数,换用不同工作负载时加速比可能变化。此外,本证据包为全文解析,部分表格中的具体百分点数值在文本中以占位形式出现,因此对个别幅度只能依据正文叙述而非逐项数字核对。

来源