一句话总结

来自清华大学、字节跳动 Seed 及合作者的研究人员提出了 SMELT，一种计算匹配的混合专家循环 Transformer，它将中间一半的层循环两次，同时匹配每 token FLOPs、总非嵌入参数和 KV 缓存，在计算最优前沿上节省了 $6.8\text{--}18.0\%$6.8–18.0% 的训练 FLOPs，并通过减少注意力汇的归纳偏置，在下游任务上取得了显著提升，尤其在代码方面。

核心贡献

• SMELT，一种循环混合专家 Transformer，将中间一半的层重复两次，同时严格匹配每 token FLOPs、非嵌入参数和 KV 缓存，从而隔离了深度复用的架构效应。
• 扩展实验和 Chinchilla 风格的定律表明，SMELT 在计算最优前沿上节省了 6.8–18.0% 的训练 FLOPs，下游基准测试的提升超过了验证损失预测，尤其在代码、长样本和多样本上下文任务上。
• 机制分析揭示，第二次循环访问持续减少了注意力汇，并将注意力质量重定向到与内容相关的 token，表明循环起到的是精炼步骤的作用，而非仅仅增加容量。

引言

循环 Transformer 重复一个共享的层块，以在不增加参数的情况下增加有效深度，为更强的推理提供参数高效的途径。然而，之前的评估通常将循环模型与参数匹配的基线进行比较，同时允许额外的每 token FLOPs 和 KV 缓存，这使得循环是否提供了超出其额外计算量的架构优势变得不明确。控制 FLOPs 的尝试常常减少了独特参数，混淆了权重共享的效果与容量损失。作者通过使用混合专家，在循环和非循环模型之间同时匹配每 token FLOPs、总参数和 KV 缓存，解决了这一问题。他们在这些预算下搜索循环设计空间，并确定了 SMELT，一种稀疏 MoE Transformer，它以较高的深度-宽度比将中间一半的层循环两次。在多个规模和稀疏度水平上，SMELT 始终比匹配的基线实现更低的损失，在 Chinchilla 前沿上节省计算，并在结构化数据和上下文学习上显示出不成比例的增益，机制证据表明第二次访问放大了第一次访问，并将注意力从汇 token 转移开。

方法

作者设计了一系列稀疏混合专家 Transformer，它们通过简单的循环机制重复使用一个连续的层块。他们将标准的非循环模型称为 Baseline，将重复 $m$m 层跨度 $r$r 次的变体称为循环 Transformer。最终的方案 SMELT 是在严格的计算匹配协议下，通过一组受控的消融实验得出的。

训练协议。 所有模型都是仅解码器的 Transformer，具有稀疏 MoE 前馈层，每个 token 被路由到其 top-8 专家。注意力使用分组查询注意力（GQA）。作者采用了一个内部的 Baseline 家族，并根据其活跃的非嵌入参数数量标记了四个匹配的规模（100M、200M、600M、1.6B）。在循环 Transformer 中，循环跨度内每个子层的残差更新按 $1/r$1/r 缩放，以防止相关的权重绑定更新在重复访问时膨胀残差流。

为了公平比较模型，作者引入了一个计算等效稀疏度指标。给定一个完全激活的控制 Baseline，它激活每个 token 的所有专家，其每 token 训练 FLOPs 为 $F_0$F0​，总非嵌入参数为 $N_0$N0​，任何具有 FLOPs $F$F 和参数 $N$N 的配置由以下特征描述：

$$N_{\mathrm{act}}^{\mathrm{eq}} = \frac{F}{F_0} N_0, \qquad S = 1 - \frac{N_{\mathrm{act}}^{\mathrm{eq}}}{N}.$$Nacteq​=F0​F​N0​,S=1−NNacteq​​.

这里 $N_{\mathrm{act}}^{\mathrm{eq}}$Nacteq​ 通过 FLOPs 比率缩放控制的参数数量，$S$S 是非活跃份额，类似于参数比率稀疏度，但根据 FLOPs 计算以考虑注意力成本。匹配到 $S=0$S=0 的循环 Transformer 并非密集的；它仍然将 token 路由到更大专家池的 top-8 子集。实验报告了在 $S \approx 85\%$S≈85%、$95\%$95% 和 $97\%$97% 下的匹配 Baseline/循环对。

