Looped Transformer：共享权重循环迭代，重塑深度模型效率边界

最近这段时间，Transformer 相关的研究方向出现了一个挺有意思的收敛趋势：大家不再一味把模型往“更深、更宽”的方向堆，而是开始回头思考一个非常基础的问题——同一组参数，能不能在推理时被反复使用多次？这个思路落到工程上，就是圈里一直在聊的 Looped Transformer。简单说，它不靠增加层数来提升能力，而是把同一个 Transformer Block 循环迭代若干次，让模型在固定参数规模下获得更深的“有效计算深度”。

这个标题里其实压缩了四条研究线：循环架构本身、循环 MoE 的稀疏化改造、缩放定律在循环模型上的重新拟合，以及优化器和扩散生成场景下的落地适配。这些方向单独拎出来都值得写一篇长文，合在一起看，恰好是“循环化”这个趋势从理论分析到工程实践的完整闭环。这篇文章我打算从架构设计讲起，把每条线的思路、原理、实验取舍和踩坑经验都拆开聊，希望能给正在做相关方向或者准备入坑的同学一些参考。

1. 循环 Transformer 的核心思想与架构基础

1.1 为什么要把同一个 Block 循环用

先聊最底层的问题：循环 Transformer 到底在解决什么。以往的 Transformer 模型，参数量决定能力上限，层数决定感受野和抽象层次，两者都要靠堆资源来换。但 Looped Transformer 的做法完全不同——它把同样的 Block 权重在时间维度上反复执行，用“时间换深度”。

好，这里要先纠正一个常见误解。很多人第一次听到“同一个 Block 循环使用”，第一反应是“这不就是把推理时的 depth 变大嘛”。玩法其实更细：模型通常由三类子层组成，最前面是 Input Projection 把 token 映射到隐藏维度，中间是一组循环执行的共享权重 Block，最后接 Output Projection 把隐藏状态映射回词表。共享权重的这个核心 Block 内部会包含 Attention 和 FFN，而且通常会做权重的归一化处理和残差连接调整。这种结构最大的好处是参数效率极高，因为不管循环多少次，真正参与存储的只有中间这一组权重。

我一开始是被这样一组对比数字吸引的：一个 12-layer 的普通 Transformer，大概需要 12 组独立的 Attention 和 FFN 权重；而一个共享权重的循环 Transformer 只需要 1 组权重，却可以在推理时迭代 8 步甚至 16 步来获得超过 12 层的表达能力。在参数量受限的场景（比如边缘端部署、小规模实验预算），这种取舍非常划算。

但“参数省了”不等于“能力白拿”。循环展开后会带来自动求导路径的急剧加深，以及训练时梯度需要穿过整个循环链，这给优化器带来了不小压力。所以后面会单独开一节讲优化器相关对策，这里先记住结论：架构上省参，优化上补课，两件事必须同时做。

1.2 循环块的关键设计细节

真正决定 Looped Transformer 成败的，是循环块内部怎么设计。我看到过不少失败的复现，问题几乎都出在三个细节上：归一化的位置、残差连接的缩放，以及循环块前后的预/后处理层。

先说归一化。普通的 Pre-LN Transformer，每个子层前面都有一层 LayerNorm，这是标准做法。但循环场景下有个微妙的问题：如果共享 Block 内的 LayerNorm 参数也一起共享，那每次迭代都会用同一套均值/方差统计，这本身问题不大；真正麻烦的是，如果 Block 内部的残差分支没有做好尺度控制，多次迭代会让隐藏状态的范数越滚越大，训练很快就爆。一个稳妥的做法是在残差连接上乘一个可学习的缩放因子，或者在循环边界处加一个额外的 Norm，保证隐藏状态在进入下一次迭代之前被拉回到合理的尺度范围。

再说预/后处理层。Input Projection 和 Output Projection 的维度会影响整个模型的表达瓶颈。因为中间的循环 Block 是共享的，如果隐藏维度过小，信息在多次迭代中的传递容量会受限；如果隐藏维度过大，省参的优势会被抵消不少。我自己做实验时的经验是：在固定总参数量预算下，隐藏维度取普通 Transformer 的 1.5 到 2 倍往往能获得更好的效果，因为循环块本身参数少，可以把省下来的参数预算“补贴”给投影层。

