同等预算下,Looped Transformer更强吗?清华、字节Seed给出答案

在大模型的发展中,提升能力的主要手段一直是 "做大"—— 更多参数、更多数据、更多算力。如今,另一条思路开始受到关注:与其不断堆叠新的层,不如让已有的层再跑一遍。这就是 Looped Transformer。普通 Transformer 的每一层只执行一次,而 Looped Transformer 会把其中一部分层重复执行,让模型在不增加参数的情况下获得更深的计算。
这个想法在推理和数学等任务上已经展现出不小的潜力,但一个关键问题始终没有被正面回答:Looped Transformer 的收益,究竟是循环架构本身的优势,还是因为多算了几遍、实际上用了更多的算力?清华大学、字节跳动 Seed 等机构的研究团队在论文 SMELT 中正面回答了这个问题。他们同时对齐算力、参数量和 KV cache,首次在多个模型尺度上进行了公平比较。

论文链接:https://arxiv.org/abs/2609.01343
本文的第一作者为清华大学博士生王少文,师从李建教授,研究方向为大语言模型预训练,曾在 ICML、NeurIPS、ACL、EMNLP 等国际会议上发表多篇论文,并获 LIT@ICLR 2026 最佳论文奖。
预算匹配框架
公平比较需要同时控制三个变量:决定训练和推理成本的每 token FLOPs、约束知识容量的总参数量、以及限制可服务上下文长度的 KV cache。
论文利用 MoE 将参数量和计算量解耦,具体的匹配方式如下:1)缩窄或变浅模型来抵消 Loop 架构循环执行的计算,2)增加专家数来补全总参数量,3)再调整注意力配置(GQA Ratio / Head Dim)来匹配 KV cache。

图 1 (a)|SMELT 的核心做法:将中间 50% 的层重复执行两次,同时匹配计算量、参数量和 KV cache。
最优循环配方
按上述控制预算的规则,论文在小规模上依次搜索循环范围、执行深宽比和循环次数,得到了以下结论:
1)循环中间 50% 的层收益最大。 论文把循环层比例从 0%(不循环)一路调到 100%(全部循环)。结果显示,比例在 50% 左右时,验证损失最低。作者推测,模型首尾的层分别承担输入编码和输出预测等特殊功能。因此,这些层保留独立参数,可能比参与共享更有利。

表 3|循环中间约 50% 的层效果最好。
2)Looped 模型的最优执行深宽比高于非 Looped 模型。固定 50% 的循环范围后,论文进一步比较了不同的执行深宽比(执行深度与模型宽度之比)。结果显示,Looped 模型的最优执行深宽比更大。论文推测,这可能与参数共享有关。额外的执行深度不会引入新的参数;同一组共享参数还会同时获得来自 “浅层” 和 “深层” 的梯度优化信号。因此,更深的配置也更容易训练。

表 4|Looped Transformer 更适合执行更深的配置。
3)循环两次收益最大。 论文在一到四次之间扫描循环次数,发现两次最优。由于 FLOPs 匹配的约束,每多循环一次模型就要再缩窄一些,因此三次和四次反而因模型过窄而退步。

表 5|循环两次效果最好;继续增加循环次数反而退步。
综合以上三条,论文将最终的配方命名为 SMELT:在稀疏 MoE 模型中,将中间一半的层循环执行两次,同时采用更大的执行深宽比。
Scaling Law 拟合额外收益
论文将 SMELT 配方扩展到四个模型规模(100M / 200M / 600M / 1.6B 激活参数),并在每个规模下设置四个稀疏度级别(通过增加专家数来提高总参,同时保持每 token 激活的参数量不变),最大总参达到 54B。每个组合都训练了一对匹配的基线与 SMELT。结果发现:在全部 16 个匹配配置上,SMELT 都取得了更低的验证损失。

图 3|在所有模型规模和稀疏度下,SMELT 都能以相近计算量获得更低的验证损失。
为了量化优势随规模的变化趋势,论文为两种架构分别拟合了Chinchilla式Scaling Law:
。其中 E 是不可约损失;中间的容量项由每token FLOPs F 和稀疏度 S 控制,F 越大或 S 越高,这一项越小;最后的数据项由训练token数 D 控制,数据越多越小。基线和SMELT各自拥有独立的六个系数(E, A, K, a, b, c)。

