论文

LoopMoE:将迭代计算与语言建模专家混合相结合

LoopMoE: Unifying Iterative Computation with Mixture-of-Experts for Language Modeling

模型架构混合架构

摘要

专家混合 (MoE) 和循环架构沿两个正交轴(即参数容量和有效深度)缩放模型。然而,主流的循环架构依赖于将参数计数与每个词元的 FLOP 结合起来的密集主干,这使得不可能在匹配的预算下隔离迭代计算的影响。为此,我们提出了 LoopMoE,一种循环 MoE 语言模型,通过两种设计将稀疏路由与迭代权重共享计算集成在一起。第一个是 IterAdaLN,它通过联合调节迭代索引和每个词元隐藏状态的调制信号来解决权重共享对称性。第二个是容量平衡策略,它恢复经过良好调整的非循环参考的注意力到 FFN 活动参数的比率。这些设计共同实现了在相同的总参数、每个词元 FLOP 和活动子层比率下对循环 MoE 与普通 MoE 进行严格控制的头对头评估。在九个下游基准测试中,LoopMoE 相对于其匹配的普通 MoE 的平均改进从 3B 等级的 1 分以上增加到 9B 等级的大约 3 分。这些结果提供了初步证据,表明迭代稀疏计算的好处可能会随着规模的扩大而增强,从而将 LoopMoE 定位为可扩展循环语言模型的有前景的架构。