SMELT：算力匹配的混合专家循环 Transformer 缩放定律 SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers
三重预算对齐下，循环中间层两次的 MoE 架构可节省 6.8%–18% 训练算力
前置知识
循环 Transformer（Looped Transformer）
让同一块层（权重完全共享）连续执行多次，以增加每 token 获得的“有效深度”的架构。经典布局是 prelude–recur–coda：非共享的入口块、循环执行的共享中间块、非共享的出口块。循环 $r$ 次时每 token 计算量随执行次数增长，但存储参数不变；循环块内的残差更新通常按 $1/r$ 缩放，以避免多次权重绑定更新把残差流越撑越大。
本文的核心问题就是：在把循环带来的额外 FLOPs 完全补齐之后，权重共享的循环结构是否仍有超出“额外计算支出”的架构净收益。不理解循环的代价结构就无法理解预算匹配的必要性。
混合专家模型（Mixture-of-Experts, MoE）
把 Transformer 的 FFN 层替换为多个“专家”子网络，由路由器为每个 token 只选择 top-k 个（本文 top-8）专家计算，从而把总参数量与每 token FLOPs 解耦：可以存储大量参数但每次只激活一小部分，稀疏度 $S$ 描述未激活的参数份额。
MoE 的解耦特性是本文预算匹配得以成立的关键：循环模型通过收窄隐藏维度支付第二次执行的 FLOPs，再靠增加专家数量找回失去的总容量，这是稠密模型做不到的。
Chinchilla 缩放定律
Hoffmann 等人建立的损失与参数量、数据量之间的幂律关系 $L(N,D)=E+AN^{-a}+KD^{-c}$（不可约项 + 容量项 + 数据项），并给出算力最优的参数/数据分配。本文将其扩展为 $L(F,S,D)=E+A(1-S)^b/F^a+K/D^c$，用每 token FLOPs $F$ 与等效稀疏度 $S$ 替代参数量 $N$ 以适配 MoE。
本文为 Baseline 和 SMELT 各拟合一条缩放面，通过比较前沿指数 $\gamma=ac/(a+c)$ 和定义在最优前沿上的 CE Gain，把架构优势量化为具体算力节省百分比。
KV 缓存与分组查询注意力（GQA）
自回归推理必须缓存每个序列位置的键/值张量，其大小决定了可部署服务的上下文长度。GQA 让多个查询头共享一组 KV 头；通过调整注意力头尺寸与 GQA 分组比例，可以在几乎不改变 KV 缓存大小的前提下调节每 token 计算量。
循环执行中间层会让 token 多次经过注意力层，本文利用头尺寸与 GQA 比例的微调使 SMELT 的 KV 缓存与 Baseline 相差小于 4%，是三重预算匹配的第三根支柱。
注意力汇聚（attention sink）
自回归 Transformer 无论语义内容如何，都会把大量注意力质量集中在序列开头的 token（如 BOS）上，充当“无处可看时的默认去处”。已有文献把它当作训练完成后模型的静态属性，且发现 sink 质量随层数加深而增强。
本文最重要的机制发现是：同一次前向传播内，第二次循环执行会大幅削减 sink 质量并把注意力转向内容相关 token（Dyck 任务上 BOS 质量从 0.60 降到 0.02），这是解释循环增益与 ICL 优势的关键证据。
等效计算稀疏度 $S$
本文定义的度量：以每 token 激活全部专家的 Baseline 为 100% 利用率参照（其每 token FLOPs 与总参数为 $F_0, N_0$），令 $N^{\mathrm{eq}}_{\mathrm{act}}=(F/F_0)N_0$、$S=1-N^{\mathrm{eq}}_{\mathrm{act}}/N$。它用 FLOPs 比值刻画一个 MoE 配置“等效激活”的参数份额，并计入随文档长度变化的注意力成本。
所有实验网格的四个稀疏等级（$S{=}0, 85\%, 95\%, 97\%$）都靠这个量对齐，且它直接进入缩放式的容量项 $(1-S)^b$，是读懂实验设置与拟合系数的前提。
研究动机
循环 Transformer 声称用远小于大模型的参数量达到可比甚至更强的推理与算术能力，但以往评测几乎没有在算力上公平比较过。Huginn（3.5B 参数）和 Ouro（1.4–2.6B）都在固定存储参数量的前提下增加循环深度：把 12 层模型循环成 24 层执行，权重只存一半，每 token FLOPs 和 KV 缓存却是 24 层模型的水平，报告的收益因此混杂了“额外计算”的贡献。反方向的努力也有：Schwethelm 等人在稠密模型上固定每 token FLOPs，发现 $r$ 次循环只相当于 $r^{0.46}$ 个独立层，循环反而更差——但该协议同时压缩了独特参数量，劣势可能来自容量损失而非循环本身；Prairie 等人固定参数量研究循环，但 FLOPs 与 KV 缓存又随循环次数失控增长。三个预算——每 token FLOPs（训练与推理成本）、总参数（知识容量）、KV 缓存（可服务上下文长度）——从未被同时钉住，“循环在算力匹配下是否仍是净收益”悬而未决。
本文的目标是本文的目标是在每 token FLOPs、总非嵌入参数量、KV 缓存三者同时对齐的前提下，严格回答循环 Transformer 是否具备超越额外计算支出的架构优势，并把答案变成可外推的缩放定律而非单点观察。具体路径为：在 200M 规模通过三个系统消融（循环哪些层、模型多深、循环几次）找出最优循环配置，形成可复用的配方 SMELT；把配方扩展到四个规模（最大 54B 非嵌入参数）与四个稀疏度等级，验证优势是否随规模存在；为两种架构分别拟合 Chinchilla 式缩放面，量化在算力最优前沿上 SMELT 能节省多少训练算力；检验优势是否迁移到下游基准、如何按领域/样本长度/上下文示例数分布；最后用机制性探针回答“第二次执行到底算出了什么不同的东西”。
与已有工作不同的是，本文的独特切入是用 MoE 化解三重预算匹配的“三难”困境。稠密模型每 token FLOPs 正比于 $L\cdot H^2$，循环增加执行深度后只能在收窄宽度或减少层数之间取舍，无论哪种都会同时损失参数与 KV 缓存。MoE 把总参数量与每 token FLOPs 解耦：循环模型可以收窄隐藏维度 $H$ 来支付第二次执行的 FLOPs，再靠提高每层专家数量找回失去的总容量，最后单独调整注意力头尺寸与 GQA 比例恢复 KV 缓存，把三项预算的残余失配分别压到 <4%、<1%、<4%。在此之上，本文不做单点对比，而是在 4 规模 × 4 稀疏度的完整网格（0.1B–54B 非嵌入参数）上为两种架构各拟合一条缩放面，把“某个尺寸下更好”升级为“前沿随算力下降更快”的归因明确、可外推的结论，并定义 CE Gain 把架构优势折算成具体的算力节省。这是块级循环方向上第一个同时关闭三重预算的多尺度研究。
核心方法
直觉上，权重绑定的循环把“存储的深度”兑换成“执行的深度”，问题在于这笔交换在账面等价时是否真的划算。技术路线分四步。第一步，建立预算匹配协议：以每 token 全激活的 Baseline 为参照定义等效计算稀疏度 $S$，对每个网格点构造一对模型——SMELT 收窄 $H$、增加专家数、微调注意力头几何，使每 token FLOPs、总非嵌入参数、KV 缓存三个比值都贴近 1。第二步，在 200M 规模做三个消融（循环跨度、有效深度、循环次数），锁定三条规则形成 SMELT 配方：循环中间 50% 的层两次，物理深度取 Baseline 最优值，循环块内残差更新按 $1/r$ 缩放。第三步，把配方铺满 4 规模 × 4 稀疏度网格，用 WSD 调度训练（稳定段 205B token，再从 6 个检查点分叉余弦退火分支、各续训 10B 全新 token），共 32 次训练、96 个匹配对、192 个评估端点。第四步，为两种架构分别拟合 $L(F,S,D)=E+A(1-S)^b/F^a+K/D^c$，在算力最优前沿上计算 CE Gain，并用下游基准残差分析与机制探针交叉验证优势的真实性与来源。
核心创新是“三重预算匹配 + 分架构缩放面”的实验设计及其归因逻辑。与已有工作的本质区别：Huginn、Ouro 等固定参数比较（收益混杂额外 FLOPs）；Schwethelm 等固定 FLOPs 但牺牲参数（劣势可能源于容量损失）；本文利用 MoE 的参数-FLOPs 解耦同时钉住三个预算，使性能差异只能归因于架构本身。配方层面也有三个经消融支撑的决策：其一，只循环中间一半的层——首尾层承担特化角色需要独立参数，50% 跨度处验证损失最低（如 $S\approx95\%$ 时 1.8517，优于全循环的 1.8601）；其二，循环模型可以承受比 Baseline 更大的有效深度-宽度比——物理深度同为 12 但执行 18 层最好，作者假设共享层从多次访问获得多深度梯度信号，使更深的串行计算更易优化；其三，FLOPs 封顶下循环两次最优，三次、四次迫使模型更瘦而退化。加上循环块内残差按 $1/r$ 缩放防止权重绑定更新撑大残差流，以及把 CE Gain 定义为同损失下的算力节省比例 $1-C_{\mathrm{tgt}}/C_{\mathrm{ref}}$，架构优势第一次变成可复现、可外推、带置信区间的量化结论。
方法步骤详情
第一步，训练协议与数据：所有模型共享 AdamW + WSD 调度，稳定段以全局 batch 256 序列（约 1M token/步）训练 196,075 步、205B token；从第 1 万/2 万/5 万/10 万/15 万/196,075 步分叉 6 条余弦退火分支，每条续训 10B 未见过的 token，形成 6 个数据端点；验证集覆盖 Code、Math/STEM、Knowledge、Finance、Web 五类 39 个来源，序列打包为 4096 token 上下文并用段级注意力掩码。第二步，预算匹配：按 $S=1-N^{\mathrm{eq}}_{\mathrm{act}}/N$、$N^{\mathrm{eq}}_{\mathrm{act}}=(F/F_0)N_0$ 定标；以 200M、$S\approx95\%$ 为例，Baseline 为 $L{=}12$、$H{=}1280$、每层 192 专家（1.33e9 FLOPs/token、3.87e9 参数），SMELT 循环中间 6 层两次共执行 18 层，把 $H$ 降到 1056、专家升到 288，最终 FLOPs 1.37e9（+2.9%）、参数 3.89e9（+0.4%）、KV 缓存差 <4%。第三步，消融锁定配方：跨度扫描（0–12 层）在 50% 处损失最低；深度扫描显示 SMELT 最优物理深度与 Baseline 一致（12）而有效深度为 18；次数扫描中 2×（损失 1.9257）优于 3×（1.9385）与 4×（1.9360）。第四步，全网格扩展与拟合：4 规模（100M/200M/600M/1.6B 激活参数，物理深度 10/12/20/30）× 4 稀疏度（$S{=}0,85\%,95\%,97\%$），在 $1.3\times10^{19}$–$2.2\times10^{21}$ FLOPs 窗口内各取 72 个端点，用 Huber 损失（$\delta=10^{-3}$）加 L-BFGS-B 分别拟合六系数缩放式，再由 $C=FD$ 约束解出算力最优 $F^*$、$D^*$ 与前沿 $L^*$，计算 CE Gain 并用 cell bootstrap 给出 95% 区间。第五步，下游与机制评估：DCLM Core/Completion 22 任务（few-shot 任务 10 个随机种子）与 MMLU 5-shot；用 Baseline 的 96 个端点拟合验证损失到各基准分数的四参数 S 形校准曲线（式 6），计算 SMELT 的残差 $\delta_i$；最后在 1M token 验证样本上探测专家路由重叠、残差写入范数与方向、Q/K/V/PV/O 的跨访问余弦、注意力 token 重叠与 sink 变化，并用 Dyck 括号匹配任务做 token 级注意力追踪。
技术新颖性
技术新颖性体现在四个层面。其一，实验设计：首次在块级循环方向同时关闭每 token FLOPs、总参数与 KV 缓存三重预算，并覆盖 0.1B–54B 非嵌入参数的完整缩放阶梯——Table 1 的对比显示此前没有任何工作同时做到这三点，把单尺度点观察变成了可外推的缩放定律结论。其二，方法论：为两种架构分别拟合六个系数（无共享循环指示变量），以同损失算力节省（CE Gain，与算法进步文献中的 $C_{\mathrm{ref}}/C_{\mathrm{tgt}}$ 乘数等价）作为架构进步的标准度量，并配套 cell bootstrap 置信区间，避免固定尺寸比较的混淆。其三，配方层面给出可操作的工程规则：中间 50% 跨度、两次循环、更瘦而更深的几何（收窄 $H$、加专家、调头尺寸/GQA）、残差 $1/r$ 缩放，其中“循环模型可承受更大深度-宽度比”由多深度梯度信号假设支撑。其四，机制层面首次报告单次前向传播内部的动态观察：第二次访问降低注意力 sink 并把质量转向内容相关 token，与把 sink 视为训练后静态属性的既有文献互补；“检索坐标（Q/K）跨访问保留、读取内容（V）更新”的分工也是新证据。
实验结果
核心结论是：在算力公平比较下循环仍是净收益，且优势随规模增长。第一，全网格 16 个 cell 中 SMELT 验证损失全部低于 Baseline；拟合显示 SMELT 的容量指数 $a=0.3892$、数据指数 $c=0.7011$ 均大于 Baseline 的 0.3703 与 0.6594，前沿指数 $\gamma=ac/(a+c)$ 为 0.250 对 0.237（高 5.5%），两者的不可约项 $E$ 差 0.005 nats 小于拟合 RMSE，归为拟合噪声。第二，算力节省：$C=10^{20}$ 时 CE Gain 为 6.8%–10.0%，$C=10^{21}$ 时升至 14.7%–18.0%（外推 $10^{22}$ 约 19.6%–23.5%），前沿差距随算力对数增长（每十倍约 +8 个百分点）；而最优 token/参数分配（TPP）两种架构几乎相同（差 <6%，均远高于 Chinchilla 的约 20），说明收益来自同分配下更低的损失而非重新分配预算。第三，下游迁移超出验证损失预测：DCLM Completion 96/96 全胜、DCLM Core 83/96、MMLU 29/30；残差分析表明验证损失相同时 SMELT 仍系统性占优，如 1.6B 规模 DCLM Completion 残差 +12.3 millinats、Core +1.2pp、MMLU +1.6pp，且残差随规模单调增大。第四，收益集中于结构化数据与长样本：按领域的 CE Gain 为 Code 20.4% > Finance 16.8% > Math/STEM 16.6% > Knowledge 14.9% > Web 14.8%；长样本（512–4096 token）增益是短样本的 1.52 倍（加参数对照 0.98、加专家对照 0.88）；有上下文示例时差距从 0.9pp 扩至 1.9pp，Dyck 语言 $k{=}32$ 时 29.8% 对 26.4%。第五，机制探针给出一致的内部签名：第二次访问重用专家核心子集（高稀疏下重叠 2–3/8，远超随机 $8^2/n$）、写入更大且方向对齐的残差更新（范数比 1.2–3.5×、跨访问余弦 0.56 对无关对 0.16，是放大而非覆写）、保留检索坐标（Q/K 余弦 0.89–0.93；top-8 注意 token 重叠 56–66% 对同层异头对照 28–34%）、并几乎消除注意力 sink（Dyck 头 BOS 质量 0.60→0.02、示例答案质量 0.24→0.85）；sink 削减在通用预训练数据上同样成立，并逆转了 Baseline 随深度增强 sink 的正常趋势。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 验证损失·算力最优前沿（$C=10^{21}$, $S\approx85\%$） | CE Gain（同损失节省的训练 FLOPs） | 18.0% [8%, 28%] | 0%（参照架构） | 节省 18.0% 训练算力 |
| 验证损失·算力最优前沿（$C=10^{21}$, $S\approx95\%$） | CE Gain | 15.8% [10%, 25%] | 0%（参照架构） | 节省 15.8% 训练算力 |
| 验证损失·算力最优前沿（$C=10^{20}$，三个稀疏度） | CE Gain | 6.8%–10.0% | 0%（参照架构） | 节省 6.8%–10.0% 训练算力 |
| DCLM Completion（96 个匹配对，gold 答案 token 交叉熵） | 胜率（越低越好） | 96/96 全胜 | — | 100% 胜率 |
| DCLM Core（96 个匹配对，22 任务 centered accuracy） | 胜率（越高越好） | 83/96 | — | 86.5% 胜率 |
| MMLU 5-shot（Baseline 高于随机 10pp 以上的 30 对） | 准确率胜率 | 29/30 | — | 96.7% 胜率 |
| BB Dyck Languages（$k=32$ 个示例） | 精确匹配准确率 | 29.8% | 26.4% | +3.4 pp |
| 验证损失（200M, $S\approx85\%$，循环次数消融） | 验证损失（nats，越低越好） | 1.9257（2× 循环） | 1.9445（不循环） | -0.0188 nats |
| 按领域 CE Gain（$C=10^{21}$, $S\approx95\%$） | Code 领域 CE Gain | 20.4% | Knowledge/Web 约 14.8%–14.9% | 结构化数据增益最大，高出约 5.5 pp |
局限与改进
作者明确承认的限制有四点：设计消融只在 200M 规模（至 3.9B 非嵌入参数）进行，最优循环跨度或次数在更大规模可能不同；只研究了最简单的循环形式（连续块、完全共享权重），未覆盖每访问 LoRA、自适应递归深度等变体；预算匹配对齐的是算术 FLOPs、参数与 KV 缓存大小而非墙钟成本，串行重复执行与稀疏路由可能带来需要系统级优化弥补的硬件效率差距；机制探针是描述性的，距离因果解释仍有距离。我补充几点观察：其一，每种配置只训练一次，报告的 ±0.1–0.2pp 标准误只反映 10 个评测种子的方差而非训练方差，DCLM Core 输掉 13 对、且跨度扫描中 DCLM 指标与验证损失排序相反（95% 稀疏下 Core 峰值在 67% 跨度而验证损失谷值在 50%），提示下游指标噪声不小；其二，$10^{22}$ FLOPs 的 CE Gain 是超出拟合窗口的外推值，bootstrap 区间宽达 [0%, 51%]；其三，匹配仍有最高 3.9% 的 FLOPs 残余失配与 3.6% 的 KV 失配，$S{=}0$ 因与稀疏 regime 不同被排除在拟合外，缩放面在稠密与极高稀疏情形的适用性未知；其四，全部实验基于内部语料与专有代码栈，外部无法复核数据去污染过滤的效果；其五，按 Baseline 损失四分位分组时 CE Gain 呈 U 形（Q1 18.0%、Q3 13.8%、Q4 回升到 15.9%），最难来源为何回弹作者也没有给出解释。
独立分析的弱点
弱点一：配方的规模外推性未经验证——三个消融全部在 200M 完成，而 Prairie 等已发现最优循环次数随算力增长，$r{=}2$ 在 $10^{22}$+ FLOPs 可能不再最优；改进方向是在 600M 复验跨度与次数消融，或把 $r$ 作为变量纳入缩放律联合拟合。弱点二：墙钟成本被忽视——中间一半的层执行两次意味着该段串行延迟近乎翻倍，若硬件效率差距超过 15% 的算力节省，实践收益会被抵消；改进方向是与编译器/内核协同设计（循环块内算子融合、跨 batch 并行执行两次访问），或引入 token 级自适应递归让简单 token 提前退出以省推理成本。弱点三：可复现性差——模型基于专有代码栈与内部语料，论文明确说明无法发布；改进方向是用开放语料在小尺度复现整条流水线并开源。弱点四：机制与收益的因果链不完整——sink 削减与 ICL 增益的关联只在 Dyck 上通过头消融定位了“Dyck 头”，未做因果干预实验（例如强制保留 sink、或给非循环模型加 sink 抑制正则看收益是否消失）；改进方向是设计这类干预实验。弱点五：评测方差控制不足——单训练种子 + 10 评测种子，无法排除 83/96 中部分胜负来自训练噪声；改进方向是关键 cell 至少 3 个训练种子。弱点六：只覆盖预训练，未研究循环结构对 SFT/RLHF 后训练、指令跟随与对齐任务的影响，而这些都是实际部署的必经环节。
未来方向
作者提出的方向包括：把更丰富的循环变体放到预算匹配协议下检验——每访问 LoRA 适配器（Relaxed Recursive Transformers）、token 级自适应递归深度（Mixture-of-Recursions）、学习式停机（ACT）、块选择共享、跨 token 状态复用——看它们能否保持或放大 CE Gain；在更大算力（如 $10^{25}$ FLOPs）训练以锚定缩放面外推；系统级优化以关闭墙钟差距；以及对第二次访问的收益机制建立因果解释。基于本文成果还可以延伸：其一，循环与推理时计算的组合——循环模型天然支持“多想一步”的测试时扩展，可与 self-consistency、best-of-n 等策略统一建模，形成训练-推理联合的算力分配理论；其二，把 CE Gain 分解到子结构（注意力 vs FFN、哪些头负责 sink 重定向），指导混合架构设计，例如只循环注意力层或只在特定深度区间循环；其三，验证 sink 削减能否被显式正则化移植到非循环模型，从而“免费”获得部分收益；其四，研究循环对长上下文外推、检索增强与多轮对话等输入含丰富可检索结构场景的适配——本文的长样本与 ICL 结果已给出强先兆；其五，把缩放律中的稀疏度项 $(1-S)^b$ 与循环次数联合拟合，给出给定算力下“专家数 × 循环次数”的联合最优分配表。
复现评估
完整复现难度较高，但部分复现可行。不利因素：预训练语料与代码栈均为内部专有，论文明确说明完整架构与训练栈不能发布，也未提及开源模型或代码计划；网格共 32 次训练，最长单次约 215B token，1.6B 激活规模的单次训练 FLOPs 量级达 $10^{21}$，需要工业级算力（估计数千 GPU 日量级）。有利因素：论文在附录 A（Table 9/10）完整公开了 16 个 cell 的配置——隐藏维度、物理深度、专家数、激活/总参数、FLOPs 比值与 KV 比值（残余失配最大 3.9%/1.0%/3.6%）；正文给出训练细节（AdamW、WSD 稳定段 196,075 步、batch 256 序列约 1M token/步、4096 token 打包、段级掩码、6 个退火分支各 10B token）与完整评测协议（DCLM Core 22 任务、few-shot 10 种子、centered accuracy 公式、MMLU 5-shot）；配方本身规则简单明确：循环中间 50% 的层两次、残差 $1/r$ 缩放、收窄 $H$ + 加专家 + 调头几何做预算匹配。务实的复现路径：用开放语料（如 DCLM 数据集）在 100M–200M 规模复刻匹配协议与三个消融，验证“50% 跨度、2 次循环、更大深度-宽度比”的定性结论，这一步用数十 GPU 日即可完成；缩放定律级别的完整复现则需要把网格铺到 600M–1.6B，成本接近一项中型预训练项目。
论文图表