最后一个容易被忽视的设计是循环步数的动态范围。训练时固定步数还是允许动态步数，会显著影响收敛形状。早期的实验通常固定 4 步或者 6 步训练；后期的一些工作会在推理时增加步数以获得更高精度，这其实等价于一个测试时计算预算的动态调节开关，很有意思。

2. 循环 MoE：把稀疏激活塞进循环结构

2.1 MoE 化循环块的设计动机

循环 Transformer 虽然省参数，但它的计算量并不省——每次迭代都要完整跑一遍 Attention 和 FFN，迭代次数一多，推理成本和训练成本都线性上涨。于是就会自然想到一个方向：能不能让循环块内部变得稀疏化，比如只在部分迭代中激活部分参数？

这就是“循环 MoE”的思路。传统 MoE 是在一层内放多个 Expert，每个 token 根据路由结果只激活其中若干 Expert；而循环 MoE 把 Expert 平铺在“时间维度”上——同一个循环块内部可以包含多组共享权重里的不同子模块，每次迭代通过路由决定用哪几条路径。

这个想法在实践中有个非常直观的好处：传统 MoE 为了增加专家数量，必须增加物理参数量，这对显存很不友好；循环 MoE 则是在固定物理参数量下，通过迭代次数去“伪造”出一个更大的专家池。比如一个循环块内含 4 个 FFN Expert，迭代 4 次，理论上就等价于一个 4 层 × 4 Expert 的 MoE 模型，但物理参数量只相当于 4 个 Expert 的规模。

2.2 路由策略与负载均衡的实操取舍

路由策略是循环 MoE 的重头戏，这里有几个实操层面的坑。首先，路由在循环结构中天然具有“逐步精炼”的性质：早期迭代的路由结果不必特别精准，因为后面还有多次迭代可以纠正；但传统 MoE 只有一层机会，路由错了这个 token 的专家选择就定死了。这个差异让循环 MoE 可以使用更激进的路由温度参数，我在实验中用 softmax 温度 0.8 到 1.2 都试过，比传统 MoE 更稳定。

其次，负载均衡损失（Load Balancing Loss）在循环 MoE 里需要加一个时间维度的惩罚项。否则很容易出现一种现象：前几次迭代所有 token 都往同一个 Expert 跑，后面几次迭代才分散开，造成所谓的“时间塌陷”。我在自己的实验里发现，如果对每个迭代步的专家利用率分别计算负载均衡损失，然后取平均，效果比只看总利用率好很多。因为总利用率均衡不代表每步都均衡。

路由的另一个重要变体是“可学习停止”。既然循环步数本身可以被动态决定，那么路由可以兼任“是否继续迭代”的判断器。对简单样本，模型可以提前跳出循环，节省计算；对复杂样本，则继续迭代。这种做法在推理时优势很大，但也意味着训练时要处理变长计算图，对框架和显存控制要求更高。

2.3 循环 MoE 对比标准 MoE 的优劣边界

做了大量对比实验后，我的感受是：循环 MoE 不是标准 MoE 的替代品，而是互补品。标准 MoE 擅长在超大规模参数下扩展专家数量，每个 token 可以选择非常多样化的专家组合；循环 MoE 的优势场景是参数预算中等（比如 1B 以下）、同时希望推理时能动态调整计算量。

循环 MoE 劣势也很明显：因为它共享权重，每一次迭代都使用同一批 Expert，多轮迭代之间其实是在“反复锤打”同一组参数而不是引入新的能力维度。换句话说，标准 MoE 增加“广度”，循环 MoE 增加“深度”。负责任地说，凡是任务本身需要真正多样化的知识分布，标准 MoE 会占优；凡是任务需要多步推理链或者迭代精炼（比如数学推理、代码执行模拟），循环 MoE 的性价比会更高。

我还试过把两种做法的混合体——循环块内部用 MoE 结构，循环步数之外再叠加普通层。实验结果中规中矩，收益不明显但参数量上去了，现阶段不建议把这当默认方案。

3. 缩放定律：循环模型如何重新定义规模尺度

3.1 从层数缩放走向“步数缩放”

