将稀疏专家与迭代计算结合,实现参数与计算量双重可控的高效语言模型。
LoopMoE: Unifying Iterative Computation with Mixture-of-Experts for Language Modeling

- 通过迭代自适应层归一化和容量平衡策略,实现权重共享下的稀疏路由。
- 在3B到9B模型规模下,相比同规模纯MoE模型,平均性能提升1至3点。
- 适合追求高扩展性、低计算成本的大型语言模型研究者使用。
Mixture-of-Experts (MoE) 与循环架构分别从参数容量和有效深度两个正交维度扩展模型。然而主流循环架构依赖密集主干,导致参数量与每令牌浮点运算量耦合,无法在相同预算下独立评估迭代计算的效果。为此,我们提出 LoopMoE,一种融合稀疏路由与迭代权值共享计算的循环 MoE 语言模型,包含两项设计:其一是 IterAdaLN,通过联合依赖迭代次数与每令牌隐藏状态的调制信号打破权值共享对称性;其二是容量平衡策略,恢复了经过良好调优的非循环参考模型中注意力层与前馈层的活跃参数比例。上述设计使首次在总参数量、每令牌浮点运算量及活跃子层比例完全一致的前提下,实现循环 MoE 与普通 MoE 的严格对比。在九个下游基准测试中,LoopMoE 相较于匹配的普通 MoE,平均性能提升从30亿规模的1点以上增至90亿规模的约3点。结果初步表明,迭代稀疏计算的优势可能随规模增大而增强,为可扩展循环语言模型提供了有前景的架构方向。
原文摘要 · Abstract (English)
Mixture-of-Experts (MoE) and looped architectures scale models along two orthogonal axes, namely parameter capacity and effective depth. However, mainstream looped architectures rely on dense backbones that couple parameter count with per-token FLOPs, which makes it impossible to isolate the effect of iterative computation under matched budgets. To this end, we present LoopMoE, a looped MoE language model that integrates sparse routing with iterative weight-shared computation through two designs. The first is IterAdaLN, which resolves weight-sharing symmetry via a modulation signal jointly conditioned on the iteration index and the per-token hidden state. The second is a capacity-balancing strategy that recovers the attention-to-FFN active parameter ratio of well-tuned non-looped references. Together, these designs enable the first strictly controlled, head-to-head evaluation of a looped MoE against a Vanilla MoE under identical total parameters, per-token FLOPs, and active sublayer ratios. Across nine downstream benchmarks, LoopMoE's average improvement over its matched vanilla MoE increases from over 1 point at the 3B scale to approximately 3 points at the 9B scale. These results provide initial evidence that the benefits of iterative sparse computation may strengthen with scale, positioning LoopMoE as a promising architecture for scalable looped language models.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。