超越两次循环：循环专家混合的可扩展配方

循环 Transformer 引入循环深度作为 LLM 的新缩放轴：通过重复应用共享 Transformer 块，它们可以在不增加参数数量的情况下增加有效深度。然而，在 FLOPs 匹配比较下，对于大型 MoE LLM 来说，循环的好处仍然不清楚。主要原因是额外迭代带来的收益会迅速减少，甚至可能变成退化，因此花在循环上的额外 FLOP 几乎没有带来实质性的改进。因此，之前的工作通常集中在两个循环上。我们确定了扩展循环 MoE 的两个主要障碍。首先，循环继承并放大了深度灾难：随着残差更新的累积，隐藏状态方差随着每次迭代而增长，这会破坏深度递归并导致表示漂移。其次，循环 MoE 会遭受专家选择崩溃的困扰：路由器在循环中重复选择相同的专家，因此额外的迭代会增加计算量，但不会增加计算多样性。在这一诊断的指导下，我们提出了 LOOM，它建立在一个原则之上：每个循环都应该贡献新的计算，同时保持循环状态稳定。 LOOM 通过缩放残差更新来限制方差增长并在每个循环中重新注入输入嵌入来稳定循环，并通过参与不同专家的每个循环路由器和向前传送早期输出的循环残差来使其多样化。 100M-1.7B 模型的实验表明，可以稳定地扩展到 9-12 个循环。在近 iso-FLOP 下，700M 模型在 5 个循环时表现最佳，与非循环基线相比，困惑度从 18.36 降低到 16.54，平均零样本准确率从 38.84% 提高到 39.53%。在没有 FLOP 匹配的情况下，在 60B 词元上训练的 1.7B 模型在 9 个循环时达到峰值，将困惑度从 9.62 降低到 7.77，并将平均零样本准确率从 42.4% 提高到 47.7%。代码可用 https://github.com/hed-ucas/LOOM.