所有运行都使用 AdamW，采用预热-稳定-衰减（WSD）调度。稳定阶段以恒定学习率训练 196,075 步（每步约 1M token，全局批量大小 256）。从每个稳定运行中，在不同步骤启动六个余弦衰减分支，每个分支在额外的 10B token 上衰减，从而为每个配置产生六个 token 时间范围。预训练数据是内部语料库；稳定阶段消耗 205B token，每个分支添加 10B 新 token，因此最长的时间范围达到约 215B token，无重复。匹配对在相同的 token 序列上训练。评估使用 39 个来源的验证损失、DCLM Core 套件（22 个任务，平均中心准确率）、DCLM Completion（自由形式答案任务上的 token 加权交叉熵）和 MMLU 5-shot。已知的评估项从预训练语料库中过滤掉。

匹配作为计算分配问题。 为了将性能差异归因于架构而非资源预算，作者在循环 Transformer 和 Baseline 之间紧密匹配三个量：总参数（知识容量）、每 token FLOPs（训练/推理成本）和 KV 缓存（部署上下文长度）。将 $m$m 层的跨度循环 $r$r 次，将有效深度增加到 $L_{\mathrm{eff}} = L + (r-1)m$Leff​=L+(r−1)m，并按比例提高 FLOPs，而总参数保持不变。为了保持 FLOPs 固定，模型必须变窄（减少隐藏维度 $H$H）或变短（减少 $L$L）。变窄也会缩小每个专家的 FFN，减少总参数；为了补偿，增加每层的专家数量。这些调整被联合应用，使得每 token FLOPs、总参数和 KV 缓存保持在 Baseline 的百分之几以内。匹配过程在每个稀疏度水平 $S$S 上重复。

消融实验与 SMELT 方案。 在匹配框架就绪后，作者在 200M 规模上进行了三项消融实验，以确定最优循环配置。

首先，他们探究哪些层应该循环。固定物理深度为 $L=12$L=12 并使用两次传递，他们扫描了循环的连续中间段的长度，从 0（Baseline）到 12（全循环）。最佳性能来自循环中间一半的层，这一结果在稀疏度水平上均成立。

其次，他们扫描了两种架构的深度-宽度比，同时保持循环跨度为 50% 并使用两次传递。循环 Transformer 在与 Baseline 相同的物理深度下实现了最佳验证损失，但它维持了更大的有效深度-宽度比，因为共享层从多次访问中接收梯度贡献，使得额外的串行计算更容易优化。

第三，他们改变循环传递次数，从一次到四次。额外的访问消耗 FLOPs，在匹配预算下迫使宽度更窄。两次循环产生最佳性能；更多循环会降低结果。

这三项发现定义了 SMELT，一种中间层循环两次的稀疏 MoE Transformer：（1）循环中间一半的层而非整个堆栈，（2）使循环模型的有效深度-宽度比大于 Baseline，以及（3）恰好循环两次。然后，该方案被扩展到更大的模型规模进行主要实验。

实验

SMELT 方案将中间一半的层循环两次，并具有更大的有效深度-宽度比，扩展到 54B 参数，并在匹配的计算预算下与非循环基线进行比较。拟合的缩放定律揭示了 SMELT 更陡峭的计算-损失前沿，节省了 6.8–18% 的训练计算，且随着规模增长而增加。下游基准测试放大了验证损失的改进，尤其是在代码等结构化数据上，并且增益随着样本长度和上下文示例数量的增加而增加。机制分析表明，第二次访问重用了核心专家子集，写出更大但对齐的残差更新，保持注意力模式同时改变值投影，并减少注意力汇以更好地检索相关演示。