缩放定律（Scaling Laws）是这几年大模型领域的基础共识：模型性能与参数量、数据量、计算量之间存在稳定的幂律关系。但 Looped Transformer 直接挑战了这个框架里的一个重要假设——模型深度是确定层数，性能随层数增加而平滑提升。循环架构下，深度不再直接对应参数量，而是对应迭代步数，这等于把一个连续的“缩放宽高比”问题变成了“时间展开”问题。

更新后的缩放曲线大致长这样：在固定参数量下，性能随循环步数增加而上升，但会收敛到某个平台；这个平台位置取决于模型宽度、数据质量和优化器设置。换句话说，循环模型存在一个“有效深度饱和点”，超过这个点再增加步数，收益会变得很小，但计算成本仍在线性增加。

这个饱和点不是一个固定值。在小规模数据集上，它可能出现在 4 到 6 步；数据集足够大、优化充分的情况下，可以推到 10 到 14 步。所以做缩放实验时，不要只测一个固定步数就下结论——选择一个过低的步数会把循环模型的真实能力低估不少。

3.2 FLOPs 预算约束下的最优配置选择

缩放定律分析的一个重要用途，是在固定 FLOPs 预算下找出最优的配置组合。对普通 Transformer，这个组合通常是“深度与宽度等比扩大”；对循环 Transformer，最优组合变成了“参数量、宽度、迭代步数的三维搜索”。

我复现实验结果时使用过一组相对清晰的对照：在同一份数据上训练不同配置，并固定总计算量为一个固定量级，之后对比不同配置的最终损失。结论是：中等宽度乘以中等步数通常优于“很窄乘以很多步”或“很宽乘以很少步”。原因是前者能在参数容量和计算深度之间取得更好的平衡——太窄时循环块内的信息瓶颈过强，太宽时参数增加导致欠拟合。

另外一个比较反直觉的发现是：在低资源预算下（比如训练数据只有 1 亿 token），循环模型的缩放曲线比普通 Transformer 更早进入饱和区。也就是说，循环模型更适合中等以上的数据量；数据太少时，共享权重带来的参数效率优势无法抵消优化难度上升的劣势。这一点对刚开始做相关实验的人尤其重要：如果你想用一个小数据集快速验证循环模型的想法，效果很可能没有理论预期那么好，不是方法错了，是数据量还不到甜区。

3.3 缩放视角下“小模型模拟大模型”的价值

循环模型的缩放性质还带了一个非常实用的副产品：训练一个参数较小的循环模型，在推理时增大迭代步数，可以逼近参数量大得多的非循环模型的能力。这个特性很有工程价值，尤其考虑到大模型推理时的显存压力很大，如果能在固定模型权重下通过增加少量时间成本来换取部分深度增益，那是相当划算的交易。

不过这个特性只在一定范围内成立。我实测下来的粗略判断：一个 3 亿参数的循环模型，迭代 12 步左右，大约能摸到 10 亿参数普通 Transformer 的能力下限；但要再往上追，就比较吃力了。原因在于参数容量本身设置了能力天花板，循环无法无中生有地创造知识容量，它提升的是“同一批参数被用得更充分”的效率。

所以，如果团队在模型部署上有明确的显存预算上限，循环模型值得做小规模验证；如果目标是冲击超大模型且预算宽裕，循环模型定位成主模型的辅助模块可能更现实。

4. 优化器与训练稳定性：循环模型最容易被低估的关卡

4.1 梯度路径变长带来的稳定性问题

训练循环 Transformer 时最容易翻车的点，不是模型架构设计，而是优化器和梯度流的稳定性。循环展开 N 步，梯度的反向传播路径相当于 N 倍层数的普通 Transformer。举例说明，一个 6 步循环块等效于 6 倍深度的网络，普通 8 层模型变成等效 48 层深度，梯度爆炸/消失的风险指数上升。

我自己的经验是三个对策缺一不可：梯度裁剪阈值要比普通 Transformer 严格，建议设到 0.5 左右（普通模型通常用 1.0）；优化器尽量使用带解耦权重衰减的 AdamW，其中 beta2 可以适当调低，让优化器对梯度的剧烈变化更敏感；循环块内部尽量采用 Pre-LN 结构，并且在每次迭代开始的输入处对隐藏状态做一次 RMSNorm，防止范数漂移。

