Looping Beyond Twice: A Scalable Recipe for Looped Mixture-of-Experts
Di He, Pengxiang Li, Da Chang, Qingyan Meng, Lu Yin, Shiwei Liu
cs.LG, cs.AI
2026-10-01
循环 MoE 过了两次容易失稳。LOOM 用残差缩放、输入回注和分轮路由稳住它：等算力下 700M、5 轮的验证困惑度从 18.36 降到 16.54，不锁算力的 1.7B、9 轮从 9.62 降到 7.77。
同一组 Transformer 块沿深度重复执行，有效层数增加，参数不随轮次复制。这叫 looped Transformer。稠密模型把它当作参数量和数据之外的缩放轴。搬到 MoE 上，每个 token 只激活少数专家，多循环几轮等于用同一套专家反复改表示。
收益掉得很快，再加轮次还会变差，已有 looped MoE 大多停在 2 次。curse of depth 会被循环放大：Pre-LN 把归一化放在残差支路里，残差逐轮累加，隐藏状态方差变大，表示漂移。另一件是 expert selection collapse。各轮共用一个 router，由它给 token 挑专家，选出的专家几乎不变，多花的计算没有新组合。
原则是每一轮贡献新计算，同时稳住循环状态。骨干是 Llama 式 pre-norm MoE decoder，一组 M 层重复 H 遍。注意力和专家跨轮共享，router 按轮分开。有效深度 M×H，新增参数主要是 router。
λ=0.5 固定、不学习。残差乘 γ=λ/√(H×M)，注意力和 MoE 输出都乘它，MoE 支路再套 RMSNorm。1/√H 管跨轮累积，1/√M 管层间方差。只压方差不够，深循环仍会把输入冲淡，所以从第 2 轮起按 gt=λ/√(t×M) 混入输入嵌入：新状态=(1−gt)×上一轮输出+gt×输入。gt 随轮次变小。
专家选择靠分轮 router。每层每一轮一个 sigmoid router，float32 计算，专家共享。约 350M、9 轮上，分轮之后专家负载的跨轮相似度下降、波动变大。Looping Residual 用 β=0.5 的指数滑动平均保存缩放后的注意力输出，全局记忆跨轮保留，局部记忆每轮清零，并用这两个摘要替换注意力的直接残差，然后再做 MoE 残差。
反传以 K=3 分段。每当轮次是 3 的倍数或已到最后一轮，就计算语言模型损失、反传、更新，再截断状态与全局记忆的梯度。9 轮对应一条序列上的 3 次监督和 3 次更新。
数据是 FineWeb-Edu，序列长度 1024。零样本七项：OpenBookQA、WinoGrande、ARC-Challenge、ARC-Easy、HellaSwag、SIQA、PIQA。
等 FLOPs 的 700M 骨干为 M=10、隐藏维度 512、80 个路由专家、无共享专家，训练 10B token，全程反传。f=H×(6+3k)（单位 d²），不循环时 top-k=26、f=84，加轮次就减小 top-k。
| 设置 | 验证困惑度 | 七项平均 |
| 不循环，k=26 | 18.36 | 38.84% |
| 2 轮，k=12 | 17.37 | 39.00% |
| 3 轮，k=8 | 16.91 | 39.34% |
| 4 轮，k=5 | 16.58 | 39.53% |
| 5 轮，k=4 | 16.54 | 39.53% |
| 6 轮，k=3 | 16.57 | 39.50% |
5 轮困惑度 16.54，低于不循环的 18.36 和 2 轮的 17.37。平均准确率 4 轮和 5 轮都是 39.53%，6 轮 39.50%，相对 38.84% 高 0.69 个点。5 轮上 WinoGrande 从 51.38% 降到 50.20%，SIQA 从 35.72% 降到 34.65%。3、5、6 轮 f=90，约多 7% 乘加。
top-k 不减时，轮次就是额外 FLOPs。350M（激活约 141M）、10B token：不循环 20.07 / 37.96%，9 轮 14.80 / 41.15%，12 轮 14.86 / 41.42%。100M（激活约 70M）、5B token：不循环 26.26 / 36.73%，12 轮 19.26 / 39.11%。350M 上不加任何技巧，6 轮困惑度到 1312；只缩放或只回注在 3 轮已是 23.80 和 42.76，基线为 20.07，更深出现不可恢复的 loss 尖峰。
1.7B（激活约 0.63B）、60B token、top-k=8：
| 轮次 | 有效深度 | 验证困惑度 | 七项平均 |
| 1 | 15 | 9.62 | 42.4% |
| 6 | 90 | 7.91 | 46.7% |
| 9 | 135 | 7.77 | 47.7% |
| 12 | 180 | 7.84 | 47.5% |
9 轮最好，困惑度 9.62 降到 7.77，平均准确率 42.4% 升到 47.7%。HellaSwag 37.1% 到 48.8%，ARC-Easy 50.2% 到 57.7%，12 轮回落。
第 5000 步、350M、9 轮的消融：完整配方 18.62。去掉残差缩放、输入回注、MoE RMSNorm，变为 24.63、26.05、26.50。去掉 Looping Residual 为 19.25，共享 router 为 19.83。
分段反传用于非等 FLOPs 实验。350M、microbatch 为 1，轮次 3 到 12 时分段 NPU 显存约 12.6 GiB，全程反传从 12.5 GiB 到 19.3 GiB；12 轮时间从 26.1 小时到 12.6 小时。9 轮全程反传困惑度 81.04、准确率 34.3%，分段为 14.80 和 41.1%。全程反传在最深设置缩小了 micro-batch，显存并未严格对齐。
等 FLOPs 能成立，是因为基线 top-k 宽到 26，可以把「多激活专家」换成「同一个块多走几遍」。top-2 或 top-8 的 MoE 没有这个额度，循环直接按轮数乘 FLOPs。在这个交换里，4 到 5 轮困惑度好于不循环，也好于 2 轮；下游只多 0.69 个点，5 轮后持平。
1.7B 的 9 轮多出 5.3 个点，有效深度从 15 增到 135。1.7B×9 被记成大约 15B 的展开规模。参数仍是 1.7B，旁边没有同 FLOPs 的非循环模型。
λ、β、K 在 100M 到 1.7B 都是 0.5、0.5、3。训练侧更实质的收获是分段反传：12 轮时间大约减半，显存不随轮次上涨。推理若本来就是小 top-k，这份配方不会自动省计算。
正文没有单列局限。等 FLOPs 下准确率只升 0.69 个点，WinoGrande 和 SIQA 还下降。用七项短常识题支持「循环深度是高效缩放轴」，证据偏弱，困惑度这边更硬。序列长度 1024，没有数学、代码、长上下文，也没有算法任务或长度外推的实验。
1.7B 的涨幅伴随大约 9 倍的深度计算，并且每 3 轮一次语言建模损失和参数更新。等 FLOPs 那组关掉分段反传，收益就小一档。深度、中间监督和更新次数没有拆开。
分轮 router 与 Looping Residual 只在第 5000 步消融，困惑度差 0.6 到 1.2，稳定化组件则是 6 到 8。路由更散是图 4 里的事实，它是不是主因还没有验证。λ、β、K 没有敏感性扫描，也未与 Mix-LN、mHC 或不动点循环做同设置对比。规模停在 1.7B 和 60B token。