Layerwise Recurrent Router for Mixture-of-Experts (2408.06793v1)

Abstract: The scaling of 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 LLMs 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

Layerwise Recurrent Router for Mixture-of-Experts

The paper introduces a novel approach for enhancing the Mixture-of-Experts (MoE) framework by developing a Layerwise Recurrent Router (RMoE). This method aims to address issues of parameter inefficiency noted in many MoE models despite their scale, such as a 52B parameter MoE performing on par with a significantly smaller 6.7B parameter standard model as highlighted by the authors referencing prior work. The central hypothesis is that existing MoE routers, which operate independently across layers, fail to leverage historical routing data, leading to non-optimal token-expert allocations and inefficient parameter utilization.

Methodological Advancements

RMoE differentiates itself by implementing a Gated Recurrent Unit (GRU) that links routing decisions across layers. This mechanism is intended to create a dependency on previous layer decisions, theoretically improving expert selection and token routing efficiency. The introduction of this GRU is purported to prevent the convergence of routing decisions to suboptimal, token-id-dependent mappings seen in some MoE models. Moreover, the recurrent connections in RMoE are meant to be computed efficiently, not imposing prohibitive computational costs contrary to traditional sequence-level recurrences.

The unique architecture involves projecting hidden states using separate projectors for different layers to fit into the GRU, which then generates routing decisions influenced by previous layers' routing choices. This approach also integrates a decoupled computation stage to make it compatible with existing MoE methods while providing enhanced cross-layer information sharing.

Empirical Evaluation

Extensive experiments were conducted on various LLMing tasks, including both pre-training and fine-tuning scenarios. The results consistently demonstrated that models utilizing RMoE outperformed a range of baseline models, including those employing fixed or more complex router configurations like HyperMoE and SMoE-MLP. Remarkably, the RMoE achieves these improvements with minimal increment in computational costs.

Key findings from the analyses suggest that RMoE's enhanced performance stems not only from parameter enhancements but significantly from the cross-layer recurrent information sharing and additional gradient propagation pathways introduced by the GRU. This new gradient pathway is crucial for the learning dynamics of MoE models, leading to more well-distributed routing decisions across layers.

Implications and Future Directions

The implications of RMoE are significant for both theoretical and practical advancements in large-scale neural networks. Theoretically, it provides a path forward for more efficient utilization of parameters in sparse expert models by addressing core inefficiencies in current MoE implementations. Practically, RMoE offers a template for integrating structured routing strategies in LLMs without incurring substantial computational overhead.

Future work may explore extending the recurrent mechanism to other components of neural architectures or integrating it with emergent MoE configurations that emphasize expert precision and task-specific routing. Additionally, optimizing the implementation of RMoE in distributed settings could enable its application across broader AI-driven domains, potentially extending to multimodal systems and beyond. This paper lays substantial groundwork for innovations in efficient neural computation strategies through modular architectural improvements.