比较反直觉的一点是：在一些循环模型中，学习率反而可以比普通 Transformer 设得更大。因为共享权重使得同一个参数收到多次梯度更新，有效学习率其实比名义上看到的更大，所以初始学习率可以适度降低而不是提高。我见过很多失败的复现，都是因为把学习率按普通 Transformer 的默认值 3e-4 去调，结果训练很快就发散，其实把学习率降到 1e-4 到 2e-4 反而更稳。

4.2 EMA 与权重平均的实战价值

对于循环模型，权重的指数滑动平均（EMA，Exponential Moving Average）是一个被严重低估的优化技巧。普通 Transformer 的 EMA 效果通常不太明显，但在循环架构下，共享权重被反复迭代使用，权重的微小扰动会被放大 N 倍。使用 EMA 后的模型，即使原始权重在训练后期有些震荡，平滑后的权重往往能显著提升评估指标。从我自己的实测结果看，EMA 在循环模型上的收益通常比普通 Transformer 大 1 到 2 个点，几乎是白捡的优化，推荐纳入默认配置。

EMA 的衰减系数建议设到 0.999 到 0.9999 之间，训练足够长时用偏大的值；如果训练步数不多，0.999 就够了。实践上还有一个细节：EMA 权重最好在验证评估时使用，而不是在训练过程中替代原始权重继续更新，否则会破坏训练的收敛性。

4.3 显存优化与重计算策略

循环模型的训练显存压力集中在隐藏状态序列的保存上——每次迭代都要保留一份完整的激活值，反向传播才能逐站回传梯度。N 次迭代就需要保留 N 份激活，显存消耗随步数线性增长。

实用的解决方案是梯度检查点（Activation Checkpointing），即在正向传播时只保留每次迭代的输入 hidden state，不保留中间激活；反向传播时重新计算正向过程并临时构建图。这样显存占用从 O(步数 × 序列长度 × 隐藏维度) 降为 O(序列长度 × 隐藏维度)，代价是训练时间增大约 30% 到 50%。

如果显存依然吃紧，还有一个更激进的做法：用分布式重计算，只保留每两步或者每三步的检查点，中间步的激活全部舍弃，回传时重新计算。这个策略的显存占用和重计算开销达到了更细粒度的平衡，适合在 7B 以上规模的循环模型中使用。

另外，混合精度训练在循环模型里需要额外小心。低精度训练特别是 Float16 下，循环块内多次迭代容易出现精度漂移，因为我上一条提到过隐藏状态范数会逐步累积。建议使用 BFloat16 而不是 Float16，并且对循环块输入输出之间的残差流保持高精度（Float32），只在 Attention 和 FFN 的计算子层切换到低精度。这个操作我在多次复现中验证过，对稳定性和精度都有明显帮助。

5. 扩散生成场景下的 Looped Transformer 应用

5.1 扩散模型与循环迭代的天然契合

扩散模型的采样过程本质上是多步去噪：从一个纯噪声开始，经过数十乃至上百步的迭代，逐步恢复出干净的图像或数据。这正好和 Looped Transformer“同一组权重反复精炼”的核心机制不谋而合。传统的扩散模型主干通常使用 U-Net 或 DiT（Diffusion Transformer），U-Net 靠不同尺度的编解码结构去捕捉不同频率的信息，DiT 则完全是标准 Transformer 层堆叠。如果把 DiT 中的多个 Transformer Block 换成共享权重的循环块，路径就是现成的。

外围做法上，做得巧妙的地方在于了循环块的迭代步数直接对应扩散模型的去噪步数节奏。刚开始使用大噪声时模型只需要粗糙地去除低频污染，迭代少一点；到后期精细纹理生成阶段，迭代更多步来打磨细节。这种做法等于让模型的“有效深度”随着采样进度动态变化，比固定深度的 DiT 更贴合扩散过程本身的特性。

5.2 DiT 架构循环化的改造要点

把普通 DiT 循环化，不是简单地把多个 Block 替换成一个共享 Block 就能完事。有四个地方必须单独处理。

第一，AdaLN 条件注入（Adaptive LayerNorm）。DiT 的核心特征是用时间步和类别标签来调制 LayerNorm 的 scale 和 shift 参数，循环化后这个注入点是在每次迭代都要进行，还是只在第一次迭代进行？我实测下来，每次迭代都注入效果更稳定，因为去噪过程每一阶段需要的调制不同。

