Layerwise Recurrent Router for Mixture-of-Experts

📄 arXiv: 2408.06793v2 📥 PDF

作者: Zihan Qiu, Zeyu Huang, Shuang Cheng, Yizhi Zhou, Zili Wang, Ivan Titov, Jie Fu

分类: cs.CL

发布日期: 2024-08-13 (更新: 2025-03-19)

期刊: ICLR2025

🔗 代码/项目: GITHUB

💡 一句话要点

提出层间循环路由RMoE，提升混合专家模型参数效率

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家模型 层间循环路由 参数效率 门控循环单元 大型语言模型

📋 核心要点

  1. 现有MoE模型参数效率低，各层路由独立决策，忽略了历史路由信息，导致token-expert组合次优。
  2. RMoE利用GRU建立层间路由决策依赖，通过层间循环实现高效并行计算，提升参数效率。
  3. 实验表明，RMoE优于多种基线模型，且能与其他MoE架构兼容，有效提升专家选择和多样性。

📝 摘要（中文）

大型语言模型（LLMs）的扩展极大地提升了它们在各种任务中的能力，但这种增长需要与高效的计算策略相匹配。混合专家（MoE）架构因其能够在不显著增加训练成本的情况下扩展模型大小而脱颖而出。尽管具有优势，但当前的MoE模型通常表现出参数效率低下。例如，一个基于MoE的预训练LLM，拥有520亿参数，其性能可能与一个标准的67亿参数模型相当。作为MoE的关键部分，当前不同层的路由器独立地分配tokens，而不利用历史路由信息，这可能导致次优的token-expert组合和参数效率问题。为了缓解这个问题，我们引入了用于混合专家的层间循环路由器（RMoE）。RMoE利用门控循环单元（GRU）来建立连续层之间路由决策的依赖关系。这种层间循环可以为输入tokens高效地并行计算，并引入可协商的成本。我们广泛的经验评估表明，基于RMoE的语言模型始终优于一系列基线模型。此外，RMoE集成了与现有方法正交的新计算阶段，从而可以与其他MoE架构无缝兼容。我们的分析将RMoE的收益归因于其有效的跨层信息共享，这也改善了专家选择和多样性。我们的代码位于https://github.com/qiuzh20/RMoE。

🔬 方法详解

问题定义：现有混合专家模型（MoE）存在参数效率低下的问题。具体来说，不同层的路由模块独立进行token分配，忽略了层与层之间的路由信息依赖关系。这种独立性可能导致次优的token-expert组合，使得模型在达到相似性能时需要更多的参数。

核心思路：RMoE的核心思路是引入层间循环机制，使得每一层的路由决策能够利用之前层的路由信息。通过在层之间传递信息，模型可以更好地理解token的上下文，从而做出更明智的专家选择，提高参数利用率。

技术框架：RMoE在MoE架构的每一层路由模块中引入了一个门控循环单元（GRU）。对于每个输入token，GRU接收当前层的路由logits以及上一层的隐藏状态作为输入，输出更新后的隐藏状态和路由logits。更新后的logits用于选择该层对应的专家。整个过程可以并行地应用于所有token。

关键创新：RMoE的关键创新在于引入了层间循环依赖，打破了传统MoE模型中各层路由独立决策的模式。这种层间信息共享使得模型能够更好地利用上下文信息，从而提高路由效率和参数利用率。此外，RMoE的设计与现有MoE架构正交，可以方便地集成到各种MoE模型中。

关键设计：RMoE使用标准的GRU单元来建模层间依赖关系。GRU的输入包括当前层的路由logits和上一层的隐藏状态。GRU的输出包括更新后的隐藏状态和路由logits。损失函数与原始MoE模型保持一致，没有引入额外的损失项。实验中，GRU的隐藏层大小与路由logits的维度相同。作者开源了代码，方便复现和进一步研究。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，RMoE在多个语言模型任务上均优于基线模型。例如，在相同的参数规模下，基于RMoE的语言模型能够达到更高的perplexity。此外，RMoE还能够提高专家选择的多样性，避免过度依赖少数几个专家，从而提高模型的泛化能力。代码已开源，方便研究人员复现和进一步研究。

🎯 应用场景

RMoE可应用于各种需要大规模语言模型的场景，例如机器翻译、文本生成、对话系统等。通过提高参数效率，RMoE能够降低训练和推理的计算成本，使得更大规模的模型能够在资源受限的环境中部署。此外，RMoE的模块化设计使其易于集成到现有的MoE架构中，具有广泛的应用前景。

📄 摘要（原文）

The scaling of large language models (LLMs) has revolutionized their capabilities in various tasks, yet this growth must be matched with efficient computational strategies. The Mixture-of-Experts (MoE) architecture stands out for its ability to scale model size without significantly increasing training costs. Despite their advantages, current MoE models often display parameter inefficiency. For instance, a pre-trained MoE-based LLM with 52 billion parameters might perform comparably to a standard model with 6.7 billion parameters. Being a crucial part of MoE, current routers in different layers independently assign tokens without leveraging historical routing information, potentially leading to suboptimal token-expert combinations and the parameter inefficiency problem. To alleviate this issue, we introduce the Layerwise Recurrent Router for Mixture-of-Experts (RMoE). RMoE leverages a Gated Recurrent Unit (GRU) to establish dependencies between routing decisions across consecutive layers. Such layerwise recurrence can be efficiently parallelly computed for input tokens and introduces negotiable costs. Our extensive empirical evaluations demonstrate that RMoE-based language models consistently outperform a spectrum of baseline models. Furthermore, RMoE integrates a novel computation stage orthogonal to existing methods, allowing seamless compatibility with other MoE architectures. Our analyses attribute RMoE's gains to its effective cross-layer information sharing, which also improves expert selection and diversity. Our code is at https://github.com/qiuzh20/RMoE .