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

TL;DR SMELT 在匹配每 token FLOPs、参数量和 KV cache 的前提下，让 MoE Transformer 中间层循环两次，实现 6.8–18% 训练计算节省，且下游增益随序列长度与上下文示例数放大。

问题

问题背景

当前大模型 scaling 研究聚焦于计算最优训练 与 稀疏激活 的效率提升，循环 Transformer 和 Mixture-of-Experts 是两个独立热点。

现有方法局限

以往循环 Transformer 的评测通常在固定模型深度或参数规模下对比，未严格匹配 per-token FLOPs，导致循环引入的额外计算被误判为架构收益。此外，循环与 MoE 的组合在 非嵌入参数、KV cache 和每 token FLOPs 三方预算匹配下的消融研究几乎空白，也缺少对应的 Chinchilla-style scaling law，使行业难以判断循环带来的是真实深度复用还是变相增加计算量。

为什么这个问题难/重要

要同时匹配这三个预算，需重新规划层数、专家数与隐藏维度，实验网格与计算成本很高；而要得到可信的 compute-optimal frontier 还需训练多个规模。然而，在计算最优曲线上能节省 6.8–18% 的训练 FLOPs，对大规模预训练意味着可观的开销下降，因此该问题具有很强的工程价值。

行业类比：类似代码生成模型处理长上下文时，若能在第二次访问中把注意力从 BOS/分隔符等 sink token 转向与逻辑相关的 token，就能在不增加显存与计算的前提下改善生成质量。

核心洞察

• 在严格匹配 per-token FLOPs、总参数量和 KV cache 的预算下，循环中间层仍能带来计算效率提升。这挑战了以往循环 Transformer 评估中“额外深度等于额外 FLOPs”的隐含假设，证明循环本身具有架构优势，而不仅仅是计算量增加的结果。SMELT 通过在中部 MoE 层循环两次，在计算最优前沿上节省训练 FLOPs，表明缩放定律需要将架构选择作为独立变量，而非仅依赖参数和 token 数量。
• 循环层的第二次访问系统性地减少 attention sink，并将注意力质量重新分配给与内容相关的 token。这一机制发现连接了宏观效率增益与微观注意力动态，解释了为何 SMELT 在长样本和更多上下文示例上增益更大。相比只报告性能提升的同类工作，该分析提供了可操作的归纳偏置线索，可能启发未来注意力机制设计，例如显式抑制 sink token 或分层注意力重用。
• SMELT 的增益在代码等结构化任务上最大，且下游性能提升超过验证损失预测。这表明循环带来了超越 perplexity 的泛化优势，可能源于循环层对算法性推理的归纳偏好。对实际工程而言，这意味着在评估新架构时，仅看验证损失可能低估其价值，需要结合下游任务和样本长度进行多维评估。

方法

方法详解

输入与总览

输入为 token 序列，经过 embedding 得到隐藏状态。SMELT 采用 稀疏 Mixture-of-Experts (MoE) Transformer 架构，核心机制是将中间一半的层循环两次（即 middle layers Loop Twice），从而增加有效深度而不增加参数。

关键模块与设计决策

• 预算匹配：通过“compute-equivalent sparsity”策略，在循环层数增加的同时调整 MoE 专家数量或层数，使得 per-token FLOPs、总非嵌入参数、KV cache 与未循环 Baseline 严格匹配。
• 循环位置与次数：消融表明循环中间一半层优于循环全部层；循环两次优于三次或四次。这暗示模型偏好更大的 有效深度-宽度比。
• 专家路由：MoE 层中每个 token 由 router 选择 top-k 专家，第二次访问时 router 重用核心专家子集并多样化其余部分（据机制分析）。
• 注意力模式：第二次访问显著降低 attention sink，将注意力质量重定向到内容相关 token，这是性能增益的一个内在归纳偏置。

输出与训练

输出为下一个 token 预测的 logits，使用标准语言建模损失训练。作者为 SMELT 和 Baseline 分别拟合 Chinchilla 风格缩放定律，发现 SMELT 的损失随计算量下降更快，在计算最优前沿可节省 6.8%–18.0% 的训练 FLOPs。

与同类方法的差异

以往循环 Transformer 研究通常在固定模型大小下比较，混淆了架构优势与额外 FLOPs；SMELT 在严格匹配计算、参数与缓存预算的前提下，证明循环本身能带来可测量的训练与下游收益。

实验

实验设计

研究团队在 稀疏 Mixture-of-Experts (MoE) Transformer 中引入循环机制, 提出 SMELT(Sparse MoE Transformer, middle layers Loop Twice). 核心设计是循环中间一半层两次, 同时严格匹配每 token FLOPs、非嵌入参数量和 KV cache 三个预算. 通过一系列消融实验, 确定循环中间层优于全循环, 两圈优于三圈或四圈, 且循环 Transformer 偏好更大的有效深度-宽度比. 随后, 将 SMELT 扩展到四个规模(最大 54B 非嵌入参数), 并为每个架构拟合 Chinchilla 风格缩放定律.

