Layerwise Recurrent Router for Mixture-of-Experts

2024-08-13

arXiv.org

AI Summary

Existing Mixture-of-Experts (MoE) models suffer from low parameter efficiency: a 52B-parameter MoE model performs comparably to only a 6.7B-parameter dense model, primarily because layer-wise routers make independent decisions—ignoring historical routing patterns—and thus yield suboptimal token-to-expert assignments. To address this, we propose Inter-layer Recursive Routing (IRR), the first approach to incorporate Gated Recurrent Units (GRUs) into MoE routing, enabling cross-layer modeling of routing dependencies while preserving parallel computation. IRR is architecture-agnostic and orthogonal to existing MoE designs, ensuring broad compatibility. Extensive experiments demonstrate that IRR consistently outperforms mainstream MoE baselines across multiple language modeling benchmarks. It significantly improves parameter utilization efficiency and expert activation rates, elevating the performance of the 52B MoE model to match or exceed that of more optimally configured dense models.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Planning, Routing, and Scheduling: Planning with Language Models
• Natural Language Processing: (Large) Language Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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 .

Problem

• Improves parameter efficiency in Mixture-of-Experts models.
• Enhances token-expert combinations using historical routing information.
• Introduces cross-layer information sharing for better expert selection.

Innovation

• Layerwise Recurrent Router for Mixture-of-Experts
• GRU for cross-layer routing dependencies
• Seamless compatibility with existing MoE architectures