循环 Transformer 的研究主要循环中间层，跨度为 33–50%，并发现权重共享改善了推理或困惑度。密集模型通常仅匹配参数数量，FLOPs 和 KV 缓存不受控制，而基于 MoE 的循环模型可以同时匹配参数和每 token FLOPs，取得了强大的基准测试结果。缩放分析表明，重复层的贡献是次线性的（r^0.46 有效贡献），并且最优循环次数随计算量增长，而机制实验揭示，第二次循环访问减少了注意力汇质量，将焦点重定向到信息性 token。中间层循环是主导设计，在所有比较的研究中跨度为 33–50%。只有 MoE 循环模型同时匹配总参数和每 token FLOPs，在 3B 规模上赢得 8/9 基准，在 9B 上赢得 9/9。密集循环模型匹配参数但不匹配 FLOPs 或 KV 缓存，使得循环的额外计算不受控制。密集循环模型的缩放定律发现，r 次重复的贡献类似于 r^0.46 个独特层，并且最优循环次数 r 随计算预算增加。第二次权重绑定访问大幅减少注意力汇质量，将注意力从初始 token 转移到上下文学习任务中的演示答案。

为了隔离架构效应，循环 Transformer 在三个预算上与 Baseline 匹配：通过缩小隐藏维度或层数控制每 token FLOPs，通过增加专家恢复总参数，通过头大小或 GQA 比率对齐 KV 缓存。经过这些调整后，典型的残差不匹配在 FLOPs 和 KV 缓存上保持在 4% 以下，在总参数上保持在 1% 以下。调整隐藏维度或层数通常将每 token FLOPs 的不匹配保持在 4% 以下。增加专家将总参数恢复到基线的 1% 以内。

当循环跨度接近 50% 时，两个稀疏度水平的验证损失最小化，因此采用 50% 跨度作为默认值。DCLM 指标并不始终跟踪验证损失，使得验证损失成为更可靠的选择标准。循环中间块优于全栈循环，这与初始层和最终层的专门角色一致。在 85% 和 95% 稀疏度水平下，验证损失在 50% 循环跨度处达到最低点。在 95% 稀疏度下，DCLM Core 在 67% 跨度处达到峰值，而验证损失在 50% 跨度处最佳，说明了指标分歧。

在匹配的计算和参数预算下，物理深度为 12、有效深度为 18 的循环 Transformer 在验证损失、DCLM Core 准确率和 DCLM Completion 困惑度上均优于所有 Baseline 配置。Baseline 在物理深度 12 处实现最佳验证损失，循环 Transformer 的最优物理深度与此匹配，而其更大的有效深度表明了对更高有效深度-宽度比的偏好。这表明通过多次传递重复使用层，比同等深度的独立层堆栈能够实现更多的串行计算。具有 12 个物理层、执行 18 次的循环 Transformer 在所有测试配置中实现了最佳验证损失、最高 DCLM Core 和最低 DCLM Completion。Baseline 的验证损失最优值出现在物理深度 12；循环 Transformer 匹配此物理深度，但在有效深度为 18 时达到峰值性能，显示出更大的有效深度-宽度比。

在匹配的 FLOPs 下，将中间层循环两次产生最低的验证损失和最高的 DCLM Core 分数，优于单次访问和更深的循环次数。添加第三次或第四次访问迫使模型变薄，导致所有指标的性能下降。两次访问实现了最佳验证损失和 DCLM Core，超过了基线和三次、四次访问配置。三次和四次访问相对于两次访问均出现退化，且它们之间没有一致的顺序，表明额外循环迫使的较薄模型损害了性能。

实验在受控的计算、参数和 KV 缓存预算下评估了循环 Transformer，比较了密集和混合专家变体。将中间 50% 的层循环两次，产生了最佳的验证损失和下游性能，优于单次访问和更深的重复，同时减少了注意力汇质量。混合专家循环模型同时匹配总参数和每 token FLOPs，在 3B 和 9B 规模上几乎赢得了所有基准测试，而密集循环模型仅匹配参数，并表现出次线性缩放，其中重复层的贡献类似于 r^0.46 个独特层。总体而言，最优循环次数随计算预算增加，重复使用层使得有效深度-宽度比高于堆叠独立层。