SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers
Shaowen Wang, Ge Zhang, Kairong Luo, Yuhao Wu, Shaofan Liu, Jiaheng Liu, Wenhao Huang, Shen Yan, Jian Li
cs.LG
2026-09-01
在匹配每token FLOPs、总参数和KV cache后，把MoE中间一半层循环两次，缩放律显示最优前沿可省6.8%到18.0%训练算力，下游增益超过验证损失所能预测的幅度。
Looped Transformer 用共享层反复执行来增加有效深度,但多数论文拿固定参数量去比,循环多出来的 FLOPs 和 KV cache 没被扣掉。Schwethelm 等人反过来钉死每 token FLOPs,循环模型又因为少了独立参数而吃亏。真正要回答的问题是:三项预算同时对齐时,循环还有没有结构优势。
Mixture-of-Experts 让这件事变得可做。循环带来的额外层执行,可以用变窄隐层来买单;总参数靠加专家数补回来;KV cache 用更小的头维度和更高的 GQA 比例压住。ByteDance Seed 与清华、TokenWave 把这条路走成一套可缩放配方,名叫 SMELT。
配方来自 200M 主动参数规模上的三组消融,最后锁成三条规则:
循环段的残差更新按 1/r 缩放,避免两次相关写入把残差流撑爆。每个 MoE 层 top-8 路由。训练用内部语料和 WSD 调度,最长约 215B token,四个主动参数规模(100M/200M/600M/1.6B,最大 54B 非嵌入参数)交叉约 85%、95%、97% 三种等效稀疏度,共 96 对匹配端点。两边各自拟合一条带稀疏度修正的 Chinchilla 曲面。
在 10^20 到 10^21 FLOPs 的拟合窗口里,SMELT 的容量指数 a=0.3892、数据指数 c=0.7011,都高于 Baseline 的 0.3703 和 0.6594。前沿指数 γ 高 5.5%。同一损失下,10^20 FLOPs 省 6.8% 到 10.0%,10^21 省 14.7% 到 18.0%。
| 设置 | 指标 | 结果 |
| 10^21 FLOPs, S≈97% | 训练算力 | 少用 14.7% |
| DCLM Completion | 96 对匹配 | 96 胜 |
| DCLM Core | 96 对 | 83 胜 |
| MMLU(Baseline 高于随机 10 点以上) | 30 对 | 29 胜 |
下游增益超过验证损失能解释的部分。按领域看,Code 的算力节省最大。按长度看,512–4096 token 桶的相对增益是 32–256 token 桶的 1.52 倍。16 个少样本任务上,k=0 时差距 0.9 个百分点,有示范后到 1.9。Dyck 括号匹配里,第二次访问把段首 BOS 上的注意力质量从 0.60 降到 0.02,示范答案从 0.24 升到 0.85。机制上,两次访问的路由和注意力位置高度重叠,残差写入变大且方向对齐,更像精炼而不是另起炉灶。
对 MoE 预训练来说,这是一条几乎不增加服务缓存的深度复用路径:参数、FLOPs、KV 都对齐。增益集中在代码、长样本和 in-context learning,跟「第二次访问在做检索精炼」的内部观察一致。这是渐进改进,不是新范式,但两边缩放律斜率已经分开,规模越大纸面节省越多。
配方消融全在 200M(最多 3.9B 非嵌入参数)上做,更大模型上最优循环跨度和次数可能变。匹配的是算术 FLOPs 而不是墙钟,串行重执行和稀疏路由可能在真实硬件上吃掉一部分纸面节省。循环形式也最朴素:连续块、权重完全共享,没有 per-visit adapter 或 token 级停机。10^22 的外推置信区间已经跨过零。每个配置只训一次,DCLM 的误差条只反映评测种子,不反映训练波动。