第二，位置编码。普通 Transformer 的位置编码是一次性的，循环模型中每次迭代看到的是同一个序列，位置编码如果重复使用会造成信息冗余干扰。更强悍的做法是在每次迭代时结合迭代步数引入一个“时间步感知的缩放因子”，把迭代序号编码成向量注入到 Attention 的 bias 里。这个细节对生成质量影响比想象中更大，一个不加只是简单共享的循环扩散模型，生成的图像会出现重复纹理。

第三，最后一层输出头。扩散模型最终要预测噪声或者样本本身，循环块输出之后不要再叠加多余的层，否则容易让训练时梯度负担过重。

第四，随机丢弃迭代步训练。为了让模型不依赖固定的迭代次数，训练时让步数随机采样，比如从 4 到 10 之间随机选，可以显著提高不同采样步数下的鲁棒性。这是扩散模型本身一个经典训法（随机时间步采样）在循环维度上的翻版。

5.3 推理步数调节与质量增益实测

在扩散生成场景下，循环模型另外的玩法在于控制“测试时计算预算”。固定权重不变，推理时从 6 步提升到 12 步，生成图像的 FID 分数通常会有可观察的下降（意味着质量更好）。这个趋势和缩放定律一样会遇到饱和：步数超过一定阶段，继续增加反而可能引入伪影。

从我的经验看，一个约 5 亿参数的循环 DiT，迭代 8 到 10 步，可以达到约 12 亿参数普通 DiT 相近的生成质量，而参数量不到后者的一半。代价是推理时间变长，因为循环迭代步数增加，实际计算量与普通模型的深层网络相当。

这个“参数量换时间”的特性和图像生成场景很契合：模型分发成本低，用户端因为推理时间适中，体验不会太差。另外，在扩散这种多步采样天然并行的场景，可以把前几步使用较少循环步数、后几步使用较多循环步数，进一步控制平均推理开销。这种非均匀步数策略在普通 DiT 上不好实现，因为普通模型每个 Block 本身是固定深度，无法在一条采样轨迹中动态改变层数；循环模型则天然支持这种操作。

6. 实验配置与避坑经验总结

6.1 一个可供参考的基准实验配置

对于刚开始切入这个方向的人，我给一个经过了多次调优的基准配置，作为起点比较稳妥。它不一定是最优的，但可以帮避免一些低级崩溃。

这个配置在中等规模（3 亿到 5 亿参数，约 50 亿训练 token）上表现比较稳，训练崩溃率很低。不推荐一上来就上 12 步甚至 16 步的训练配置，除非显存和数据量都能跟上。

6.2 常见训练崩溃问题定位表

循环模型训练比普通模型更容易出现“看起来在收敛、突然就炸了”或者“loss 一直在原地不动”的疑难杂症。归类一下常见症状和排查思路。

6.3 后续可以继续深挖的方向

这个方向还有几个值得跟进的问题。第一个是理论层面的：循环模型的有效深度饱和现象到底由什么决定，目前更多是经验的表征，理论上还缺少像普通 Transformer 那样相对清晰的 scaling law 推导。第二个是模型架构的通用性：现在循环化的实验主要集中在语言模型和扩散生成上，对于序列推荐、时间序列预测这类同样强调多步推理的任务，很可能也适用。第三个是训练策略的改进：目前的优化器基本沿用普通架构的策略，是否能针对循环共享权重设计专门的二阶优化器或自适应学习率方案，还缺乏系统性研究。

我个人的判断是，这个方向未来一年可能还会有一个爆发期，尤其是当更多人意识到“参数量不变、计算量可调”这条路线在边缘部署和推理成本控制上的价值之后。与此同时，严谨的对比实验和训练稳定性方法论会成为区分有效与无效工作的分水岭。如果你正在考虑用这个架构做自己的项目，建议从规模可控的扩散生成任务切入，既有清晰的任务指标，又能直观感受迭代步数的调控效果。

最后再分享一个小技巧：在实验初期给循环步数设一个可变的随机采样范围交替训练，同时用固定步数做验证，这种对步数区间内的泛化能力约束，能明显提升最终模型在不同推理预算下的稳定性。我踩过好几次“训练固定 6 步、推理 10 步效果反而变差”的坑，加上随机步数之后就顺了很多，你可以直接试试。