关键发现

• SMELT 的验证损失随计算量下降更快, 在计算最优前沿可节省 6.8%–18.0% 的训练 FLOPs.
• 下游基准性能提升超过验证损失预测, 在 Code 任务上增益最大, 且随样本长度和上下文示例数量增加而增长.
• 机制分析表明, 第二次访问层减少了注意力汇 (attention sink), 将注意力质量重定向到内容相关 token, 这可能是性能提升的归纳偏置.

与基线对比

传统 Looped Transformer 评估常固定模型大小, 混淆架构优势与额外 FLOPs. 本研究在预算匹配下对比 unlooped Baseline, 证明循环本身带来实质增益, 而非依赖更多计算. SMELT 在相同资源下实现更优损失下降, 且下游收益超出验证损失预测, 表明循环引入的归纳偏置可能有效提升泛化能力.

行业影响

落地场景

SMELT 的核心优势在于训练算力节省与长上下文/代码任务增益，适用于：

• 大模型预训练：云厂商或 AI 实验室训练 MoE 模型时采用 SMELT 架构，在相同算力预算下获得更优损失曲线，适合从头训练或继续预训练。
• 代码助手与长文档处理：SMELT 在下游代码基准、长样本、多 in-context 示例任务上表现突出，可显著提升代码生成、法律/金融文档分析等产品的准确率。
• 成本敏感的 API 服务：通过匹配 per-token FLOPs 与 KV cache，在同等推理成本下获得更强性能，适合对成本敏感的推理部署。

商业价值

• 降本：训练阶段节省 6.8–18.0% FLOPs，直接降低 GPU 集群租赁或折旧成本，缩短模型迭代周期。
• 体验提升：下游增益超出验证损失预测，尤其在 Code 和长上下文任务上，可提高代码助手付费转化率、降低文档分析错误率。
• 风险可控：循环结构可能影响推理并行度，但论文匹配了 KV cache 和 FLOPs，实际推理成本可控制在同等水平；需结合推理优化确保时延不增加。

与现有产品/工作流的接口

• 模型架构改动：SMELT 仅改变 Transformer 层的复用方式，可在现有 MoE 训练框架（如 Megatron-DeepSpeed、Fairseq）中通过简单修改 layer 循环实现。
• 训练流程：无需改变优化器、数据 pipeline，只需调整层堆叠配置；支持从头训练或继续预训练。
• 推理部署：需处理重复层的前向计算，可复用已有模型并行和专家并行策略；注意循环展开带来的内存和调度开销，可结合编译优化。

具体 use case：例如在 电商搜索/推荐 场景中，使用 SMELT 训练模型处理长商品描述和用户行为序列，并利用多个 in-context 示例提升推荐准确率；在 代码助手 产品中，利用 SMELT 对长代码库上下文和多示例 few-shot 的增益，减少 API 调用成本并提高生成质量。

局限

• **架构特定性**：SMELT 的结论仅在稀疏 MoE Transformer 上验证，未在稠密 Transformer 或其他循环模式（如全面循环、交替循环）上测试。MoE 的动态路由和稀疏激活可能掩盖了循环带来的参数共享效应，因为专家在两次循环中可能被不同 token 激活，这与稠密层共享权重有本质差异。此外，论文仅探索了中间一半层循环两次这一固定策略，未考虑不同层循环次数不同的混合方案，限制了结论的普适性。
• **规模与预算匹配粒度**：最大扩展到 54B 非嵌入参数，虽然拟合了 Chinchilla 式缩放定律，但未在更大规模（如 100B+ 或 production scale）验证。per-token FLOPs 匹配依赖于模型配置的离散化，实际训练中的 wall-clock 时间、显存占用和通信开销未考虑。在 compute-optimal frontier 上报告的 6.8--18.0% FLOPs 节省可能受数据重复、学习率调度等超参数影响，且该增益是否在非 compute-optimal 的实际训练场景中保持未知。
• **下游评估有限**：下游任务集中在 Code、ICL 和 Dyck language 等结构化数据，未覆盖多语言、多模态、复杂推理等。机制分析聚焦于 attention sink 的减弱，但未能完整解释为何第二次循环能带来超越验证损失的增益。此外，循环虽然匹配 per-token FLOPs，但增加了计算图深度，可能不利于推理延迟优化，论文未对此进行工程上的讨论。与更复杂的循环 Transformer 变体（如共享权重但调整位置编码）相比，缺少直接对比。