循环 Transformer 的核心思路,是反复执行共享的层,从而增加模型的有效深度。但如果只在固定模型规模下比较,循环结构带来的优势可能只是因为消耗了更多 FLOPs。SMELT 关注的正是这个容易被忽略的公平性问题:在计算预算对齐后,层循环是否仍然有效?

三项预算同时对齐

论文将循环机制引入稀疏 MoE Transformer,并尽量匹配每个 token 的 FLOPs、非嵌入参数总量以及 KV cache。在多组消融实验后,作者提出 SMELT 配方,即将模型中间一半的层循环执行两次,同时让这三项预算都与不循环的 Baseline 对齐。这样比较的重点就从“更大的模型或更多的计算”转向了架构本身。

作者进一步在四种规模上扩展 SMELT,最大达到 54B 非嵌入参数,并为循环和不循环架构分别拟合 Chinchilla 风格的 scaling law。结果显示,在相同训练计算预算下,SMELT 的损失下降速度更快;在计算最优前沿上,它可节省 6.8% 至 18.0% 的训练 FLOPs。这个优势也会迁移到下游基准,而且并不能完全由验证损失预测。代码任务上的收益最明显,并且会随着样本长度和上下文示例数量增加而扩大。

可能的机制解释

机制分析显示,第二次访问中间层后,模型的 attention sink 会减弱,注意力质量更多转向与内容相关的 token。论文将其视为一种可能解释性能提升的归纳偏置:共享层的再次处理不只是重复计算,也可能改变模型分配注意力的方式。不过,材料并未说明这一机制已经被充分证明为收益的唯一原因,因此更适合把它看作对实验结果的解释线索。

我的判断

这项工作的价值在于把“层循环有效”放回严格的预算比较中检验,避免用额外计算掩盖架构收益。结果对 MoE Transformer 尤其具有参考意义,代码场景也显示出较清晰的应用方向。但它并不意味着循环结构在所有模型、任务或规模上都必然占优:论文只报告了给定架构、预算匹配方式和实验范围内的结果,且机制分析仍是解释性的。对工程实践而言,SMELT 更像一个值得复现实验的结构配方,而不是可以无条件替换现有模型的结论。