SMELT: 计算匹配的MoE循环Transformer的缩放定律

摘要

循环Transformer通过迭代一组共享层来增加有效深度，但大多数评估在固定模型规模下进行比较，从而将架构优势与额外FLOPs混为一谈。我们在混合专家Transformer上研究循环，同时严格控制每词元FLOPs、总非嵌入参数和KV缓存与基线匹配。通过一系列消融实验，我们得出一种称为SMELT（Sparse MoE Transformer，中间层循环两次）的方案：在三项预算均与无循环基线匹配的条件下，将中间一半的层循环两次。我们将SMELT扩展到四种规模，非嵌入参数最多达54B，并为每种架构分别拟合了一条Chinchilla式缩放定律。随着计算量增加，SMELT的损失下降更快，在计算最优前沿上可节省6.8%–18.0%的训练FLOPs。这一优势可迁移至下游基准，且幅度超过验证损失的预测范围；在代码任务上最为显著，并随样本长度和上下文示例数量的增加而增强。机制分析表明，第二次访问降低了注意力下沉现象，并将注意力权重重新分配至内容相关词元上，这一归纳偏置可能是上述性能提升的基础。这些结果表明，即使在预算匹配条件下，循环也能改进Transformer，为将深度复用转化为可测量收益提供了一种实用方案。

English

Looped Transformers increase effective depth by iterating a shared block of layers, but most evaluations compare at fixed model size, conflating architectural advantage with extra FLOPs. We study looping on Mixture-of-Experts Transformers while closely matching per-token FLOPs, total non-embedding parameters, and KV cache. Through a series of ablations, we arrive at a recipe we call SMELT (Sparse MoE Transformer, middle layers Loop Twice), which loops the middle half of layers twice while matching the unlooped Baseline on all three budgets. We scale SMELT across four sizes up to 54B non-embedding parameters and fit a separate Chinchilla-style scaling law for each architecture. SMELT's loss drops faster with compute, saving 6.8--18.0\% of training FLOPs on the compute-optimal frontier. The advantage transfers to downstream benchmarks beyond what validation loss predicts, is largest on Code, and grows with sample length and the number of in-context examples. Mechanistic analysis shows that the second visit reduces the attention sink and redirects mass toward content-relevant tokens, an inductive bias that may underlie the observed performance gains. These results show that looping can improve Transformers even under budget matching, offering a practical recipe that turns depth reuse into measurable gains.

SMELT: 计算匹配的MoE循环Transformer的缩放定律

SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers

摘要

Support