表 6|分别拟合基线和 SMELT 得到的 Scaling Law 系数。
拟合结果显示,SMELT 的容量指数 a(0.3892 vs. 0.3703)和数据指数 c(0.7011 vs. 0.6594)都大于基线,意味着 SMELT 的 loss 随算力和数据的增长下降得更快。在 Compute-Optimal 前沿上,SMELT 达到相同 loss 所需的训练算力始终更少,节省幅度在 6.8%–18.0% 之间,且随算力预算增大而持续扩大。

图 1 (b)|SMELT 的计算最优前沿更低;在 10²¹ FLOPs 时可节省 14.7% 的训练计算量。
下游任务收益
论文评测了 DCLM 和 MMLU 等下游任务,SMELT 在几乎所有匹配配置上都取得了更高的分数。

图 6|SMELT 在 96 组配置中的 83 组上取得更高的 DCLM Core 分数。

图 7|在明显高于随机水平的 30 组配置中,SMELT 在 29 组上取得更高的 MMLU 准确率。
更值得关注的是,这些下游收益不仅仅来自 loss 的降低。论文先用基线数据拟合了验证损失与下游分数的关系曲线,然后把 SMELT 的数据点放到同一张图上。结果发现,在相同的验证损失下,SMELT 的下游分数始终高于基线,而且这一差距随规模增大而增大。

图 9|即使验证损失相同,SMELT 在多数下游任务上仍有额外收益,而且通常随模型规模增大。
论文进一步分析了循环的收益在不同领域和任务上的分布。
在验证损失上,结构化程度越高的领域,收益越大:Code 领先,Finance 和 Math/STEM 次之,Web 最低。同时长文本的收益也显著高于短文本。作为对照,增加专家数或参数量并不会产生这种长度上的差异,说明这是循环带来的特有性质。

图 10 (a)|SMELT在所有数据领域上都有收益,其中代码数据收益最大。

图 11(a)|文本越长,SMELT相对基线的收益越明显。
在下游任务上,给出的 in-context 示例越多,SMELT 相对基线的优势越大,在高度依赖示例的 Dyck Languages 任务上表现得尤为突出。

图 11 (b–c)|上下文示例越多,SMELT的优势越明显;这一趋势在Dyck Languages上尤其突出。
循环内部机制
接下来,论文从专家路由、残差流和注意力模式三个角度分析了循环层第二次执行与第一次执行的差异。
专家路由:第二次会复用第一次的部分专家,两次选中的专家集合重叠度远高于随机水平。

图 12|第二次执行会复用部分专家,同时选择一些新专家;实际重叠远高于随机水平。
残差流更新:第二次执行的残差流更新幅度始终大于第一次。两次更新的方向有所重叠,但并不相同。这说明,第二次执行在第一次的基础上作了更大幅度的修正。

图 13|第二次执行向残差流写入的更新整体更大。
注意力检索:两次执行中,Q 和 K 各自的余弦相似度高达 0.89–0.93,说明模型关注的位置大体相同。变化主要发生在 V 及其后续输出。换句话说,第二次执行仍然 “看向同一处”,但 “读到的内容” 已经不同。

图 16|两次执行通常看向相似的位置,但读取和写回的内容发生了更大变化。
Case Study on Dyck Language:Dyck language 是一个括号匹配任务,模型需要从给定示例中找到相关答案的位置,才能作出正确预测。论文发现,关键注意力头在第一次执行时集中关注 BOS,而相同的头第二次执行则关注答案位置。

图 18|在 Dyck 任务中,第二次执行把注意力从 BOS 转移到示例答案。
在基线模型中,attention sink 会随深度增加而增强。Loop 架构却呈现相反趋势:第二次执行后,attention sink 反而减弱。论文据此推测,第一次执行已经建立了初步表征,第二次就不再需要依赖 sink,转而把注意力用于实际的内容检索。

图 19 (b)|普通模型的 attention sink 随深度增强,而 SMELT 的第二次执行会减弱 attention sink 。
小结
论文在预算匹配的条件下,从 Scaling Law、下游评测和机制分析三个层面给出了一致的结论:Loop 架构具有独特的优势。现有的局限包括消融实验规模有限以及未涉及实际运行时间层面的系统效率优化。
未来,这套预算匹配方案和 SMELT 配方有望成为 Loop 社区的一套标准实验设置。以此为基础,后续工作可以在公平条件下继续探索更多架构变体。