XMoE: Sparse Models with Fine-grained and Adaptive Expert Selection

TL;DR

XMoE addresses computational waste in sparse Mixture-of-Experts by decomposing FFNs into many small experts and routing tokens with a novel threshold-based router that adapts the number of engaged experts per token. The approach yields substantial FLOPs reductions (over 50% in MoE layers) while maintaining or improving performance on language modeling and machine translation, and it can be applied to dense models to enable sparse inference. The method includes an auxiliary balancing loss and a complexity-aware design, and demonstrates favorable comparisons against established MoE baselines like Top-1, Top-k, and Switch Transformer. Overall, XMoE advances efficient conditional computation by enabling fine-grained expert selection and adaptive resource allocation, with practical implications for large-scale transformers and sparse inference regimes.

Abstract

Sparse models, including sparse Mixture-of-Experts (MoE) models, have emerged as an effective approach for scaling Transformer models. However, they often suffer from computational inefficiency since a significant number of parameters are unnecessarily involved in computations via multiplying values by zero or low activation values. To address this issue, we present \tool, a novel MoE designed to enhance both the efficacy and efficiency of sparse MoE models. \tool leverages small experts and a threshold-based router to enable tokens to selectively engage only essential parameters. Our extensive experiments on language modeling and machine translation tasks demonstrate that \tool can enhance model performance while decreasing the computation load at MoE layers by over 50\% without sacrificing performance. Furthermore, we present the versatility of \tool by applying it to dense models, enabling sparse computation during inference. We provide a comprehensive analysis and make our code available at https://github.com/ysngki/XMoE.

Figures (8)

• Figure 1: Average percentage of positive values in the FFN layers after the activation function. All models are decoder-only Transformers with 12 layers.
• Figure 2: Overview of an MoE layer in XMoE, where tokens are routed to small experts by an adaptive router.
• Figure 3: Overview of the threshold-based router. The number of tokens processed per expert is determined according to both the total token number and the capacity factor.
• Figure 4: Test perplexity (PPL) with regard to the the normalized FLOPs in the MoE layer during inference. We adjust FLOPs by modifying the capacity factor $\gamma$.
• Figure 5: Average number of required experts with regards to training steps across different layers.